package ai

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

	"opsxdb/server/internal/aitools"
)

const (
	defaultAgentMaxRounds  = 8
	maxToolResultFeedChars = 12000
	maxStepArgsChars       = 200
	maxGridRows            = 200
)

// AgentReq 侧栏 Agent 请求（与 Ask 的 /api/ai/chat 分路径）。
type AgentReq struct {
	Prompt       string    `json:"prompt"`
	ConnectionID string    `json:"connectionId,omitempty"`
	Driver       string    `json:"driver,omitempty"`
	Database     string    `json:"database,omitempty"`
	Model        string    `json:"model,omitempty"`
	MaxRounds    int       `json:"maxRounds,omitempty"`
	Stream       bool      `json:"stream,omitempty"`
	SessionID    string    `json:"sessionId,omitempty"`
	Continue     bool      `json:"continue,omitempty"` // 真·续跑：注入 Prior
	Prior        []Message `json:"prior,omitempty"`    // 含 tool 轨迹的历史（不含本轮 user）
	RunID        string    `json:"runId,omitempty"`    // 可选；空则自动生成
}

// AgentStep 挂在回复下的步骤时间线一项。
type AgentStep struct {
	Tool         string            `json:"tool"`
	Args         string            `json:"args,omitempty"`
	OK           bool              `json:"ok"`
	Summary      string            `json:"summary"`
	Error        string            `json:"error,omitempty"`
	Ms           int64             `json:"ms"`
	NeedConfirm  bool              `json:"needConfirm,omitempty"`
	ConfirmSQL   string            `json:"confirmSql,omitempty"`
	ConfirmBind  string            `json:"confirmBind,omitempty"`
	Pending      bool              `json:"pending,omitempty"`
	Result       *AgentQueryResult `json:"result,omitempty"`
}

// AgentRes Agent 最终回答 + 步骤 + 运行结果契约。
type AgentRes struct {
	Content          string      `json:"content"`
	SQL              string      `json:"sql,omitempty"`
	Steps            []AgentStep `json:"steps"`
	Rounds           int         `json:"rounds"`
	RunID            string      `json:"runId,omitempty"`
	SessionID        string      `json:"sessionId,omitempty"`
	Outcome          string      `json:"runOutcome,omitempty"`
	IncompleteReason string      `json:"incompleteReason,omitempty"`
	Resumable        bool        `json:"resumable,omitempty"`
	StepsBudget      int         `json:"stepsBudget,omitempty"`
	// Trajectory 含 system/user/assistant.tool_calls/tool 的完整对话（供落盘续跑；JSON 默认省略大字段可选）
	Trajectory       []Message   `json:"trajectory,omitempty"`
}

// RunAgent 多轮：LLM tool_call → aitools.Runtime.Call → 回灌。
// ConfirmDanger 永不由模型自动带上（execOneTool 会剥离 confirm_danger）。
func RunAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentReq) (AgentRes, error) {
	return runAgent(ctx, llm, rt, req, nil)
}

// RunAgentStream 同 RunAgent，经 emit 边跑边推（步骤 / 正文块 / done）。
func RunAgentStream(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentReq, emit AgentEmitter) (AgentRes, error) {
	return runAgent(ctx, llm, rt, req, emit)
}

func runAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentReq, emit AgentEmitter) (AgentRes, error) {
	if rt == nil {
		return AgentRes{}, fmt.Errorf("aitools runtime required")
	}
	if llm == nil {
		return AgentRes{}, fmt.Errorf("tool chat LLM required")
	}
	prompt := strings.TrimSpace(req.Prompt)
	if prompt == "" && !(req.Continue && len(req.Prior) > 0) {
		return AgentRes{}, fmt.Errorf("prompt required")
	}
	h := DefaultHarness()
	maxRounds := h.ClampRounds(req.MaxRounds)

	runID := strings.TrimSpace(req.RunID)
	if runID == "" {
		runID = NewRunID()
	}
	sessionID := strings.TrimSpace(req.SessionID)
	if sessionID == "" {
		sessionID = NewSessionID()
		req.SessionID = sessionID
	}
	rt.RunID = runID
	// Agent 路径禁止模型自行带确认令牌
	rt.RejectCallerConfirm = true

	tools := aitools.OpenAITools()
	var msgs []Message
	if req.Continue && len(req.Prior) > 0 {
		msgs = append([]Message{}, req.Prior...)
		hasSystem := false
		for _, m := range msgs {
			if m.Role == "system" {
				hasSystem = true
				break
			}
		}
		if !hasSystem {
			msgs = append([]Message{{Role: "system", Content: buildAgentSystem(req)}}, msgs...)
		}
		if prompt != "" {
			msgs = append(msgs, Message{Role: "user", Content: prompt})
		}
		msgs = TrimMessages(msgs, h)
	} else {
		msgs = []Message{
			{Role: "system", Content: buildAgentSystem(req)},
			{Role: "user", Content: prompt},
		}
	}

	finish := func(steps []AgentStep, rounds int, msgs []Message) AgentRes {
		return AgentRes{
			Steps:       steps,
			Rounds:      rounds,
			RunID:       runID,
			SessionID:   sessionID,
			StepsBudget: maxRounds,
			Trajectory:  append([]Message{}, msgs...),
		}
	}

	var steps []AgentStep
	var callHistory []string
	rounds := 0
	for round := 0; round < maxRounds; round++ {
		if err := ctx.Err(); err != nil {
			res := finish(steps, rounds, msgs)
			res.Content = "已取消"
			res.Outcome = OutcomeCancelled
			res.Resumable = len(steps) > 0
			emitEvent(emit, AgentEvent{Type: AgentEventError, Error: "cancelled", Steps: steps, Rounds: rounds, RunID: runID, SessionID: sessionID, Outcome: OutcomeCancelled})
			return res, err
		}
		rounds = round + 1
		if tip := CheckpointHint(rounds); tip != "" {
			msgs = append(msgs, Message{Role: "system", Content: tip})
		}
		turn, err := llm.ChatWithTools(ctx, req.Model, msgs, tools)
		if err != nil {
			res := finish(steps, rounds, msgs)
			res.Outcome = OutcomeFailed
			emitEvent(emit, AgentEvent{Type: AgentEventError, Error: err.Error(), Steps: steps, Rounds: rounds, RunID: runID, SessionID: sessionID, Outcome: OutcomeFailed})
			return res, err
		}
		if len(turn.ToolCalls) == 0 {
			content := turn.Content
			if content == "" {
				content = "（模型未返回文本，也未调用工具）"
			}
			for _, part := range chunkDelta(content, 36) {
				emitEvent(emit, AgentEvent{Type: AgentEventDelta, Delta: part})
			}
			res := finish(steps, rounds, msgs)
			res.Content = content
			res.SQL = extractSQL(content)
			res.Outcome = OutcomeCompleted
			emitEvent(emit, AgentEvent{
				Type: AgentEventDone, Content: res.Content, SQL: res.SQL, Steps: steps, Rounds: rounds,
				RunID: runID, SessionID: sessionID, Outcome: OutcomeCompleted, StepsBudget: maxRounds,
			})
			return res, nil
		}

		msgs = append(msgs, Message{
			Role:      "assistant",
			Content:   turn.Content,
			ToolCalls: turn.ToolCalls,
		})

		awaiting := false
		for _, tc := range turn.ToolCalls {
			name := strings.TrimSpace(tc.Function.Name)
			argsRaw := strings.TrimSpace(tc.Function.Arguments)
			fp := ToolCallFingerprint(name, argsRaw)
			if DetectRepeatedToolCall(callHistory, fp, h.SameCallLimit) {
				step := AgentStep{
					Tool:    name,
					Args:    truncateRunes(argsRaw, maxStepArgsChars),
					OK:      false,
					Summary: "repeated tool call rejected (same args)",
					Error:   "repeated tool call rejected (same args)",
				}
				steps = append(steps, step)
				cp := step
				emitEvent(emit, AgentEvent{Type: AgentEventStep, Step: &cp, RunID: runID})
				msgs = append(msgs, Message{
					Role: "tool", ToolCallID: tc.ID, Name: name,
					Content: step.Error,
				})
				callHistory = append(callHistory, fp)
				continue
			}
			callHistory = append(callHistory, fp)

			pending := AgentStep{Tool: name, Args: truncateRunes(argsRaw, maxStepArgsChars), Pending: true}
			emitEvent(emit, AgentEvent{Type: AgentEventStepStart, Step: &pending, RunID: runID})

			step, toolMsg := execOneTool(ctx, rt, tc)
			steps = append(steps, step)
			cp := step
			ev := AgentEvent{Type: AgentEventStep, Step: &cp, RunID: runID}
			if step.Result != nil {
				ev.Result = step.Result
			}
			emitEvent(emit, ev)
			msgs = append(msgs, toolMsg)
			if step.NeedConfirm {
				awaiting = true
			}
		}
		if awaiting {
			content := "需要确认写操作后才能继续。"
			res := finish(steps, rounds, msgs)
			res.Content = content
			res.Outcome = OutcomeAwaitingConfirm
			res.Resumable = true
			emitEvent(emit, AgentEvent{
				Type: AgentEventDone, Content: content, Steps: steps, Rounds: rounds,
				RunID: runID, SessionID: sessionID, Outcome: OutcomeAwaitingConfirm, Resumable: true, StepsBudget: maxRounds,
			})
			return res, nil
		}
		msgs = TrimMessages(msgs, h)
	}

	content := fmt.Sprintf("已达最大工具轮次（%d）。可加预算继续，或缩小问题后重试。", maxRounds)
	for _, part := range chunkDelta(content, 36) {
		emitEvent(emit, AgentEvent{Type: AgentEventDelta, Delta: part})
	}
	res := finish(steps, rounds, msgs)
	res.Content = content
	res.Outcome = OutcomeIncomplete
	res.IncompleteReason = ReasonMaxSteps
	res.Resumable = true
	emitEvent(emit, AgentEvent{
		Type: AgentEventIncomplete, Content: content, Steps: steps, Rounds: rounds,
		RunID: runID, SessionID: sessionID, Outcome: OutcomeIncomplete, IncompleteReason: ReasonMaxSteps,
		Resumable: true, StepsBudget: maxRounds,
	})
	return res, nil
}

