package llm

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"strings"
	"time"

	"github.com/openai/openai-go"
	"github.com/openai/openai-go/option"
	"github.com/openai/openai-go/packages/respjson"
	"github.com/openai/openai-go/shared"
	"github.com/user/keen-code/internal/config"
	"github.com/user/keen-code/internal/tools"
)

type chatStream interface {
	Next() bool
	Current() openai.ChatCompletionChunk
	Err() error
	Close() error
}

type streamFactory func(ctx context.Context, params openai.ChatCompletionNewParams, opts ...option.RequestOption) chatStream

type sdkChatStream struct {
	stream interface {
		Next() bool
		Current() openai.ChatCompletionChunk
		Err() error
		Close() error
	}
}

func (s *sdkChatStream) Next() bool {
	return s.stream.Next()
}

func (s *sdkChatStream) Current() openai.ChatCompletionChunk {
	return s.stream.Current()
}

func (s *sdkChatStream) Err() error {
	return s.stream.Err()
}

func (s *sdkChatStream) Close() error {
	return s.stream.Close()
}

type OpenAICompatibleClient struct {
	provider                Provider
	model                   string
	thinkingEffort          string
	maxRetries              int
	client                  openai.Client
	streamImpl              streamFactory
	pendingState            []openai.ChatCompletionMessageParamUnion
	contextWindowTokenCount int
}

func NewOpenAICompatibleClient(cfg *ClientConfig) (*OpenAICompatibleClient, error) {
	baseURL := cfg.BaseURL
	if baseURL == "" {
		var err error
		baseURL, err = openAICompatibleBaseURL(cfg.Provider)
		if err != nil {
			return nil, err
		}
	}

	client := openai.NewClient(
		option.WithAPIKey(cfg.APIKey),
		option.WithBaseURL(baseURL),
	)

	c := &OpenAICompatibleClient{
		provider:                cfg.Provider,
		model:                   cfg.Model,
		thinkingEffort:          cfg.ThinkingEffort,
		maxRetries:              retryCount(cfg.MaxRetries),
		client:                  client,
		contextWindowTokenCount: cfg.ContextWindowTokens,
	}
	c.streamImpl = func(ctx context.Context, params openai.ChatCompletionNewParams, opts ...option.RequestOption) chatStream {
		return &sdkChatStream{stream: c.client.Chat.Completions.NewStreaming(ctx, params, opts...)}
	}

	return c, nil
}

func openAICompatibleBaseURL(provider Provider) (string, error) {
	switch provider {
	case Provider(config.ProviderDeepSeek):
		return deepSeekBaseURL, nil
	case Provider(config.ProviderMoonshotAI):
		return moonshotAIBaseURL, nil
	case Provider(config.ProviderZAI):
		return zaiBaseURL, nil
	case Provider(config.ProviderOpenCodeGo):
		return openCodeGoBaseURL + "/v1/", nil
	default:
		return "", fmt.Errorf("unsupported OpenAI-compatible provider: %s", provider)
	}
}

func toOpenAIMessages(messages []Message) []openai.ChatCompletionMessageParamUnion {
	result := make([]openai.ChatCompletionMessageParamUnion, 0, len(messages))
	for _, m := range messages {
		content := FormatMessageForProvider(m)
		switch m.Role {
		case RoleSystem:
			result = append(result, openai.SystemMessage(content))
		case RoleUser:
			result = append(result, openai.UserMessage(content))
		case RoleAssistant:
			am := openai.ChatCompletionAssistantMessageParam{}
			am.Content.OfString = openai.String(content)
			result = append(result, openai.ChatCompletionMessageParamUnion{
				OfAssistant: &am,
			})
		}
	}
	return result
}

func toOpenAITools(registry *tools.Registry) []openai.ChatCompletionToolParam {
	if registry == nil {
		return nil
	}

	all := registry.All()
	result := make([]openai.ChatCompletionToolParam, 0, len(all))
	for _, t := range all {
		result = append(result, openai.ChatCompletionToolParam{
			Function: openai.FunctionDefinitionParam{
				Name:        t.Name(),
				Description: openai.String(t.Description()),
				Parameters:  openai.FunctionParameters(t.InputSchema()),
				Strict:      openai.Bool(false),
			},
		})
	}

	if len(result) == 0 {
		return nil
	}
	return result
}

func extractJSONStringField(extra map[string]respjson.Field, key string) string {
	if len(extra) == 0 {
		return ""
	}
	field, ok := extra[key]
	if !ok {
		return ""
	}
	raw := field.Raw()
	if raw == "" || raw == respjson.Null {
		return ""
	}

	var value string
	if err := json.Unmarshal([]byte(raw), &value); err == nil {
		return value
	}
	return ""
}

