package agent

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

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

func (d *KimiWireDriver) handleServerRequest(envelope kimiWireEnvelope) {
	var request struct {
		Type    string          `json:"type"`
		Payload json.RawMessage `json:"payload"`
	}
	if err := json.Unmarshal(envelope.Params, &request); err != nil {
		d.writeRPCError(envelope.ID, -32602, "invalid Kimi Wire request")
		return
	}
	var payload struct {
		ID string `json:"id"`
	}
	_ = json.Unmarshal(request.Payload, &payload)
	switch request.Type {
	case "ApprovalRequest":
		d.writeRPCResult(envelope.ID, map[string]any{
			"request_id": payload.ID,
			"response":   "reject",
			"feedback":   "AgentTeam does not grant interactive Kimi Wire approvals",
		})
	case "QuestionRequest":
		d.writeRPCResult(envelope.ID, map[string]any{
			"request_id": payload.ID,
			"answers":    map[string]string{},
		})
	default:
		d.writeRPCError(envelope.ID, -32601, fmt.Sprintf("unsupported Kimi Wire request type %q", request.Type))
	}
}

func (d *KimiWireDriver) writeRPCResult(id json.RawMessage, result any) {
	d.writeRPCResponse(map[string]any{"jsonrpc": "2.0", "id": json.RawMessage(id), "result": result})
}

func (d *KimiWireDriver) writeRPCError(id json.RawMessage, code int, message string) {
	d.writeRPCResponse(map[string]any{
		"jsonrpc": "2.0",
		"id":      json.RawMessage(id),
		"error":   map[string]any{"code": code, "message": message},
	})
}

func (d *KimiWireDriver) writeRPCResponse(payload any) {
	data, err := json.Marshal(payload)
	if err != nil {
		return
	}
	_ = d.writeFrame(append(data, '\n'))
}