func buildAgentSystem(req AgentReq) string {
	var b strings.Builder
	b.WriteString("你是 OpsxDB 数据库 Agent。你可以通过工具查询用户已白名单的连接。\n")
	b.WriteString("规则：\n")
	b.WriteString("1. 需要查库时必须调用工具；不要假装已执行 SQL。\n")
	b.WriteString("2. 默认只读：execute_query 以 SELECT/SHOW/EXPLAIN/DESCRIBE/WITH…SELECT 为主；Redis 以读命令为主。\n")
	b.WriteString("3. 写操作若被拒绝并提示需要确认，向用户说明须在界面确认后执行；不要声称已写入。\n")
	b.WriteString("4. 先确认 connection_id（可用 opsxdb_list_connections）；用户给出偏好连接时优先用它。\n")
	b.WriteString("5. 用中文简明回答；可执行 SQL 放在 ```sql 代码块中。\n")
	b.WriteString("6. 工具结果可能截断，注意说明。\n")
	if req.ConnectionID != "" {
		b.WriteString(fmt.Sprintf("\n用户当前偏好连接 id: %s", req.ConnectionID))
		if req.Driver != "" {
			b.WriteString(fmt.Sprintf("（%s）", req.Driver))
		}
		if req.Database != "" {
			b.WriteString(fmt.Sprintf("，数据库: %s", req.Database))
		}
		b.WriteString("。\n")
	}
	return b.String()
}

func execOneTool(ctx context.Context, rt *aitools.Runtime, tc ToolCall) (AgentStep, Message) {
	name := strings.TrimSpace(tc.Function.Name)
	argsRaw := strings.TrimSpace(tc.Function.Arguments)
	if argsRaw == "" {
		argsRaw = "{}"
	}
	step := AgentStep{
		Tool: name,
		Args: truncateRunes(argsRaw, maxStepArgsChars),
	}
	start := time.Now()

	var args map[string]any
	if err := json.Unmarshal([]byte(argsRaw), &args); err != nil {
		step.OK = false
		step.Error = "invalid tool arguments JSON: " + err.Error()
		step.Summary = step.Error
		step.Ms = time.Since(start).Milliseconds()
		return step, Message{
			Role:       "tool",
			ToolCallID: tc.ID,
			Name:       name,
			Content:    step.Error,
		}
	}
	// 永不让模型自行带确认令牌
	delete(args, "confirm_danger")
	delete(args, "confirmDanger")

	out, err := rt.Call(ctx, name, args)
	step.Ms = time.Since(start).Milliseconds()
	if err != nil {
		step.OK = false
		step.Error = err.Error()
		step.Summary = err.Error()
		return step, Message{
			Role:       "tool",
			ToolCallID: tc.ID,
			Name:       name,
			Content:    err.Error(),
		}
	}

	if m, ok := out.(map[string]any); ok {
		if nc, _ := m["needConfirm"].(bool); nc {
			step.OK = false
			step.NeedConfirm = true
			if sql, _ := m["sql"].(string); sql != "" {
				step.ConfirmSQL = sql
			} else if cmd, _ := m["command"].(string); cmd != "" {
				step.ConfirmSQL = cmd
			}
			if bind, _ := m["confirmBind"].(string); bind != "" {
				step.ConfirmBind = bind
			}
			msg, _ := m["message"].(string)
			if msg == "" {
				msg = "requires user confirmation (confirm_danger)"
			}
			step.Error = msg
			step.Summary = msg
			feed := aitools.ResultJSON(out)
			return step, Message{
				Role:       "tool",
				ToolCallID: tc.ID,
				Name:       name,
				Content:    feed,
			}
		}
		if name == aitools.ToolExecuteQuery || name == aitools.ToolExecuteRedisCommand {
			if qr := extractQueryResult(name, args, m); qr != nil {
				step.Result = qr
			}
		}
	}

	step.OK = true
	step.Summary = summarizeToolResult(name, out)
	feed := aitools.ResultJSON(out)
	if len(feed) > maxToolResultFeedChars {
		feed = feed[:maxToolResultFeedChars] + "\n…(truncated)"
	}
	return step, Message{
		Role:       "tool",
		ToolCallID: tc.ID,
		Name:       name,
		Content:    feed,
	}
}