func extractReasoningDelta(extra map[string]respjson.Field) string {
	for _, key := range []string{"reasoning_content", "reasoning", "reasoning_text"} {
		if value := extractJSONStringField(extra, key); value != "" {
			return value
		}
	}
	return ""
}

func emitChunk(eventCh chan<- StreamEvent, content string) {
	if content == "" {
		return
	}
	eventCh <- StreamEvent{
		Type:    StreamEventTypeChunk,
		Content: content,
	}
}

const defaultMaxRetries = 10

func retryCount(maxRetries int) int {
	if maxRetries <= 0 {
		return defaultMaxRetries
	}
	return maxRetries
}

func isRetryableError(err error) bool {
	if err == nil {
		return false
	}
	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
		return false
	}
	var apiErr *openai.Error
	if errors.As(err, &apiErr) {
		switch apiErr.StatusCode {
		case http.StatusTooManyRequests,
			http.StatusInternalServerError,
			http.StatusBadGateway,
			http.StatusServiceUnavailable,
			http.StatusGatewayTimeout:
			return true
		default:
			return false
		}
	}
	return true
}

func (c *OpenAICompatibleClient) buildAssistantMessage(message openai.ChatCompletionMessage, reasoningContent string) openai.ChatCompletionAssistantMessageParam {
	assistant := openai.ChatCompletionAssistantMessageParam{}
	if message.Content != "" {
		assistant.Content.OfString = openai.String(message.Content)
	}
	if len(message.ToolCalls) > 0 {
		assistant.ToolCalls = make([]openai.ChatCompletionMessageToolCallParam, len(message.ToolCalls))
		for i, tc := range message.ToolCalls {
			assistant.ToolCalls[i] = openai.ChatCompletionMessageToolCallParam{
				ID: tc.ID,
				Function: openai.ChatCompletionMessageToolCallFunctionParam{
					Name:      tc.Function.Name,
					Arguments: tc.Function.Arguments,
				},
			}
		}
	}
	if reasoningContent != "" {
		assistant.SetExtraFields(map[string]any{
			"reasoning_content": reasoningContent,
		})
	}
	return assistant
}

func emitMissingFinalContent(
	eventCh chan<- StreamEvent,
	fullContent string,
	streamedContent string,
) {
	if fullContent == "" {
		return
	}

	// We stream delta content live as it arrives. At stream end, the accumulator
	// also exposes the full final content. Emit only the missing tail to avoid
	// duplicate UI text while still handling providers that send little/no deltas.
	if strings.HasPrefix(fullContent, streamedContent) {
		if tail := fullContent[len(streamedContent):]; tail != "" {
			emitChunk(eventCh, tail)
		}
		return
	}

	if streamedContent == "" {
		emitChunk(eventCh, fullContent)
	}
}

func (c *OpenAICompatibleClient) shouldLogRawChunks() bool {
	return c.provider == Provider(config.ProviderOpenCodeGo) && isOpenCodeGoKimiModel(c.model)
}

func (c *OpenAICompatibleClient) collectTurn(
	ctx context.Context,
	params openai.ChatCompletionNewParams,
	eventCh chan<- StreamEvent,
	requestOpts ...option.RequestOption,
) (openai.ChatCompletionMessage, string, string, bool, openai.CompletionUsage, error) {
	stream := c.streamImpl(ctx, params, requestOpts...)
	var acc openai.ChatCompletionAccumulator
	var reasoningContent strings.Builder
	var streamedContent strings.Builder

	for stream.Next() {
		chunk := stream.Current()
		if c.shouldLogRawChunks() {
			slog.Debug("OpenCode Go Kimi stream chunk", "chunk", chunk.RawJSON())
		}
		acc.AddChunk(chunk)

		if len(chunk.Choices) == 0 {
			continue
		}

		delta := chunk.Choices[0].Delta
		if delta.Content != "" {
			streamedContent.WriteString(delta.Content)
			emitChunk(eventCh, delta.Content)
		}

		// reasoning_content/reasoning are OpenAI-compatible extensions not modeled by openai-go.
		// Capture it during streaming because the SDK accumulator does not retain JSON metadata.
		reasoningDelta := extractReasoningDelta(delta.JSON.ExtraFields)
		reasoningContent.WriteString(reasoningDelta)
		if reasoningDelta != "" {
			eventCh <- StreamEvent{
				Type:    StreamEventTypeReasoningChunk,
				Content: reasoningDelta,
			}
		}
	}
	_ = stream.Close()

	if err := stream.Err(); err != nil {
		return openai.ChatCompletionMessage{}, "", "", false, openai.CompletionUsage{}, fmt.Errorf("stream error: %w", err)
	}
	if len(acc.ChatCompletion.Choices) == 0 {
		return openai.ChatCompletionMessage{}, "", "", false, openai.CompletionUsage{}, nil
	}

	return acc.ChatCompletion.Choices[0].Message, reasoningContent.String(), streamedContent.String(), true, acc.ChatCompletion.Usage, nil
}