func (d *KimiWireDriver) handleNotification(method string, params json.RawMessage) {
	if method != "event" {
		return
	}
	var event struct {
		Type    string          `json:"type"`
		Payload json.RawMessage `json:"payload"`
	}
	if json.Unmarshal(params, &event) != nil {
		return
	}
	sessionID := d.currentSessionID()
	switch event.Type {
	case "TurnBegin":
		inputID := kimiRuntimeInputIDFromPayload(event.Payload)
		d.mu.Lock()
		d.turnActive = true
		d.wireTurnEnded = false
		d.turnEndEmitted = false
		accepted := d.promptAccepted
		expectedInputID := d.promptInputID
		if accepted != nil && expectedInputID != "" && inputID == expectedInputID {
			d.promptAccepted = nil
		} else {
			accepted = nil
		}
		d.mu.Unlock()
		if accepted != nil {
			accepted <- nil
		}
		d.emit(translatedEvent{
			Kind:      consts.ParsedEventKindThinking,
			SessionID: sessionID,
			Metadata:  map[string]any{codexAuthoritativeTurnStartMetadataKey: true},
		})
		if inputID != "" {
			d.emit(translatedEvent{Kind: consts.ParsedEventKindRuntimeInputAcked, RuntimeInputID: inputID, SessionID: sessionID})
		}
	case "SteerInput":
		inputID := kimiRuntimeInputIDFromPayload(event.Payload)
		if inputID != "" {
			d.mu.Lock()
			applied := d.steerWaiters[inputID]
			if applied != nil {
				delete(d.steerWaiters, inputID)
			}
			d.mu.Unlock()
			if applied != nil {
				applied <- nil
			}
			d.emit(translatedEvent{Kind: consts.ParsedEventKindRuntimeInputAcked, RuntimeInputID: inputID, SessionID: sessionID})
		}
	case "TurnEnd":
		d.mu.Lock()
		d.completedTurns++
		d.wireTurnEnded = true
		d.markTurnTerminalLocked()
		emitError, emitTurnEnd, terminalSessionID := d.takePromptTerminalEventsLocked()
		d.mu.Unlock()
		if emitError != "" {
			d.emit(translatedEvent{
				Kind:      consts.ParsedEventKindError,
				Message:   emitError,
				SessionID: terminalSessionID,
				Metadata:  map[string]any{"will_retry": false},
			})
		}
		if emitTurnEnd != "" {
			d.emit(translatedEvent{Kind: consts.ParsedEventKindTurnEnd, SessionID: terminalSessionID, StopReason: emitTurnEnd})
		}
	case "StepInterrupted":
		d.emit(translatedEvent{Kind: consts.ParsedEventKindThinking, Text: "Kimi step interrupted", SessionID: sessionID})
	case "StepRetry":
		d.emit(translatedEvent{Kind: consts.ParsedEventKindThinking, Text: kimiRetryText(event.Payload), SessionID: sessionID})
	case "StepBegin":
		var step struct {
			Number int `json:"n"`
		}
		if json.Unmarshal(event.Payload, &step) == nil {
			d.emit(translatedEvent{
				Kind:      consts.ParsedEventKindThinking,
				Text:      fmt.Sprintf("Kimi step %d", step.Number),
				SessionID: sessionID,
				Metadata:  map[string]any{"kimi_event": "step_begin", "step": step.Number},
			})
		}
	case "CompactionBegin":
		d.emit(translatedEvent{
			Kind:      consts.ParsedEventKindRuntimeCompacting,
			SessionID: sessionID,
			Metadata:  map[string]any{"kimi_event": "compaction_begin"},
		})
	case "CompactionEnd":
		d.mu.Lock()
		d.contextCompacted = true
		d.mu.Unlock()
		d.emit(translatedEvent{
			Kind:      consts.ParsedEventKindThinking,
			Text:      "Kimi context compaction completed",
			SessionID: sessionID,
			Metadata: map[string]any{
				"kimi_event":                          "compaction_end",
				contextCompactionConfirmedMetadataKey: true,
			},
		})
	case "MCPLoadingBegin", "MCPLoadingEnd":
		d.emit(translatedEvent{
			Kind:      consts.ParsedEventKindThinking,
			Text:      "Kimi MCP loading " + strings.TrimPrefix(strings.ToLower(event.Type), "mcploading"),
			SessionID: sessionID,
			Metadata:  map[string]any{"kimi_event": strings.ToLower(event.Type)},
		})
	case "Notification":
		var notification struct {
			Title    string `json:"title"`
			Body     string `json:"body"`
			Severity string `json:"severity"`
		}
		if json.Unmarshal(event.Payload, &notification) == nil {
			text := strings.TrimSpace(notification.Title)
			if body := strings.TrimSpace(notification.Body); body != "" {
				if text != "" {
					text += ": "
				}
				text += body
			}
			if text != "" {
				d.emit(translatedEvent{
					Kind:      consts.ParsedEventKindThinking,
					Text:      text,
					SessionID: sessionID,
					Metadata:  map[string]any{"kimi_event": "notification", "severity": notification.Severity},
				})
			}
		}
	case "SubagentEvent":
		d.handleKimiSubagentEvent(event.Payload, sessionID)
	case "ContentPart":
		d.handleKimiContentPart(event.Payload, sessionID)
	case "ToolCall":
		d.handleKimiToolCall(event.Payload, sessionID)
	case "ToolResult":
		d.handleKimiToolResult(event.Payload, sessionID)
	case "PlanDisplay":
		var plan struct {
			Content  string `json:"content"`
			FilePath string `json:"file_path"`
		}
		if json.Unmarshal(event.Payload, &plan) == nil && strings.TrimSpace(plan.Content) != "" {
			d.emit(translatedEvent{
				Kind:      consts.ParsedEventKindThinking,
				Text:      plan.Content,
				SessionID: sessionID,
				Metadata:  map[string]any{"kimi_event": "plan", "file_path": plan.FilePath},
			})
		}
	case "StatusUpdate":
		d.handleKimiStatusUpdate(event.Payload, sessionID)
	}
}

