package llm

import (
	"context"
	"fmt"
	"os"

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

	coretypes "github.com/enmanuel/agents/pkg/llm"
)

// NewOpenAIComplete returns a CompleteFunc backed by the OpenAI-compatible API.
// Works with OpenAI, Ollama, vLLM, LMStudio — just change baseURL.
func NewOpenAIComplete(apiKeyEnv, baseURL string) coretypes.CompleteFunc {
	return func(ctx context.Context, req coretypes.CompletionRequest) (coretypes.CompletionResponse, error) {
		apiKey := os.Getenv(apiKeyEnv)
		if apiKey == "" {
			apiKey = "ollama" // Ollama doesn't require a real key
		}

		cfg := openai.DefaultConfig(apiKey)
		if baseURL != "" {
			cfg.BaseURL = baseURL
		}
		client := openai.NewClientWithConfig(cfg)

		msgs := make([]openai.ChatCompletionMessage, 0, len(req.Messages)+1)
		if req.SystemPrompt != "" {
			msgs = append(msgs, openai.ChatCompletionMessage{
				Role:    openai.ChatMessageRoleSystem,
				Content: req.SystemPrompt,
			})
		}
		for _, m := range req.Messages {
			role := openai.ChatMessageRoleUser
			switch m.Role {
			case coretypes.RoleAssistant:
				role = openai.ChatMessageRoleAssistant
			case coretypes.RoleSystem:
				role = openai.ChatMessageRoleSystem
			case coretypes.RoleTool:
				role = openai.ChatMessageRoleTool
			}
			msgs = append(msgs, openai.ChatCompletionMessage{
				Role:    role,
				Content: m.Content,
			})
		}

		openReq := openai.ChatCompletionRequest{
			Model:       req.Model,
			Messages:    msgs,
			MaxTokens:   req.MaxTokens,
			Temperature: float32(req.Temperature),
		}

		resp, err := client.CreateChatCompletion(ctx, openReq)
		if err != nil {
			return coretypes.CompletionResponse{}, fmt.Errorf("openai completion: %w", err)
		}
		if len(resp.Choices) == 0 {
			return coretypes.CompletionResponse{}, fmt.Errorf("openai: empty choices")
		}

		return coretypes.CompletionResponse{
			Content:      resp.Choices[0].Message.Content,
			FinishReason: string(resp.Choices[0].FinishReason),
			Usage: coretypes.TokenUsage{
				InputTokens:  resp.Usage.PromptTokens,
				OutputTokens: resp.Usage.CompletionTokens,
				TotalTokens:  resp.Usage.TotalTokens,
			},
		}, nil
	}
}