package llm

import (
	"context"
	"fmt"
	"os"

	openai "github.com/sashabaranov/go-openai"
)

var providers = map[string]string{
	"deepseek": "https://api.deepseek.com/v1",
	"kimi":     "https://api.moonshot.cn/v1",
	"ollama":   "http://localhost:11434/v1",
	"openai":   "https://api.openai.com/v1",
}

var defaultModels = map[string]string{
	"deepseek": "deepseek-chat",
	"kimi":     "moonshot-v1-8k",
	"ollama":   "qwen2.5",
	"openai":   "gpt-4o",
}

type Client struct {
	c     *openai.Client
	model string
}

func New(provider, model, baseURL, apiKeyEnv string) (*Client, error) {
	if baseURL == "" {
		var ok bool
		baseURL, ok = providers[provider]
		if !ok {
			baseURL = providers["deepseek"]
		}
	}
	if model == "" {
		model = defaultModels[provider]
		if model == "" {
			model = "deepseek-chat"
		}
	}
	apiKey := os.Getenv(apiKeyEnv)
	if apiKey == "" {
		apiKey = "ollama" // ollama doesn't need a real key
	}
	cfg := openai.DefaultConfig(apiKey)
	cfg.BaseURL = baseURL
	return &Client{c: openai.NewClientWithConfig(cfg), model: model}, nil
}

type Message = openai.ChatCompletionMessage

func NewMsg(role, content string) Message {
	return Message{Role: role, Content: content}
}

// Stream calls the LLM and streams tokens to the callback. Returns full response.
func (c *Client) Stream(ctx context.Context, msgs []Message, onToken func(string)) (string, error) {
	req := openai.ChatCompletionRequest{
		Model:    c.model,
		Messages: msgs,
		Stream:   true,
	}
	stream, err := c.c.CreateChatCompletionStream(ctx, req)
	if err != nil {
		return "", fmt.Errorf("llm stream: %w", err)
	}
	defer stream.Close()

	var full string
	for {
		resp, err := stream.Recv()
		if err != nil {
			break
		}
		delta := resp.Choices[0].Delta.Content
		full += delta
		if onToken != nil {
			onToken(delta)
		}
	}
	return full, nil
}