package agent

import (
	"encoding/json"
	"fmt"
	"log/slog"
	"strings"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
	"github.com/tidbcloud/AgentTeam/desktop/internal/consts"
)

func (d *ClaudeDriver) rememberToolCallName(toolCallID, toolName string) {
	toolCallID = strings.TrimSpace(toolCallID)
	if toolCallID == "" {
		return
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.toolCallNames == nil {
		d.toolCallNames = make(map[string]string)
	}
	d.toolCallNames[toolCallID] = strings.TrimSpace(toolName)
}

func (d *ClaudeDriver) toolCallName(toolCallID string) string {
	toolCallID = strings.TrimSpace(toolCallID)
	if toolCallID == "" {
		return ""
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	return strings.TrimSpace(d.toolCallNames[toolCallID])
}

func (d *ClaudeDriver) rememberClaudeTranscriptBoundary(sessionID string, uuid string) {
	trimmedSessionID := strings.TrimSpace(sessionID)
	trimmedUUID := strings.TrimSpace(uuid)
	if trimmedSessionID == "" && trimmedUUID == "" {
		return
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	if trimmedSessionID != "" {
		d.sessionID = trimmedSessionID
	}
	if trimmedUUID != "" {
		d.lastTranscriptUUID = trimmedUUID
	}
}

func (d *ClaudeDriver) rememberClaudeReplayRecoveryBoundary(sessionID string, messageID string) {
	trimmedMessageID := strings.TrimSpace(messageID)
	if trimmedMessageID == "" {
		return
	}
	trimmedSessionID := strings.TrimSpace(sessionID)
	d.mu.Lock()
	defer d.mu.Unlock()
	if trimmedSessionID == "" {
		trimmedSessionID = strings.TrimSpace(d.sessionID)
	}
	if d.replayRecoveryBoundaries == nil {
		d.replayRecoveryBoundaries = make(map[string]claudeRecoveryBoundary)
	}
	d.replayRecoveryBoundaries[trimmedMessageID] = claudeRecoveryBoundary{
		sessionID:    trimmedSessionID,
		boundaryUUID: strings.TrimSpace(d.lastTranscriptUUID),
	}
}

func (d *ClaudeDriver) rememberSession(sessionID string) {
	trimmedSessionID := strings.TrimSpace(sessionID)
	if trimmedSessionID == "" {
		return
	}
	d.mu.Lock()
	d.sessionID = trimmedSessionID
	d.mu.Unlock()
}

func claudeToolResultText(raw json.RawMessage) string {
	raw = json.RawMessage(strings.TrimSpace(string(raw)))
	if len(raw) == 0 || string(raw) == "null" {
		return ""
	}
	var text string
	if err := json.Unmarshal(raw, &text); err == nil {
		return text
	}
	var blocks []struct {
		Type string `json:"type"`
		Text string `json:"text"`
	}
	if err := json.Unmarshal(raw, &blocks); err == nil {
		parts := make([]string, 0, len(blocks))
		for _, block := range blocks {
			if strings.TrimSpace(block.Text) != "" {
				parts = append(parts, block.Text)
			}
		}
		return strings.Join(parts, "\n")
	}
	var object struct {
		Text string `json:"text"`
	}
	if err := json.Unmarshal(raw, &object); err == nil && strings.TrimSpace(object.Text) != "" {
		return object.Text
	}
	return string(raw)
}

// ParseLine parses one Claude stream-json line into normalized events.

// ParseLine parses one Claude stream-json line into normalized events.
func (d *ClaudeDriver) ParseLine(line string) []drivercontract.ParsedEvent {
	if !strings.HasPrefix(strings.TrimSpace(line), "{") {
		return nil
	}

	var raw struct {
		Type            string  `json:"type"`
		Subtype         string  `json:"subtype"`
		IsError         bool    `json:"is_error"`
		SessionID       string  `json:"session_id"`
		UUID            string  `json:"uuid"`
		IsReplay        bool    `json:"isReplay"`
		ParentToolUseID *string `json:"parent_tool_use_id"`
		Message         struct {
			StopReason string `json:"stop_reason"`
			Content    []struct {
				Type      string          `json:"type"`
				ID        string          `json:"id"`
				Text      string          `json:"text"`
				Thinking  string          `json:"thinking"`
				Name      string          `json:"name"`
				Input     json.RawMessage `json:"input"`
				ToolUseID string          `json:"tool_use_id"`
				Content   json.RawMessage `json:"content"`
				IsError   bool            `json:"is_error"`
			} `json:"content"`
			Model string `json:"model"`
			Usage *struct {
				InputTokens              int `json:"input_tokens"`
				CacheReadInputTokens     int `json:"cache_read_input_tokens"`
				CacheCreationInputTokens int `json:"cache_creation_input_tokens"`
				OutputTokens             int `json:"output_tokens"`
			} `json:"usage"`
		} `json:"message"`
		// Top-level usage field: Claude Code emits usage at top level in result events.
		Usage *struct {
			InputTokens              int `json:"input_tokens"`
			CacheReadInputTokens     int `json:"cache_read_input_tokens"`
			CacheCreationInputTokens int `json:"cache_creation_input_tokens"`
			OutputTokens             int `json:"output_tokens"`
		} `json:"usage"`
		// Top-level modelUsage: maps "provider/model" to per-model stats.
		ModelUsage map[string]struct {
			InputTokens              int `json:"inputTokens"`
			OutputTokens             int `json:"outputTokens"`
			CacheReadInputTokens     int `json:"cacheReadInputTokens"`
			CacheCreationInputTokens int `json:"cacheCreationInputTokens"`
		} `json:"modelUsage"`
		// Top-level fields for error diagnostics.
		DurationMs    int             `json:"duration_ms"`
		DurationAPIMs int             `json:"duration_api_ms"`
		NumTurns      int             `json:"num_turns"`
		TotalCostUSD  float64         `json:"total_cost_usd"`
		Errors        []string        `json:"errors"`
		Result        json.RawMessage `json:"result"`
		Response      struct {
			Subtype   string `json:"subtype"`
			RequestID string `json:"request_id"`
			Error     string `json:"error"`
		} `json:"response"`
	}

	if err := json.Unmarshal([]byte(line), &raw); err != nil {
		return nil
	}
	if raw.Type == "user" && raw.IsReplay && strings.TrimSpace(raw.UUID) != "" {
		d.rememberClaudeReplayRecoveryBoundary(raw.SessionID, raw.UUID)
	}
	if raw.Type == "user" || raw.Type == "assistant" {
		d.rememberClaudeTranscriptBoundary(raw.SessionID, raw.UUID)
	} else {
		d.rememberSession(raw.SessionID)
	}

	var events []drivercontract.ParsedEvent
	switch raw.Type {
	case "control_response":
		if raw.Response.Subtype == "success" {
			d.completeClaudeControlWaiter(raw.Response.RequestID, nil)
		} else if raw.Response.Subtype != "" {
			errMsg := strings.TrimSpace(raw.Response.Error)
			if errMsg == "" {
				errMsg = "claude control request failed"
			}
			d.completeClaudeControlWaiter(raw.Response.RequestID, fmt.Errorf("%s", errMsg))
		}
	case "system":
		if raw.Subtype == "init" && raw.SessionID != "" {
			events = append(events, drivercontract.ParsedEvent{
				Kind:      consts.ParsedEventKindSessionInit,
				SessionID: raw.SessionID,
			})
		} else if raw.Subtype == "compact_boundary" {
			d.mu.Lock()
			if d.compactionInProgress {
				d.compactBoundaryObserved = true
			}
			d.mu.Unlock()
			events = append(events, drivercontract.ParsedEvent{
				Kind:      consts.ParsedEventKindRuntimeCompacting,
				SessionID: raw.SessionID,
			})
		}
	case "assistant":
		for _, block := range raw.Message.Content {
			switch block.Type {
			case consts.ParsedEventKindThinking:
				events = append(events, drivercontract.ParsedEvent{
					Kind: consts.ParsedEventKindThinking,
					Text: block.Thinking,
				})
			case consts.ParsedEventKindText:
				events = append(events, drivercontract.ParsedEvent{
					Kind: consts.ParsedEventKindText,
					Text: block.Text,
				})
			case consts.DriverEventTypeToolUse:
				inputStr := ""
				if block.Input != nil {
					inputStr = string(block.Input)
				}
				toolCallID := strings.TrimSpace(block.ID)
				if toolCallID != "" {
					d.rememberToolCallName(toolCallID, block.Name)
				}
				events = append(events, drivercontract.ParsedEvent{
					Kind:       consts.ParsedEventKindToolCall,
					ToolName:   block.Name,
					ToolCallID: toolCallID,
					ToolInput:  inputStr,
				})
			}
		}
	case "user":
		if raw.IsReplay && strings.TrimSpace(raw.UUID) != "" {
			// Claude documents replayed user messages as acknowledgements. Map
			// that strongest correlated stdio signal to AgentTeam's historical
			// runtime_input_acked ownership transition; it is not a
			// per-UUID proof that a model request consumed the body.
			events = append(events, drivercontract.ParsedEvent{
				Kind:           consts.ParsedEventKindRuntimeInputAcked,
				RuntimeInputID: strings.TrimSpace(raw.UUID),
				SessionID:      strings.TrimSpace(raw.SessionID),
			})
		}
		for _, block := range raw.Message.Content {
			if block.Type != "tool_result" {
				continue
			}
			toolCallID := strings.TrimSpace(block.ToolUseID)
			if toolCallID == "" && raw.ParentToolUseID != nil {
				toolCallID = strings.TrimSpace(*raw.ParentToolUseID)
			}
			text := strings.TrimSpace(claudeToolResultText(block.Content))
			if text == "" {
				continue
			}
			metadata := map[string]any(nil)
			if block.IsError {
				metadata = map[string]any{"toolError": true}
			}
			events = append(events, drivercontract.ParsedEvent{
				Kind:       consts.ParsedEventKindToolOutput,
				ToolName:   d.toolCallName(toolCallID),
				ToolCallID: toolCallID,
				Text:       text,
				Metadata:   metadata,
			})
		}
	case "result":
		d.mu.Lock()
		compactionInProgress := d.compactionInProgress
		compactBoundaryObserved := d.compactBoundaryObserved
		d.compactionInProgress = false
		d.compactBoundaryObserved = false
		d.mu.Unlock()
		abnormalResult := raw.IsError || len(raw.Errors) > 0 || claudeResultSubtypeDetail(raw.Subtype) != ""
		resultErrors := claudeResultErrorTexts(raw.Errors, raw.Result, abnormalResult)
		if abnormalResult {
			attrs := []any{
				"subtype", raw.Subtype,
				"is_error", raw.IsError,
				"session_id", raw.SessionID,
				"duration_ms", raw.DurationMs,
				"duration_api_ms", raw.DurationAPIMs,
				"num_turns", raw.NumTurns,
				"total_cost_usd", raw.TotalCostUSD,
			}
			if len(raw.Errors) > 0 {
				attrs = append(attrs, "errors", strings.Join(raw.Errors, "; "))
			}
			slog.Warn("claude result reported an error", attrs...)
		}
		if raw.Usage != nil {
			// Resolve model: prefer modelUsage keys (which contain "provider/model"),
			// fall back to empty (recordUsage will use the process-level fallback).
			// modelUsage is long-lifecycle cumulative and may list several models when
			// a turn fans out to subagents; only attribute when exactly one model is
			// named, otherwise leave model empty rather than mis-attribute the whole
			// request to an arbitrary model.
			var model string
			if len(raw.ModelUsage) == 1 {
				for m := range raw.ModelUsage {
					model = m
				}
			}
			// Top-level usage is a single request (not cumulative); emit it directly.
			// Anthropic splits input into input_tokens (non-cached) plus
			// cache_read_input_tokens and cache_creation_input_tokens. Normalize to
			// the shared contract: input = total input including cached, cached =
			// cache_read subset, so total = input + output matches the provider total.
			events = append(events, drivercontract.ParsedEvent{
				Kind:  consts.ParsedEventKindUsage,
				Model: model,
				InputTokens: uint64(raw.Usage.InputTokens) +
					uint64(raw.Usage.CacheCreationInputTokens) +
					uint64(raw.Usage.CacheReadInputTokens),
				CachedInputTokens: uint64(raw.Usage.CacheReadInputTokens),
				OutputTokens:      uint64(raw.Usage.OutputTokens),
			})
		}
		if abnormalResult {
			transientError := claudeResultTransientRuntimeError(resultErrors)
			metadata := mergeTransientRuntimeErrorMetadata(nil, transientError)
			if compactionInProgress && !compactBoundaryObserved {
				if metadata == nil {
					metadata = map[string]any{}
				}
				metadata[contextCompactionFailedMetadataKey] = true
			}
			events = append(events, drivercontract.ParsedEvent{
				Kind:       consts.ParsedEventKindError,
				Message:    claudeResultErrorMessage(raw.Subtype, raw.Message.StopReason, resultErrors),
				SessionID:  raw.SessionID,
				StopReason: raw.Message.StopReason,
				Metadata:   metadata,
			})
		}
		var turnEndMetadata map[string]any
		if abnormalResult {
			turnEndMetadata = map[string]any{"claude_result_abnormal": true}
		}
		if compactionInProgress {
			if turnEndMetadata == nil {
				turnEndMetadata = map[string]any{}
			}
			if compactBoundaryObserved {
				turnEndMetadata[contextCompactionConfirmedMetadataKey] = true
			} else {
				turnEndMetadata[contextCompactionFailedMetadataKey] = true
			}
		}
		events = append(events, drivercontract.ParsedEvent{
			Kind:       consts.ParsedEventKindTurnEnd,
			SessionID:  raw.SessionID,
			StopReason: raw.Message.StopReason,
			Metadata:   turnEndMetadata,
		})
	}

	return events
}

// InterruptActiveRun asks Claude Code stream-json to abort the currently
// running query while keeping the CLI session alive. Claude Code exposes this
// through the SDK control_request protocol, the same path used by print-mode
// clients for an interrupt/Escape action.