package setup

import (
	"context"
	"time"

	"forge.cadoles.com/wpetit/clearcase/internal/config"
	"forge.cadoles.com/wpetit/clearcase/internal/core/service"
	"github.com/bornholm/genai/llm"
	"github.com/bornholm/genai/llm/provider"
	"github.com/bornholm/genai/llm/retry"
	"github.com/pkg/errors"

	_ "github.com/bornholm/genai/llm/provider/all"
	"github.com/bornholm/genai/llm/provider/openrouter"
)

func NewForgeManagerFromConfig(ctx context.Context, conf *config.Config) (*service.ForgeManager, error) {
	client, err := provider.Create(ctx,
		provider.WithChatCompletionOptions(provider.ClientOptions{
			Provider: provider.Name(conf.LLM.Provider.Name),
			BaseURL:  conf.LLM.Provider.BaseURL,
			APIKey:   conf.LLM.Provider.Key,
			Model:    conf.LLM.Provider.Model,
		}))
	if err != nil {
		return nil, errors.Wrapf(err, "could not create llm client '%s'", conf.LLM.Provider.Name)
	}

	if conf.LLM.Provider.Name == string(openrouter.Name) {
		client = &extendedContextClient{
			client: client,
			extend: func(ctx context.Context) context.Context {
				// Automatically "compress" prompts when using openrouter
				// See https://openrouter.ai/docs/features/message-transforms
				ctx = openrouter.WithTransforms(ctx, []string{"middle-out"})
				return ctx
			},
		}
	}

	client = retry.Wrap(client, time.Second, 5)

	forgeFactories, err := getForgeFactories(conf)
	if err != nil {
		return nil, errors.Wrap(err, "could not get forge factories")
	}

	forgeManager := service.NewForgeManager(client, forgeFactories...)

	return forgeManager, nil
}

type extendedContextClient struct {
	client llm.Client
	extend func(ctx context.Context) context.Context
}

// Embeddings implements llm.Client.
func (c *extendedContextClient) Embeddings(ctx context.Context, input string, funcs ...llm.EmbeddingsOptionFunc) (llm.EmbeddingsResponse, error) {
	ctx = c.extend(ctx)
	return c.client.Embeddings(ctx, input, funcs...)
}

// ChatCompletion implements llm.Client.
func (c *extendedContextClient) ChatCompletion(ctx context.Context, funcs ...llm.ChatCompletionOptionFunc) (llm.ChatCompletionResponse, error) {
	ctx = c.extend(ctx)
	return c.client.ChatCompletion(ctx, funcs...)
}

// Model implements llm.Client.
func (c *extendedContextClient) Model() string {
	return c.Model()
}

var _ llm.Client = &extendedContextClient{}