package openai_compat

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	"go.uber.org/zap"

	"iop/apps/node/internal/runtime"
)

const (
	runtimeMetadataOpenAIToolCalls        = "openai_tool_calls"
	runtimeMetadataOpenAITextToolFallback = "openai_text_tool_fallback"
	streamTraceMetadataKey                = "iop_trace_stream"
	streamTraceEnvKey                     = "IOP_OPENAI_COMPAT_TRACE_STREAM"
)

// chatStreamSession owns the accumulated stream state (usage, finish reason,
// tool-call accumulator, counted output tokens, and trace sequence) while the
// SSE payloads are decoded into RuntimeEvents.
type chatStreamSession struct {
	adapter          *Adapter
	spec             runtime.ExecutionSpec
	textToolFallback bool
	traceStream      bool
	traceSeq         int
	outputTokens     int
	finishReason     string
	usage            *runtime.UsageStats
	toolCalls        openAIToolCallAccumulator
}

// consumeStream decodes the chat completions SSE stream, emitting reasoning and
// content deltas and, on [DONE], the completion event. The scan/emit ordering
// and terminal error handling match the prior single-function Execute.
func (a *Adapter) consumeStream(ctx context.Context, spec runtime.ExecutionSpec, sink runtime.EventSink, resp *http.Response, textToolFallback bool) error {
	defer resp.Body.Close()
	session := &chatStreamSession{
		adapter:          a,
		spec:             spec,
		textToolFallback: textToolFallback,
		traceStream:      openAICompatTraceStreamEnabled(spec.Metadata),
	}
	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		line := scanner.Text()
		if !strings.HasPrefix(line, "data: ") {
			continue
		}
		payload := strings.TrimPrefix(line, "data: ")
		if payload == "[DONE]" {
			return session.emitComplete(ctx, sink)
		}
		if err := session.handlePayload(ctx, sink, payload); err != nil {
			return err
		}
	}
	if err := scanner.Err(); err != nil {
		_ = emitError(ctx, sink, spec.RunID, fmt.Sprintf("openai_compat stream scan failed: %v", err))
		return fmt.Errorf("openai_compat adapter: scan stream: %w", err)
	}
	_ = emitError(ctx, sink, spec.RunID, "openai_compat stream ended without [DONE]")
	return fmt.Errorf("openai_compat adapter: stream ended without [DONE]")
}

// emitComplete emits the terminal completion event with the accumulated finish
// reason, usage, counted output tokens, and tool calls.
func (s *chatStreamSession) emitComplete(ctx context.Context, sink runtime.EventSink) error {
	if s.traceStream {
		s.adapter.logger.Info("openai_compat provider stream done",
			zap.String("run_id", s.spec.RunID),
			zap.Int("seq", s.traceSeq+1),
		)
	}
	return sink.Emit(ctx, completeEvent(s.spec.RunID, s.finishReason, s.usage, s.outputTokens, s.toolCalls.ToolCalls(), s.textToolFallback))
}

// applyUsage records the usage stats from a chunk when present.
func (s *chatStreamSession) applyUsage(chunk chatChunk) {
	if chunk.Usage == nil {
		return
	}
	s.usage = &runtime.UsageStats{
		InputTokens:  chunk.Usage.PromptTokens,
		OutputTokens: chunk.Usage.CompletionTokens,
	}
	if d := chunk.Usage.PromptTokensDetails; d != nil {
		s.usage.CachedInputTokens = d.CachedTokens
	}
	if d := chunk.Usage.CompletionTokensDetails; d != nil {
		s.usage.ReasoningTokens = d.ReasoningTokens
	}
}