func (c *OpenAICompatibleClient) collectTurnWithRetry(
	ctx context.Context,
	params openai.ChatCompletionNewParams,
	eventCh chan<- StreamEvent,
	requestOpts ...option.RequestOption,
) (openai.ChatCompletionMessage, string, string, bool, openai.CompletionUsage, error) {
	maxRetries := retryCount(c.maxRetries)
	for attempt := 1; attempt <= maxRetries; attempt++ {
		message, reasoningContent, streamedContent, hasChoice, usage, err := c.collectTurn(ctx, params, eventCh, requestOpts...)
		if err == nil {
			return message, reasoningContent, streamedContent, hasChoice, usage, nil
		}
		if !isRetryableError(err) || attempt == maxRetries {
			return openai.ChatCompletionMessage{}, "", "", false, openai.CompletionUsage{}, err
		}

		backoff := time.Duration(attempt) * time.Second
		slog.Debug("LLM stream error, retrying", "attempt", attempt, "maxRetries", maxRetries, "backoff", backoff, "error", err)
		eventCh <- StreamEvent{Type: StreamEventTypeRetry, Error: err, Attempt: attempt}

		select {
		case <-ctx.Done():
			return openai.ChatCompletionMessage{}, "", "", false, openai.CompletionUsage{}, ctx.Err()
		case <-time.After(backoff):
		}
	}
	return openai.ChatCompletionMessage{}, "", "", false, openai.CompletionUsage{}, nil
}

func (c *OpenAICompatibleClient) injectPendingState(oaiMessages []openai.ChatCompletionMessageParamUnion) ([]openai.ChatCompletionMessageParamUnion, []openai.ChatCompletionMessageParamUnion) {
	if len(c.pendingState) == 0 {
		return oaiMessages, nil
	}

	injectedPending := append([]openai.ChatCompletionMessageParamUnion(nil), c.pendingState...)

	slog.Debug("Injecting pending state", "pending_messages", len(c.pendingState), "total_messages", len(oaiMessages))
	if prettyJSON, err := json.MarshalIndent(c.pendingState, "", "  "); err == nil {
		slog.Debug("Pending state contents:\n" + string(prettyJSON))
	}

	if len(oaiMessages) > 0 {
		last := oaiMessages[len(oaiMessages)-1]
		oaiMessages = append(oaiMessages[:len(oaiMessages)-1], injectedPending...)
		oaiMessages = append(oaiMessages, last)
	} else {
		oaiMessages = append(oaiMessages, injectedPending...)
	}
	c.pendingState = nil
	return oaiMessages, injectedPending
}

func (c *OpenAICompatibleClient) buildChatParams(oaiMessages []openai.ChatCompletionMessageParamUnion, oaiTools []openai.ChatCompletionToolParam) openai.ChatCompletionNewParams {
	params := openai.ChatCompletionNewParams{
		Model:    c.model,
		Messages: oaiMessages,
		StreamOptions: openai.ChatCompletionStreamOptionsParam{
			IncludeUsage: openai.Bool(true),
		},
	}
	if len(oaiTools) > 0 {
		params.Tools = oaiTools
	}
	if (c.provider == Provider(config.ProviderDeepSeek) || (c.provider == Provider(config.ProviderOpenCodeGo) && isOpenCodeGoDeepSeekModel(c.model))) && c.thinkingEffort != "" {
		if c.thinkingEffort == "off" {
			params.SetExtraFields(map[string]any{
				"thinking": map[string]any{
					"type": "disabled",
				},
			})
		} else {
			params.ReasoningEffort = shared.ReasoningEffort(c.thinkingEffort)
			params.SetExtraFields(map[string]any{
				"thinking": map[string]any{
					"type": "enabled",
				},
			})
		}
	}
	if (c.provider == Provider(config.ProviderZAI) || (c.provider == Provider(config.ProviderOpenCodeGo) && (isOpenCodeGoGLMModel(c.model) || isOpenCodeGoKimiModel(c.model)))) && c.thinkingEffort != "" {
		params.SetExtraFields(map[string]any{
			"thinking": map[string]any{
				"type": c.thinkingEffort,
			},
		})
	}
	if c.provider == Provider(config.ProviderOpenCodeGo) && isOpenCodeGoQwenModel(c.model) && c.thinkingEffort != "" {
		switch c.thinkingEffort {
		case "enabled":
			params.SetExtraFields(map[string]any{"enable_thinking": true})
		case "disabled":
			params.SetExtraFields(map[string]any{"enable_thinking": false})
		}
	}
	return params
}

