package setup

import (
	"context"

	"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/pkg/errors"

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

func NewForgeManagerFromConfig(ctx context.Context, conf *config.Config) (*service.ForgeManager, error) {
	client, err := provider.Create(ctx,
		provider.WithConfig(&provider.Config{
			Provider: provider.Name(conf.LLM.Provider.Name),
			BaseURL:  conf.LLM.Provider.BaseURL,
			Key:      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
			},
		}
	}

	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
}

// ChatCompletion implements llm.Client.
func (c *extendedContextClient) ChatCompletion(ctx context.Context, funcs ...llm.ChatCompletionOptionFunc) (llm.CompletionResponse, 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{}