// handlePayload decodes one SSE data payload, updating usage/finish/tool-call
// state and emitting reasoning and content deltas. A malformed chunk or a
// choice-less chunk is skipped, matching the prior continue semantics.
func (s *chatStreamSession) handlePayload(ctx context.Context, sink runtime.EventSink, payload string) error {
	a := s.adapter
	s.traceSeq++
	if s.traceStream {
		a.logger.Info("openai_compat provider stream raw chunk",
			zap.String("run_id", s.spec.RunID),
			zap.Int("seq", s.traceSeq),
			zap.Int("raw_len", len(payload)),
			zap.Bool("raw_has_open_think", hasOpenThinkTag(payload)),
			zap.Bool("raw_has_close_think", hasCloseThinkTag(payload)),
			zap.String("raw_preview", tracePreview(payload, 4000)),
		)
	}
	var chunk chatChunk
	if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
		if s.traceStream {
			a.logger.Warn("openai_compat provider stream raw chunk ignored",
				zap.String("run_id", s.spec.RunID),
				zap.Int("seq", s.traceSeq),
				zap.Error(err),
			)
		}
		return nil
	}
	s.applyUsage(chunk)
	if len(chunk.Choices) == 0 {
		return nil
	}
	choice := chunk.Choices[0]
	if choice.FinishReason != nil && *choice.FinishReason != "" {
		s.finishReason = *choice.FinishReason
	}
	if len(choice.Delta.ToolCalls) > 0 {
		s.toolCalls.AddDelta(choice.Delta.ToolCalls)
	}
	reasoning := choice.Delta.ReasoningText()
	if s.traceStream {
		a.logger.Info("openai_compat provider stream parsed chunk",
			zap.String("run_id", s.spec.RunID),
			zap.Int("seq", s.traceSeq),
			zap.String("finish_reason", s.finishReason),
			zap.Int("tool_call_delta_count", len(choice.Delta.ToolCalls)),
			zap.Int("content_len", len(choice.Delta.Content)),
			zap.Bool("content_has_open_think", hasOpenThinkTag(choice.Delta.Content)),
			zap.Bool("content_has_close_think", hasCloseThinkTag(choice.Delta.Content)),
			zap.String("content_preview", tracePreview(choice.Delta.Content, 2000)),
			zap.Int("reasoning_len", len(reasoning)),
			zap.Bool("reasoning_has_open_think", hasOpenThinkTag(reasoning)),
			zap.Bool("reasoning_has_close_think", hasCloseThinkTag(reasoning)),
			zap.String("reasoning_preview", tracePreview(reasoning, 2000)),
		)
	}
	if reasoning != "" {
		if err := sink.Emit(ctx, runtime.RuntimeEvent{
			RunID:     s.spec.RunID,
			Type:      runtime.EventTypeReasoningDelta,
			Delta:     reasoning,
			Timestamp: time.Now(),
		}); err != nil {
			return err
		}
	}
	if content := choice.Delta.Content; content != "" {
		s.outputTokens += len(strings.Fields(content))
		if err := sink.Emit(ctx, runtime.RuntimeEvent{
			RunID:     s.spec.RunID,
			Type:      runtime.EventTypeDelta,
			Delta:     content,
			Timestamp: time.Now(),
		}); err != nil {
			return err
		}
	}
	return nil
}

func openAICompatTraceStreamEnabled(metadata map[string]string) bool {
	if traceBool(os.Getenv(streamTraceEnvKey)) {
		return true
	}
	if metadata == nil {
		return false
	}
	return traceBool(metadata[streamTraceMetadataKey])
}

func traceBool(v string) bool {
	switch strings.ToLower(strings.TrimSpace(v)) {
	case "1", "true", "yes", "on":
		return true
	default:
		return false
	}
}

func tracePreview(s string, max int) string {
	if max <= 0 || len(s) <= max {
		return s
	}
	return s[:max] + "...[truncated]"
}

func hasOpenThinkTag(s string) bool {
	return strings.Contains(strings.ToLower(s), "<think")
}

func hasCloseThinkTag(s string) bool {
	return strings.Contains(strings.ToLower(s), "</think")
}

// joinOpenAIPath appends an OpenAI-compatible path to the endpoint. When the
// endpoint already targets the /v1 root it strips the leading /v1 from the path
// so a configured ".../v1" endpoint does not produce ".../v1/v1/...".
func joinOpenAIPath(baseURL, path string) string {
	u, err := url.Parse(baseURL)
	if err != nil {
		return strings.TrimRight(baseURL, "/") + path
	}
	basePath := strings.TrimRight(u.Path, "/")
	if strings.HasSuffix(basePath, "/v1") && strings.HasPrefix(path, "/v1/") {
		path = strings.TrimPrefix(path, "/v1")
	}
	u.Path = basePath + path
	return u.String()
}