func (c *OpenAICompatibleClient) exitIncomplete(eventCh chan<- StreamEvent, oaiMessages []openai.ChatCompletionMessageParamUnion, turnStartLen int, injectedPending []openai.ChatCompletionMessageParamUnion, err error, oneShot bool) {
	if !oneShot {
		c.savePendingIfAccumulated(oaiMessages, turnStartLen, injectedPending)
	}
	c.emitTerminalEvent(eventCh, oaiMessages, turnStartLen, injectedPending, err)
}

func (c *OpenAICompatibleClient) StreamChat(
	ctx context.Context,
	messages []Message,
	toolRegistry *tools.Registry,
	opts ...StreamOptions,
) (<-chan StreamEvent, error) {
	eventCh := make(chan StreamEvent)
	streamOpts := streamOptions(opts)

	go func() {
		defer close(eventCh)

		oaiMessages := toOpenAIMessages(messages)
		oneShot := streamOpts.OneShot
		var injectedPending []openai.ChatCompletionMessageParamUnion
		if !oneShot {
			oaiMessages, injectedPending = c.injectPendingState(oaiMessages)
		}

		turnStartLen := len(oaiMessages)

		if prettyJSON, err := json.MarshalIndent(oaiMessages, "", "  "); err == nil {
			slog.Debug("OpenAI messages", "messages", string(prettyJSON))
		}
		oaiTools := toOpenAITools(toolRegistry)
		requestOpts := c.requestOptions(streamOpts)
		lastInputTokenCount := 0

		for range maxToolTurns {
			params := c.buildChatParams(oaiMessages, oaiTools)

			message, reasoningContent, streamedContent, hasChoice, usage, err := c.collectTurnWithRetry(ctx, params, eventCh, requestOpts...)
			if err != nil {
				c.exitIncomplete(eventCh, oaiMessages, turnStartLen, injectedPending, err, oneShot)
				return
			}

			if !hasChoice {
				c.exitIncomplete(eventCh, oaiMessages, turnStartLen, injectedPending, nil, oneShot)
				return
			}
			if usage.PromptTokens > 0 {
				lastInputTokenCount = int(usage.PromptTokens)
			}
			if usage.PromptTokens > 0 || usage.CompletionTokens > 0 {
				eventCh <- StreamEvent{
					Type: StreamEventTypeUsage,
					Usage: &TokenUsage{
						InputTokens:     int(usage.PromptTokens),
						OutputTokens:    int(usage.CompletionTokens),
						TotalTokens:     int(usage.TotalTokens),
						CachedTokens:    int(usage.PromptTokensDetails.CachedTokens),
						ReasoningTokens: int(usage.CompletionTokensDetails.ReasoningTokens),
					},
				}
			}
			emitMissingFinalContent(eventCh, message.Content, streamedContent)
			assistant := c.buildAssistantMessage(message, reasoningContent)

			if len(message.ToolCalls) == 0 {
				eventCh <- StreamEvent{Type: StreamEventTypeDone}
				return
			}

			oaiMessages = append(oaiMessages, openai.ChatCompletionMessageParamUnion{
				OfAssistant: &assistant,
			})

			toolMsgs := c.executeTools(ctx, message.ToolCalls, toolRegistry, eventCh)
			if len(toolMsgs) > 0 {
				oaiMessages = append(oaiMessages, toolMsgs...)
			}

			if lastInputTokenCount > 0 && !contextFitsBudget(c.contextWindowTokenCount, lastInputTokenCount) {
				reducedMessages, reduction := reduceOpenAIContextForRequest(c.contextWindowTokenCount, lastInputTokenCount, oaiMessages)
				if !reduction.FitsBudget {
					slog.Debug("OpenAI context still exceeds budget after reduction", "inputTokenCount", reduction.ReducedTokenCount, "removedToolResultCount", reduction.RemovedToolResults)
					c.exitIncomplete(eventCh, oaiMessages, turnStartLen, injectedPending, fmt.Errorf(contextWindowExceededError), oneShot)
					return
				}
				oaiMessages = reducedMessages
			}
		}

		c.exitIncomplete(eventCh, oaiMessages, turnStartLen, injectedPending, nil, oneShot)
	}()

	return eventCh, nil
}