func kimiRuntimeInputIDFromPayload(payload json.RawMessage) string {
	var decoded struct {
		UserInput json.RawMessage `json:"user_input"`
	}
	if json.Unmarshal(payload, &decoded) != nil {
		return ""
	}
	var text string
	if json.Unmarshal(decoded.UserInput, &text) == nil {
		return kimiRuntimeInputIDFromText(text)
	}
	var parts []struct {
		Type string `json:"type"`
		Text string `json:"text"`
	}
	if json.Unmarshal(decoded.UserInput, &parts) != nil {
		return ""
	}
	inputID := ""
	for _, part := range parts {
		if part.Type != "text" {
			continue
		}
		partInputID := kimiRuntimeInputIDFromText(part.Text)
		if partInputID == "" {
			continue
		}
		if inputID != "" {
			return ""
		}
		inputID = partInputID
	}
	return inputID
}

func kimiRuntimeInputIDFromText(text string) string {
	match := kimiRuntimeInputMarkerPattern.FindStringSubmatch(strings.TrimSpace(text))
	if len(match) != 2 {
		return ""
	}
	return strings.TrimSpace(match[1])
}

func kimiRetryText(payload json.RawMessage) string {
	var retry struct {
		NextAttempt int     `json:"next_attempt"`
		MaxAttempts int     `json:"max_attempts"`
		WaitSeconds float64 `json:"wait_s"`
		ErrorType   string  `json:"error_type"`
	}
	if json.Unmarshal(payload, &retry) != nil {
		return "Kimi is retrying the current step"
	}
	return fmt.Sprintf("Kimi retry %d/%d after %.1fs (%s)", retry.NextAttempt, retry.MaxAttempts, retry.WaitSeconds, retry.ErrorType)
}

func (d *KimiWireDriver) handleKimiContentPart(payload json.RawMessage, sessionID string) {
	var content struct {
		Type  string `json:"type"`
		Text  string `json:"text"`
		Think string `json:"think"`
	}
	if json.Unmarshal(payload, &content) != nil {
		return
	}
	switch content.Type {
	case "text":
		if content.Text != "" {
			d.emit(translatedEvent{Kind: consts.ParsedEventKindText, Text: content.Text, SessionID: sessionID})
		}
	case "think":
		if content.Think != "" {
			d.emit(translatedEvent{Kind: consts.ParsedEventKindThinking, Text: content.Think, SessionID: sessionID})
		}
	}
}

func (d *KimiWireDriver) handleKimiToolCall(payload json.RawMessage, sessionID string) {
	var call struct {
		ID       string `json:"id"`
		Function struct {
			Name      string `json:"name"`
			Arguments string `json:"arguments"`
		} `json:"function"`
	}
	if json.Unmarshal(payload, &call) != nil || strings.TrimSpace(call.ID) == "" {
		return
	}
	d.mu.Lock()
	d.toolNames[call.ID] = call.Function.Name
	d.mu.Unlock()
	var input any = call.Function.Arguments
	if strings.TrimSpace(call.Function.Arguments) == "" {
		input = "{}"
	}
	d.emit(translatedEvent{
		Kind:       consts.ParsedEventKindToolCall,
		ToolName:   call.Function.Name,
		ToolCallID: call.ID,
		ToolInput:  input,
		SessionID:  sessionID,
	})
}

func (d *KimiWireDriver) handleKimiToolResult(payload json.RawMessage, sessionID string) {
	var result struct {
		ToolCallID string `json:"tool_call_id"`
		Return     struct {
			IsError bool            `json:"is_error"`
			Output  json.RawMessage `json:"output"`
			Message string          `json:"message"`
		} `json:"return_value"`
	}
	if json.Unmarshal(payload, &result) != nil {
		return
	}
	d.mu.Lock()
	toolName := d.toolNames[result.ToolCallID]
	delete(d.toolNames, result.ToolCallID)
	d.mu.Unlock()
	text := kimiWireOutputText(result.Return.Output)
	if text == "" {
		text = result.Return.Message
	}
	d.emit(translatedEvent{
		Kind:       consts.ParsedEventKindToolOutput,
		Text:       text,
		ToolName:   toolName,
		ToolCallID: result.ToolCallID,
		SessionID:  sessionID,
		Metadata:   map[string]any{"is_error": result.Return.IsError},
	})
}

