package llm

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"os"
	"time"

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

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

// NewOpenAIComplete returns a CompleteFunc backed by the OpenAI-compatible API.
// Works with OpenAI, Ollama, vLLM, LMStudio — just change baseURL.
func NewOpenAIComplete(apiKeyEnv, baseURL string, log *slog.Logger) 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 {
			msgs = append(msgs, toOpenAIMessage(m))
		}

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

		// Add tools if present
		if len(req.Tools) > 0 {
			openReq.Tools = toOpenAITools(req.Tools)
		}

		log.Info("llm_request",
			"provider", "openai",
			"model", req.Model,
			"messages", len(req.Messages),
			"tools", len(req.Tools),
		)

		start := time.Now()
		resp, err := client.CreateChatCompletion(ctx, openReq)
		if err != nil {
			ms := time.Since(start).Milliseconds()
			log.Error("llm_error", "provider", "openai", logger.FieldDurationMS, ms, "err", err)
			return coretypes.CompletionResponse{}, fmt.Errorf("openai completion: %w", err)
		}
		ms := time.Since(start).Milliseconds()

		if len(resp.Choices) == 0 {
			log.Error("llm_error", "provider", "openai", logger.FieldDurationMS, ms, "err", "empty choices")
			return coretypes.CompletionResponse{}, fmt.Errorf("openai: empty choices")
		}

		choice := resp.Choices[0]
		var toolCalls []coretypes.ToolCall
		for _, tc := range choice.Message.ToolCalls {
			toolCalls = append(toolCalls, coretypes.ToolCall{
				ID:        tc.ID,
				Name:      tc.Function.Name,
				Arguments: tc.Function.Arguments,
			})
		}

		log.Info("llm_response",
			"provider", "openai",
			"model", req.Model,
			logger.FieldDurationMS, ms,
			logger.FieldTokensUsed, resp.Usage.TotalTokens,
			"input_tokens", resp.Usage.PromptTokens,
			"output_tokens", resp.Usage.CompletionTokens,
			"tool_calls", len(toolCalls),
			"finish_reason", string(choice.FinishReason),
		)

		return coretypes.CompletionResponse{
			Content:      choice.Message.Content,
			ToolCalls:    toolCalls,
			FinishReason: string(choice.FinishReason),
			Usage: coretypes.TokenUsage{
				InputTokens:  resp.Usage.PromptTokens,
				OutputTokens: resp.Usage.CompletionTokens,
				TotalTokens:  resp.Usage.TotalTokens,
			},
		}, nil
	}
}

// toOpenAIMessage converts a core Message to an OpenAI ChatCompletionMessage.
func toOpenAIMessage(m coretypes.Message) openai.ChatCompletionMessage {
	role := openai.ChatMessageRoleUser
	switch m.Role {
	case coretypes.RoleAssistant:
		role = openai.ChatMessageRoleAssistant
	case coretypes.RoleSystem:
		role = openai.ChatMessageRoleSystem
	case coretypes.RoleTool:
		role = openai.ChatMessageRoleTool
	}

	msg := openai.ChatCompletionMessage{
		Role:       role,
		Content:    m.Content,
		ToolCallID: m.ToolCallID,
	}

	// Assistant messages with tool calls
	if m.Role == coretypes.RoleAssistant && len(m.ToolCalls) > 0 {
		msg.ToolCalls = make([]openai.ToolCall, len(m.ToolCalls))
		for i, tc := range m.ToolCalls {
			msg.ToolCalls[i] = openai.ToolCall{
				ID:   tc.ID,
				Type: openai.ToolTypeFunction,
				Function: openai.FunctionCall{
					Name:      tc.Name,
					Arguments: tc.Arguments,
				},
			}
		}
	}

	return msg
}

// toOpenAITools converts core ToolSpecs to OpenAI Tool format.
func toOpenAITools(specs []coretypes.ToolSpec) []openai.Tool {
	tools := make([]openai.Tool, len(specs))
	for i, s := range specs {
		tools[i] = openai.Tool{
			Type: openai.ToolTypeFunction,
			Function: &openai.FunctionDefinition{
				Name:        s.Name,
				Description: s.Description,
				Parameters:  json.RawMessage(marshalSchema(s.InputSchema)),
			},
		}
	}
	return tools
}

// marshalSchema marshals a JSON schema map to bytes. Falls back to empty object.
func marshalSchema(schema map[string]any) []byte {
	b, err := json.Marshal(schema)
	if err != nil {
		return []byte("{}")
	}
	return b
}