func readLimited(r io.Reader, limit int64) string {
	b, _ := io.ReadAll(io.LimitReader(r, limit))
	return strings.TrimSpace(string(b))
}

type chatMessage struct {
	Role       string `json:"role"`
	Content    string `json:"content"`
	ToolCalls  []any  `json:"tool_calls,omitempty"`
	ToolCallID string `json:"tool_call_id,omitempty"`
}

type chatChunk struct {
	Choices []struct {
		Delta        chatChunkDelta `json:"delta"`
		FinishReason *string        `json:"finish_reason"`
	} `json:"choices"`
	Usage *struct {
		PromptTokens        int `json:"prompt_tokens"`
		CompletionTokens    int `json:"completion_tokens"`
		PromptTokensDetails *struct {
			CachedTokens int `json:"cached_tokens"`
		} `json:"prompt_tokens_details"`
		CompletionTokensDetails *struct {
			ReasoningTokens int `json:"reasoning_tokens"`
		} `json:"completion_tokens_details"`
	} `json:"usage"`
}

type chatChunkDelta struct {
	Content          string `json:"content"`
	ReasoningContent string `json:"reasoning_content"`
	Reasoning        string `json:"reasoning"`
	ToolCalls        []any  `json:"tool_calls"`
}

func (d chatChunkDelta) ReasoningText() string {
	if d.ReasoningContent != "" {
		return d.ReasoningContent
	}
	return d.Reasoning
}

type openAIToolCallAccumulator struct {
	calls map[int]map[string]any
	order []int
}

func (a *openAIToolCallAccumulator) AddDelta(raw []any) {
	for fallbackIndex, item := range raw {
		call, ok := item.(map[string]any)
		if !ok {
			continue
		}
		index := fallbackIndex
		if parsed, ok := numericIndex(call["index"]); ok {
			index = parsed
		}
		dst := a.ensure(index)
		for key, value := range call {
			switch key {
			case "index":
				continue
			case "function":
				mergeToolCallFunction(dst, value)
			default:
				if !emptyToolCallValue(value) {
					dst[key] = value
				}
			}
		}
	}
}

func (a *openAIToolCallAccumulator) ToolCalls() []any {
	if len(a.order) == 0 {
		return nil
	}
	out := make([]any, 0, len(a.order))
	for _, index := range a.order {
		call := a.calls[index]
		if len(call) == 0 {
			continue
		}
		copyCall := make(map[string]any, len(call))
		for key, value := range call {
			copyCall[key] = value
		}
		out = append(out, copyCall)
	}
	return out
}

func (a *openAIToolCallAccumulator) ensure(index int) map[string]any {
	if a.calls == nil {
		a.calls = make(map[int]map[string]any)
	}
	if call, ok := a.calls[index]; ok {
		return call
	}
	call := make(map[string]any)
	a.calls[index] = call
	a.order = append(a.order, index)
	return call
}

func mergeToolCallFunction(call map[string]any, value any) {
	fn, ok := value.(map[string]any)
	if !ok {
		return
	}
	current, _ := call["function"].(map[string]any)
	if current == nil {
		current = make(map[string]any)
		call["function"] = current
	}
	for key, item := range fn {
		if key == "arguments" {
			if part, ok := item.(string); ok {
				current["arguments"] = currentString(current["arguments"]) + part
				continue
			}
		}
		if !emptyToolCallValue(item) {
			current[key] = item
		}
	}
}

func numericIndex(value any) (int, bool) {
	switch v := value.(type) {
	case int:
		return v, true
	case float64:
		return int(v), true
	default:
		return 0, false
	}
}

func emptyToolCallValue(value any) bool {
	if value == nil {
		return true
	}
	if text, ok := value.(string); ok {
		return text == ""
	}
	return false
}

func currentString(value any) string {
	if text, ok := value.(string); ok {
		return text
	}
	return ""
}

type modelsResponse struct {
	Data []struct {
		ID string `json:"id"`
	} `json:"data"`
}