func kimiWireOutputText(raw json.RawMessage) string {
	var text string
	if json.Unmarshal(raw, &text) == nil {
		return text
	}
	var parts []struct {
		Type string `json:"type"`
		Text string `json:"text"`
	}
	if json.Unmarshal(raw, &parts) == nil {
		var output strings.Builder
		for _, part := range parts {
			if part.Type == "text" {
				output.WriteString(part.Text)
			}
		}
		return output.String()
	}
	return strings.TrimSpace(string(raw))
}

func (d *KimiWireDriver) handleKimiStatusUpdate(payload json.RawMessage, sessionID string) {
	d.handleKimiStatusUpdateForContext(payload, sessionID, true)
}

func (d *KimiWireDriver) handleKimiStatusUpdateForContext(
	payload json.RawMessage,
	sessionID string,
	contextExposureKnown bool,
) {
	var status struct {
		MessageID  string          `json:"message_id"`
		MCPStatus  any             `json:"mcp_status"`
		TokenUsage *kimiTokenUsage `json:"token_usage"`
	}
	if json.Unmarshal(payload, &status) != nil {
		return
	}
	if status.MCPStatus != nil {
		d.emit(translatedEvent{
			Kind:      consts.ParsedEventKindThinking,
			Text:      "Kimi MCP status updated",
			SessionID: sessionID,
			Metadata:  map[string]any{"kimi_event": "mcp_status", "mcp_status": status.MCPStatus},
		})
	}
	if status.TokenUsage == nil {
		return
	}
	// kimi-cli emits one StatusUpdate per step with a stable message_id
	// (result.id); dedup by session+message_id so Wire/SSE replay does not
	// double-count token usage.
	messageID := strings.TrimSpace(status.MessageID)
	if messageID != "" {
		key := sessionID + "\x00" + messageID
		d.mu.Lock()
		if d.countedStatusUpdates == nil {
			d.countedStatusUpdates = make(map[string]struct{})
		}
		if _, seen := d.countedStatusUpdates[key]; seen {
			d.mu.Unlock()
			return
		}
		d.countedStatusUpdates[key] = struct{}{}
		contextCompacted := d.contextCompacted && contextExposureKnown
		if contextCompacted {
			d.contextCompacted = false
		}
		d.mu.Unlock()
		d.emitKimiUsage(status.TokenUsage, sessionID, messageID, contextExposureKnown, contextCompacted)
		return
	}
	d.mu.Lock()
	contextCompacted := d.contextCompacted && contextExposureKnown
	if contextCompacted {
		d.contextCompacted = false
	}
	d.mu.Unlock()
	d.emitKimiUsage(status.TokenUsage, sessionID, messageID, contextExposureKnown, contextCompacted)
}

type kimiTokenUsage struct {
	InputOther         uint64 `json:"input_other"`
	InputCacheRead     uint64 `json:"input_cache_read"`
	InputCacheCreation uint64 `json:"input_cache_creation"`
	Output             uint64 `json:"output"`
}

func (d *KimiWireDriver) emitKimiUsage(
	usage *kimiTokenUsage,
	sessionID, messageID string,
	contextExposureKnown, contextCompacted bool,
) {
	d.emit(translatedEvent{
		Kind:                    consts.ParsedEventKindUsage,
		SessionID:               sessionID,
		ModelInvocationID:       messageID,
		ContextExposureKnown:    contextExposureKnown,
		ContextExposureExcluded: !contextExposureKnown,
		ContextCompacted:        contextCompacted,
		InputTokens:             usage.InputOther + usage.InputCacheRead + usage.InputCacheCreation,
		CachedInputTokens:       usage.InputCacheRead,
		OutputTokens:            usage.Output,
	})
}