func extractQueryResult(tool string, args map[string]any, m map[string]any) *AgentQueryResult {
	cols, _ := m["columns"].([]string)
	if cols == nil {
		if raw, ok := m["columns"].([]any); ok {
			cols = make([]string, 0, len(raw))
			for _, c := range raw {
				cols = append(cols, fmt.Sprint(c))
			}
		}
	}
	var rows []map[string]interface{}
	switch v := m["rows"].(type) {
	case []map[string]interface{}:
		rows = v
	case []any:
		for _, item := range v {
			if row, ok := item.(map[string]interface{}); ok {
				rows = append(rows, row)
			} else if row, ok := item.(map[string]any); ok {
				rows = append(rows, row)
			}
		}
	}
	if len(cols) == 0 && len(rows) == 0 {
		return nil
	}
	if len(rows) > maxGridRows {
		rows = rows[:maxGridRows]
	}
	rc := len(rows)
	switch v := m["rowCount"].(type) {
	case int:
		rc = v
	case int64:
		rc = int(v)
	case float64:
		rc = int(v)
	}
	sqlText, _ := args["sql"].(string)
	if sqlText == "" {
		sqlText, _ = args["command"].(string)
	}
	return &AgentQueryResult{
		Columns:  cols,
		Rows:     rows,
		RowCount: rc,
		SQL:      sqlText,
		Tool:     tool,
	}
}

func summarizeToolResult(tool string, out any) string {
	m, ok := out.(map[string]any)
	if !ok {
		return truncateRunes(fmt.Sprint(out), 120)
	}
	lenAny := func(key string) (int, bool) {
		switch v := m[key].(type) {
		case []any:
			return len(v), true
		case []string:
			return len(v), true
		default:
			return 0, false
		}
	}
	switch tool {
	case aitools.ToolListConnections:
		if n, ok := lenAny("connections"); ok {
			return fmt.Sprintf("%d connections", n)
		}
	case aitools.ToolListDatabases:
		if n, ok := lenAny("databases"); ok {
			return fmt.Sprintf("%d databases", n)
		}
	case aitools.ToolListTables:
		if n, ok := lenAny("tables"); ok {
			return fmt.Sprintf("%d tables", n)
		}
	case aitools.ToolDescribeTable:
		if t, _ := m["table"].(string); t != "" {
			if n, ok := lenAny("columns"); ok {
				return fmt.Sprintf("%s (%d cols)", t, n)
			}
			return t
		}
	case aitools.ToolGetSchemaContext:
		if s, ok := m["schema"].(string); ok {
			return fmt.Sprintf("schema %d chars", len(s))
		}
	case aitools.ToolExecuteQuery, aitools.ToolExecuteRedisCommand:
		rc := 0
		switch v := m["rowCount"].(type) {
		case int:
			rc = v
		case int64:
			rc = int(v)
		case float64:
			rc = int(v)
		}
		return fmt.Sprintf("rows=%d", rc)
	}
	return truncateRunes(aitools.ResultJSON(out), 120)
}

func truncateRunes(s string, n int) string {
	s = strings.Join(strings.Fields(s), " ")
	r := []rune(s)
	if len(r) <= n {
		return s
	}
	return string(r[:n]) + "…"
}