package llm

import (
	"context"
	"fmt"

	anyllm "github.com/mozilla-ai/any-llm-go"
	"github.com/mozilla-ai/any-llm-go/providers"
	"github.com/mozilla-ai/any-llm-go/providers/anthropic"
	"github.com/mozilla-ai/any-llm-go/providers/ollama"
	"github.com/mozilla-ai/any-llm-go/providers/openai"
	"trankilou.fr/lassistanoque/backend/internal/domain"
)

var providerTypes = []domain.Item{
	{Value: "anthropic", Label: "Anthropic"},
	{Value: "openai", Label: "OpenAI"},
	{Value: "ollama", Label: "Ollama"},
	{Value: "openaicomp", Label: "OpenAI compatible"},
	{Value: "openrouter", Label: "Openrouter"},
}

type AnyLLMEngine struct {
}

func NewAnyLLMEngine() *AnyLLMEngine {
	return &AnyLLMEngine{}
}

func (e *AnyLLMEngine) ListProviderTypes() []domain.Item {
	return providerTypes
}

func providerFactory(provider *domain.Provider) (anyllm.Provider, error) {
	switch provider.Type {
	case "ollama":
		return ollama.New(
			anyllm.WithBaseURL(provider.URL),
		)
	case "openai":
		return openai.New(
			anyllm.WithAPIKey(provider.APIKey),
		)
	case "openaicomp":
		return openai.New(
			anyllm.WithBaseURL(provider.URL),
			anyllm.WithAPIKey(provider.APIKey),
		)
	case "anthropic":
		return anthropic.New(
			anyllm.WithAPIKey(provider.APIKey),
		)
	}
	return nil, fmt.Errorf("unknown provider type: %s", provider.Type)
}

func (e *AnyLLMEngine) ListModelsFromProvider(ctx context.Context, provider *domain.Provider) ([]string, error) {

	models := make([]string, 0)

	prov, err := providerFactory(provider)
	if err != nil {
		return nil, err
	}
	if lister, ok := prov.(anyllm.ModelLister); ok {
		response, err := lister.ListModels(ctx)
		if err != nil {
			return nil, err
		}
		for _, m := range response.Data {
			models = append(models, m.ID)
		}
	} else {
		models = append(models, "Default")
	}
	return models, nil
}

func (e *AnyLLMEngine) Stream(
	ctx context.Context,
	provider *domain.Provider,
	modelID string,
	params *domain.LLMParams,
	messages []*domain.Message,
) (*domain.Message, error) {

	p, err := providerFactory(provider)
	if err != nil {
		return nil, err
	}

	anyllmMessages := make([]anyllm.Message, 0)
	for _, m := range messages {
		anyllmMessages = append(anyllmMessages, anyllm.Message{
			Role:    m.Role,
			Content: m.Content,
		})
	}

	if params.ReasoningEffort == "" {
		params.ReasoningEffort = "none"
	}

	chunkChan, errChan := p.CompletionStream(ctx, anyllm.CompletionParams{
		Model:           modelID,
		Messages:        anyllmMessages,
		Stream:          true,
		ReasoningEffort: providers.ReasoningEffort(params.ReasoningEffort),
	})

	fullContent := ""
	role := anyllm.RoleAssistant

	for chunk := range chunkChan {
		if len(chunk.Choices) > 0 {
			content := chunk.Choices[0].Delta.Content
			reasoning := chunk.Choices[0].Delta.Reasoning
			if params.OnChunk != nil {
				if content != "" {
					fullContent += content
					params.OnChunk(&domain.Chunk{
						Done:      false,
						Role:      role,
						Content:   content,
						Reasoning: false,
					})
				} else if reasoning != nil {
					params.OnChunk(&domain.Chunk{
						Done:      false,
						Role:      role,
						Content:   reasoning.Content,
						Reasoning: true,
					})
				}
			}
		}
	}

	params.OnChunk(&domain.Chunk{
		Done:    true,
		Role:    role,
		Content: "",
	})

	if err := <-errChan; err != nil {
		return nil, err
	}

	return &domain.Message{
		Role:    role,
		Content: fullContent,
	}, nil
}