package agent

import (
	"encoding/json"
	"strings"

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

type kimiACPSessionUpdate struct {
	SessionID string          `json:"sessionId"`
	Update    json.RawMessage `json:"update"`
}

type kimiACPContentBlock struct {
	Type string `json:"type"`
	Text string `json:"text"`
}

type kimiACPToolCallItem struct {
	Type    string              `json:"type"`
	Content kimiACPContentBlock `json:"content"`
	Text    string              `json:"text"`
}

func (d *KimiACPDriver) handleACPNotification(method string, params json.RawMessage) {
	if method != "session/update" {
		return
	}
	var notification kimiACPSessionUpdate
	if json.Unmarshal(params, &notification) != nil {
		return
	}
	sessionID := strings.TrimSpace(notification.SessionID)
	if sessionID == "" {
		sessionID = d.currentSessionID()
	}
	var head struct {
		SessionUpdate string `json:"sessionUpdate"`
	}
	if json.Unmarshal(notification.Update, &head) != nil {
		return
	}
	switch head.SessionUpdate {
	case "agent_message_chunk":
		d.emitACPContentChunk(consts.ParsedEventKindText, notification.Update, sessionID)
	case "agent_thought_chunk":
		d.emitACPContentChunk(consts.ParsedEventKindThinking, notification.Update, sessionID)
	case "tool_call":
		d.emitACPToolCall(notification.Update, sessionID, false)
	case "tool_call_update":
		d.emitACPToolCall(notification.Update, sessionID, true)
	case "plan":
		d.emitACPPlan(notification.Update, sessionID)
	case "usage_update":
		d.emitACPUsage(notification.Update, sessionID)
	}
}

func (d *KimiACPDriver) emitACPContentChunk(kind string, raw json.RawMessage, sessionID string) {
	var decoded struct {
		Content kimiACPContentBlock `json:"content"`
	}
	if json.Unmarshal(raw, &decoded) != nil || decoded.Content.Type != "text" || decoded.Content.Text == "" {
		return
	}
	d.emit(translatedEvent{Kind: kind, Text: decoded.Content.Text, SessionID: sessionID})
}

func (d *KimiACPDriver) emitACPToolCall(raw json.RawMessage, sessionID string, update bool) {
	var decoded struct {
		ToolCallID string                `json:"toolCallId"`
		Title      string                `json:"title"`
		Kind       string                `json:"kind"`
		Status     string                `json:"status"`
		RawInput   json.RawMessage       `json:"rawInput"`
		Content    []kimiACPToolCallItem `json:"content"`
	}
	if json.Unmarshal(raw, &decoded) != nil {
		return
	}
	toolCallID := strings.TrimSpace(decoded.ToolCallID)
	if toolCallID == "" {
		return
	}
	d.mu.Lock()
	title := d.toolTitles[toolCallID]
	if title == "" {
		title = strings.TrimSpace(decoded.Title)
		if title != "" {
			d.toolTitles[toolCallID] = title
		}
	}
	d.mu.Unlock()
	toolName := firstNonEmptyString(title, strings.TrimSpace(decoded.Kind), "tool")

	status := strings.ToLower(strings.TrimSpace(decoded.Status))
	if status == "completed" || status == "failed" {
		d.emit(translatedEvent{
			Kind:       consts.ParsedEventKindToolOutput,
			ToolName:   toolName,
			ToolCallID: toolCallID,
			Text:       kimiACPToolContentText(decoded.Content),
			SessionID:  sessionID,
			Metadata:   map[string]any{"is_error": status == "failed"},
		})
		return
	}
	if update {
		return
	}
	var toolInput any
	if len(decoded.RawInput) > 0 {
		_ = json.Unmarshal(decoded.RawInput, &toolInput)
	}
	d.emit(translatedEvent{
		Kind:       consts.ParsedEventKindToolCall,
		ToolName:   toolName,
		ToolCallID: toolCallID,
		ToolInput:  toolInput,
		SessionID:  sessionID,
	})
}

func kimiACPToolContentText(items []kimiACPToolCallItem) string {
	var builder strings.Builder
	for _, item := range items {
		if item.Content.Type == "text" && item.Content.Text != "" {
			builder.WriteString(item.Content.Text)
			continue
		}
		if item.Text != "" {
			builder.WriteString(item.Text)
		}
	}
	return builder.String()
}

func (d *KimiACPDriver) emitACPPlan(raw json.RawMessage, sessionID string) {
	var decoded struct {
		Entries []struct {
			Content string `json:"content"`
			Status  string `json:"status"`
		} `json:"entries"`
	}
	if json.Unmarshal(raw, &decoded) != nil {
		return
	}
	lines := make([]string, 0, len(decoded.Entries))
	for _, entry := range decoded.Entries {
		content := strings.TrimSpace(entry.Content)
		if content == "" {
			continue
		}
		if status := strings.TrimSpace(entry.Status); status != "" {
			content = "[" + status + "] " + content
		}
		lines = append(lines, content)
	}
	if len(lines) > 0 {
		d.emit(translatedEvent{
			Kind:      consts.ParsedEventKindThinking,
			Text:      strings.Join(lines, "\n"),
			SessionID: sessionID,
			Metadata:  map[string]any{"kimi_event": "plan"},
		})
	}
}

func (d *KimiACPDriver) emitACPUsage(raw json.RawMessage, sessionID string) {
	var decoded struct {
		Used uint64 `json:"used"`
		Size uint64 `json:"size"`
	}
	if json.Unmarshal(raw, &decoded) != nil || decoded.Size == 0 {
		return
	}
	d.emit(translatedEvent{
		Kind:                consts.ParsedEventKindUsage,
		SessionID:           sessionID,
		ActiveContextTokens: decoded.Used,
		ContextWindowTokens: decoded.Size,
		Metadata:            map[string]any{"kimi_event": "usage_update"},
	})
}

func (d *KimiACPDriver) handleServerRequest(envelope kimiACPEnvelope) {
	switch envelope.Method {
	case "session/request_permission":
		// Full-access is set through Kimi's yolo mode. Requests that still
		// reach the client are questions or unknown permission shapes; AgentTeam has
		// no interactive ACP UI and must never choose an answer for the user.
		d.writeRPCResponse(map[string]any{
			"jsonrpc": "2.0",
			"id":      envelope.ID,
			"result": map[string]any{
				"outcome": map[string]any{"outcome": "cancelled"},
			},
		})
	default:
		d.writeRPCError(envelope.ID, -32601, "unsupported ACP request "+envelope.Method)
	}
}

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

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