package ai

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

	"opsxdb/server/internal/aitools"
)

const (
	defaultAgentMaxRounds   = 8
	maxToolResultFeedChars  = 12000
	maxStepArgsChars        = 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"`
}

// 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"`
}

// AgentRes Agent 最终回答 + 步骤。
type AgentRes struct {
	Content string      `json:"content"`
	SQL     string      `json:"sql,omitempty"`
	Steps   []AgentStep `json:"steps"`
	Rounds  int         `json:"rounds"`
}

// RunAgent 多轮：LLM tool_call → aitools.Runtime.Call → 回灌，直到最终文本或轮次耗尽。
// ConfirmDanger 永不自动 true（由 Runtime 保证）。Ask 模式不要走这里。
func RunAgent(ctx context.Context, llm ToolChat, rt *aitools.Runtime, req AgentReq) (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 == "" {
		return AgentRes{}, fmt.Errorf("prompt required")
	}
	maxRounds := req.MaxRounds
	if maxRounds <= 0 {
		maxRounds = defaultAgentMaxRounds
	}
	if maxRounds > 16 {
		maxRounds = 16
	}

	tools := aitools.OpenAITools()
	msgs := []Message{
		{Role: "system", Content: buildAgentSystem(req)},
		{Role: "user", Content: prompt},
	}

	var steps []AgentStep
	rounds := 0
	for round := 0; round < maxRounds; round++ {
		rounds = round + 1
		turn, err := llm.ChatWithTools(ctx, req.Model, msgs, tools)
		if err != nil {
			return AgentRes{Steps: steps, Rounds: rounds}, err
		}
		if len(turn.ToolCalls) == 0 {
			content := turn.Content
			if content == "" {
				content = "（模型未返回文本，也未调用工具）"
			}
			return AgentRes{
				Content: content,
				SQL:     extractSQL(content),
				Steps:   steps,
				Rounds:  rounds,
			}, nil
		}

		// assistant 消息带 tool_calls
		msgs = append(msgs, Message{
			Role:      "assistant",
			Content:   turn.Content,
			ToolCalls: turn.ToolCalls,
		})

		for _, tc := range turn.ToolCalls {
			step, toolMsg := execOneTool(ctx, rt, tc)
			steps = append(steps, step)
			msgs = append(msgs, toolMsg)
		}
	}

	return AgentRes{
		Content: fmt.Sprintf("已达最大工具轮次（%d）。请根据已执行步骤自行判断，或缩小问题后重试。", maxRounds),
		Steps:   steps,
		Rounds:  rounds,
	}, 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. 先确认 connection_id（可用 opsxdb_list_connections）；用户给出偏好连接时优先用它。\n")
	b.WriteString("4. 用中文简明回答；可执行 SQL 放在 ```sql 代码块中。\n")
	b.WriteString("5. 工具结果可能截断，注意说明。\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,
		}
	}

	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(),
		}
	}
	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 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]) + "…"
}