func (d *KimiWireDriver) handleKimiSubagentEvent(payload json.RawMessage, sessionID string) {
	var subagent struct {
		ParentToolCallID string `json:"parent_tool_call_id"`
		AgentID          string `json:"agent_id"`
		SubagentType     string `json:"subagent_type"`
		Event            struct {
			Type    string          `json:"type"`
			Payload json.RawMessage `json:"payload"`
		} `json:"event"`
	}
	if json.Unmarshal(payload, &subagent) != nil {
		return
	}
	// A subagent's per-step StatusUpdate is embedded in the SubagentEvent; feed
	// it to the usage handler so subagent tokens are counted (deduped by
	// session+message_id) instead of dropped.
	if strings.TrimSpace(subagent.Event.Type) == "StatusUpdate" {
		d.handleKimiStatusUpdateForContext(subagent.Event.Payload, sessionID, false)
	}
	d.emit(translatedEvent{
		Kind:      consts.ParsedEventKindSubagentActivity,
		Message:   strings.TrimSpace("Kimi subagent " + subagent.SubagentType + " " + subagent.Event.Type),
		SessionID: sessionID,
		Metadata: map[string]any{
			"kimi_event":          "subagent",
			"agent_id":            subagent.AgentID,
			"subagent_type":       subagent.SubagentType,
			"parent_tool_call_id": subagent.ParentToolCallID,
			"event_type":          subagent.Event.Type,
		},
	})
}

func (d *KimiWireDriver) currentSessionID() string {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.sessionID
}

func (d *KimiWireDriver) emit(event translatedEvent) {
	d.mu.Lock()
	output := d.output
	d.mu.Unlock()
	if output == nil {
		return
	}
	data, err := json.Marshal(event)
	if err != nil {
		return
	}
	output.emit(string(data))
}

func (d *KimiWireDriver) ParseLine(line string) []drivercontract.ParsedEvent {
	var event translatedEvent
	if json.Unmarshal([]byte(strings.TrimSpace(line)), &event) != nil {
		return nil
	}
	parsed := drivercontract.ParsedEvent{
		Kind:                    event.Kind,
		Text:                    event.Text,
		SectionID:               event.SectionID,
		ToolName:                event.ToolName,
		ToolCallID:              event.ToolCallID,
		ToolInput:               event.ToolInput,
		RuntimeInputID:          event.RuntimeInputID,
		SessionID:               event.SessionID,
		Message:                 event.Message,
		Metadata:                event.Metadata,
		StopReason:              event.StopReason,
		InputTokens:             event.InputTokens,
		CachedInputTokens:       event.CachedInputTokens,
		OutputTokens:            event.OutputTokens,
		Model:                   event.Model,
		ModelInvocationID:       event.ModelInvocationID,
		ContextExposureKnown:    event.ContextExposureKnown,
		ContextExposureExcluded: event.ContextExposureExcluded,
		ContextCompacted:        event.ContextCompacted,
	}
	switch event.Kind {
	case consts.ParsedEventKindSessionInit:
		if parsed.SessionID == "" {
			return nil
		}
	case consts.ParsedEventKindText, consts.ParsedEventKindThinking:
		if parsed.Text == "" && !parsedBoolMetadata(parsed.Metadata, codexAuthoritativeTurnStartMetadataKey) {
			return nil
		}
	case consts.ParsedEventKindRuntimeInputAcked:
		if parsed.RuntimeInputID == "" {
			return nil
		}
	case consts.ParsedEventKindToolCall, consts.ParsedEventKindToolOutput,
		consts.ParsedEventKindTurnEnd, consts.ParsedEventKindError,
		consts.ParsedEventKindUsage, consts.ParsedEventKindRuntimeCompacting,
		consts.ParsedEventKindSubagentActivity:
	default:
		return nil
	}
	return []drivercontract.ParsedEvent{parsed}
}