func (c *OpenAICompatibleClient) requestOptions(opts StreamOptions) []option.RequestOption {
	if c.provider != Provider(config.ProviderOpenCodeGo) || opts.SessionID == "" {
		return nil
	}
	return []option.RequestOption{
		option.WithHeader("x-opencode-session", opencodeSessionID(opts.SessionID)),
	}
}

func (c *OpenAICompatibleClient) Reset() {
	c.pendingState = nil
}

func (c *OpenAICompatibleClient) savePendingIfAccumulated(oaiMessages []openai.ChatCompletionMessageParamUnion, turnStartLen int, injectedPending []openai.ChatCompletionMessageParamUnion) {
	if len(injectedPending) == 0 && len(oaiMessages) <= turnStartLen {
		return
	}

	newDelta := []openai.ChatCompletionMessageParamUnion(nil)
	if len(oaiMessages) > turnStartLen {
		newDelta = oaiMessages[turnStartLen:]
	}

	c.pendingState = make([]openai.ChatCompletionMessageParamUnion, 0, len(injectedPending)+len(newDelta))
	c.pendingState = append(c.pendingState, injectedPending...)
	c.pendingState = append(c.pendingState, newDelta...)
}

func (c *OpenAICompatibleClient) emitTerminalEvent(eventCh chan<- StreamEvent, oaiMessages []openai.ChatCompletionMessageParamUnion, turnStartLen int, injectedPending []openai.ChatCompletionMessageParamUnion, err error) {
	if len(injectedPending) > 0 || len(oaiMessages) > turnStartLen {
		eventCh <- StreamEvent{Type: StreamEventTypeIncomplete, Error: err}
	} else if err != nil {
		eventCh <- StreamEvent{Type: StreamEventTypeError, Error: err}
	} else {
		eventCh <- StreamEvent{Type: StreamEventTypeDone}
	}
}

func (c *OpenAICompatibleClient) executeTools(
	ctx context.Context,
	toolCalls []openai.ChatCompletionMessageToolCall,
	registry *tools.Registry,
	eventCh chan<- StreamEvent,
) []openai.ChatCompletionMessageParamUnion {
	toolMessages := make([]openai.ChatCompletionMessageParamUnion, 0, len(toolCalls))

	for _, tc := range toolCalls {
		start := time.Now()
		input := map[string]any{}
		if tc.Function.Arguments != "" {
			if err := json.Unmarshal([]byte(tc.Function.Arguments), &input); err != nil {
				input = map[string]any{}
			}
		}
		slog.Debug("Tool request", "tool", tc.Function.Name, "input", input)
		eventCh <- StreamEvent{
			Type: StreamEventTypeToolStart,
			ToolCall: &ToolCall{
				Name:  tc.Function.Name,
				Input: input,
			},
		}

		var output any
		var execErr error

		if registry == nil {
			execErr = fmt.Errorf("tool registry not available")
		} else if tool, exists := registry.Get(tc.Function.Name); !exists {
			execErr = fmt.Errorf("tool %q not found", tc.Function.Name)
		} else {
			output, execErr = tool.Execute(ctx, input)
		}

		duration := time.Since(start)
		toolCall := &ToolCall{
			Name:     tc.Function.Name,
			Input:    input,
			Output:   output,
			Duration: duration,
		}

		var toolOutput string
		if execErr != nil {
			toolCall.Error = execErr.Error()
			slog.Debug("Tool response", "tool", tc.Function.Name, "error", execErr.Error(), "duration", duration)
			eventCh <- StreamEvent{
				Type:     StreamEventTypeToolEnd,
				ToolCall: toolCall,
			}
			toolOutput = fmt.Sprintf(`{"error":%q}`, execErr.Error())
		} else {
			slog.Debug("Tool response", "tool", tc.Function.Name, "duration", duration)
			eventCh <- StreamEvent{
				Type:     StreamEventTypeToolEnd,
				ToolCall: toolCall,
			}
			if output == nil {
				output = map[string]any{}
			}
			b, err := json.Marshal(output)
			if err != nil {
				toolOutput = "{}"
			} else {
				toolOutput = string(b)
			}
		}

		toolMessages = append(toolMessages, openai.ToolMessage(toolOutput, tc.ID))
	}

	return toolMessages
}
