package ai

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"regexp"
	"strings"
	"time"

	"eai_agentplatform/backend/internal/config"
)

// debugWriter 默认丢弃；当 AI_AGENT_DEBUG=1 时输出到 stdout，便于临时诊断多轮工具调用。
var debugWriter io.Writer = io.Discard

func init() {
	if v, ok := os.LookupEnv("AI_AGENT_DEBUG"); ok && v != "" && v != "0" {
		debugWriter = os.Stdout
	}
}

// Tool 一个可被 LLM 调用的工具。
type Tool interface {
	// Schema 返回 OpenAI 兼容的 function schema，用于随请求提交给 LLM。
	Schema() ToolSchema
	// Execute 执行工具，args 为模型传入的原始 JSON 参数，返回给模型的纯文本结果。
	Execute(ctx context.Context, args json.RawMessage) (string, error)
}

// Agent 带工具执行能力的对话代理。
//
// Agent 采用 OpenAI 兼容的 tool_calls 环形调用：
//  1. 携带 tools 调用 LLM
//  2. 若模型返回 tool_calls，逐一执行并把结果以 role=tool 消息回填
//  3. 再次调用 LLM，直到模型返回纯文本（无新 tool_calls）
//
// 相比裸 Client，Agent 让 LLM 能按需自主调用搜索等工具。
type Agent struct {
	client      *Client
	tools       map[string]Tool
	toolSchemas []ToolSchema
	maxRounds   int
}

// NewAgent 基于指定 AI 路由创建带工具能力的 Agent。
func NewAgent(route *config.RouteConfig) *Agent {
	return &Agent{
		client:    NewClient(route),
		tools:     make(map[string]Tool),
		maxRounds: 8,
	}
}

// RegisterTool 注册一个可供 LLM 调用的工具。
func (a *Agent) RegisterTool(t Tool) {
	a.tools[t.Schema().Function.Name] = t
	a.toolSchemas = append(a.toolSchemas, t.Schema())
}

// SetMaxRounds 调整最大工具调用轮数，避免死循环。
func (a *Agent) SetMaxRounds(n int) {
	if n > 0 {
		a.maxRounds = n
	}
}

// Run 执行一轮带工具的多轮对话，返回最终文本。
//
// messages 会被克隆后追加历史，不作为参数直接变更，避免污染调用方。
func (a *Agent) Run(ctx context.Context, messages []Message) (string, error) {
	content, _, _, err := a.RunWithUsage(ctx, messages)
	return content, err
}

// RunWithUsage 与 Run 相同，额外累计多轮 token 消耗，便于审计计费。
//
// 返回最终文本、累计 usage 与最终模型信息。usage 为多轮 tool_calls 的各次
// PromptTokens / CompletionTokens 累加（TotalTokens 仅为参考，不累加）。
func (a *Agent) RunWithUsage(ctx context.Context, messages []Message) (string, ChatResult, *config.RouteConfig, error) {
	// 复制消息历史，防止调用方切片被并发修改
	var history []Message
	history = append(history, messages...)

	var usage ChatResult
	var toolRounds int // 已执行的工具调用轮数，用于打断 tool_calls 轮转过量/死循环
	for round := 0; round < a.maxRounds; round++ {
		// 轮次间检查上层上下文（如编排 150s 兜底）是否已取消：
		// 若已取消，立即退出，避免已发起的慢调用/死循环续命拖死 wg.Wait()。
		if err := ctx.Err(); err != nil {
			return "", usage, a.client.Route(), fmt.Errorf("上下文已取消，终止工具调用: %w", err)
		}
		out, err := a.roundWithEmptyRecovery(ctx, history)
		if err != nil {
			return "", usage, a.client.Route(), err
		}
		// 累计 token 消耗（多轮加总）
		usage.Usage.PromptTokens += out.Usage.PromptTokens
		usage.Usage.CompletionTokens += out.Usage.CompletionTokens
		if out.Model != "" {
			usage.Model = out.Model
		}
		// 记录模型回复
		history = append(history, Message{Role: "assistant", Content: out.Content, ToolCalls: out.ToolCalls})

		// 无工具调用 → 返回最终文本
		if len(out.ToolCalls) == 0 {
			// 部分 provider（如 LMUAI/DeepSeek 兼容 Anthropic）会把工具调用意图以
			//「文本格式」输出（<｜｜ calls> 这类 Anthropic tool_use 转义标签），而不是
			// OpenAI 规范的 tool_calls JSON。此时 ToolCalls 为空、Content 便是那串原始标签，
			// 若直接返回会把 <｜｜ invoke name="web_search"> 之类的噪声原文展示给用户。
			// 识别到这类「工具回显」后，不保留原样：先尝试无工具纯文本收敛；若模型惯性
			// 仍输出标签（历史里已存在 assistant 工具标签消息，易诱导），则剥离标签，
			// 只把其中可读的自然语言/兜底文案交给用户。
			if out.Content != "" && looksLikeToolEcho(out.Content) && len(a.toolSchemas) > 0 {
				// 先剥离工具标签提取正文：若模型在本轮已写出正文（只是混入了工具回显标签），
				// 直接返回正文，避免接下来的无谓收敛覆盖掉真实产出（worker 尤其如此）。
				if s := stripToolEcho(out.Content); s != "" {
					return s, usage, a.client.Route(), nil
				}
				// 纯工具回显、无任何正文：才做一次无工具纯文本收敛；若收敛仍剥空，返回诚实兜底。
				if c, cerr := a.client.Generate(ctx, history); cerr == nil {
					c = strings.TrimSpace(c)
					if looksLikeToolEcho(c) {
						c = stripToolEcho(c)
					}
					if c != "" {
						return c, usage, a.client.Route(), nil
					}
				}
				// 模型在「文本回显工具调用」上打转且未产出任何结论文本：低风险下不能凭空生成
				// 数据（那会违背诚实性），因此返回明确的兜底说明，而不是把原始工具标签抛给用户，
				// 也避免伪装成成功结果。
				return toolEchoFallback, usage, a.client.Route(), nil
			}
			if out.Content == "" {
				return "", usage, a.client.Route(), fmt.Errorf("模型最终回复为空")
			}
			return out.Content, usage, a.client.Route(), nil
		}

		// 工具死循环收敛：阻止模型反复发工具调用却迟迟不写结论文本。
		// 即使每轮也带部分 content（如大纲标题），超过阈值轮仍未产出最终正文，即视为轮转过量，
		// 先一步中断并走末尾的无工具纯文本收敛，避免在 240s 预算内空耗搜索轮次。
		toolRounds++
		if toolRounds >= 5 {
			break
		}

		// 依次执行工具并回填结果
		for _, tc := range out.ToolCalls {
			tool, ok := a.tools[tc.Function.Name]
			if !ok {
				return "", usage, a.client.Route(), fmt.Errorf("模型请求了未注册的工具: %s", tc.Function.Name)
			}
			argsBytes := normalizeToolArguments(tc.Function.Arguments)
			fmt.Fprintf(debugWriter, "[agent][diag] 执行工具 %s args=%s\n", tc.Function.Name, truncate(string(argsBytes), 160))
			result, execErr := tool.Execute(ctx, argsBytes)
			fmt.Fprintf(debugWriter, "[agent][diag]   -> err=%v result_head=%q\n", execErr, truncate(result, 120))
			if execErr != nil {
				result = fmt.Sprintf("工具执行失败: %v", execErr)
			}
			history = append(history, Message{
				Role:       "tool",
				ToolCallID: tc.ID,
				Content:    result,
			})
		}
	}
	// 工具死循环/超轮收敛：用无工具纯文本基于「已完成工具结果」重生成一次，作为最终结论，
	// 避免 worker 或单问联网因模型反复 tool_calls 而不产文本。受 ctx 约束。
	// 与上方的工具回显收敛同对待：模型惯性可能使收敛结果仍是 <｜｜ …> 标签，需剥除再返回。
	if len(a.toolSchemas) > 0 {
		if content, cerr := a.client.Generate(ctx, history); cerr == nil {
			content = strings.TrimSpace(content)
			if looksLikeToolEcho(content) {
				content = stripToolEcho(content)
			}
			if content != "" {
				return content, usage, a.client.Route(), nil
			}
			// Generate 成功却只剩工具标签/空白：不凭空编造，返回诚实兜底而非升级 error。
			return toolEchoFallback, usage, a.client.Route(), nil
		}
	}
	return "", usage, a.client.Route(), fmt.Errorf("工具调用超过最大轮数 %d，可能死循环", a.maxRounds)
}

// looksLikeToolEcho 判断文本是否「工具调用回显」——即模型把工具调用意图以
// 文本/标签形式输出（Anthropic style tool_use 或其转义），而非规范 tool_calls。
// 命中后调用方应避免把这串原始标签当最终答案返回。
//
// 判据保持保守：必须命中强特征才判定，避免误伤正常含这些词组的正文。
//   - 「<｜」（DeepSeek 对 Anthropic tool_use 块的非标准 Unicode 转义）几乎不会
//     出现在正常中文自然语言里，出现即基本确定为工具回显；
//   - invoke name="..."（Anthropic 工具调用块的标准文本形式）。
func looksLikeToolEcho(s string) bool {
	if strings.Contains(s, "｜") {
		return true
	}
	if strings.Contains(s, "invoke name=") {
		return true
	}
	return false
}

// 工具回显剥离用的正则：
//   - line：任何以 <｜｜>（转义形式，含开闭标签 <｜｜/>、<｜｜ invoke/…</>）开头的整行，
//     逐行删除（工具标签通常独占一行，不留可见噪声）；
//   - invokePara：标准 Anthropic 文本形式 <invoke>…</invoke>、<parameter>…</parameter> 行；
//   - blank：折叠连续空行。
var (
	reToolEchoLine = regexp.MustCompile(`(?m)^\s*</?[｜|][｜|][^>]*>.*$`)
	reToolInvoke   = regexp.MustCompile(`(?m)^\s*</?invoke\b.*$`)
	reToolPara     = regexp.MustCompile(`(?m)^\s*</?parameter\b.*$`)
	reToolBlank    = regexp.MustCompile(`\n{3,}`)
)

// toolEchoFallback：当模型深陷「文本回显工具调用」并最终未能产出结论文本时的诚实兜底。
// 不凭空编造数据（对齐项目「搜索诚实性」价值观），也不把工具标签抛给用户。
const toolEchoFallback = "当前检索未能获取有效内容。为避免给出未经核实的信息，建议直接访问权威来源（如央行、交易所、官方机构官网）核实，或稍后重试。"

// 文本格式工具调用解析正则：
//   - invokeEcho：一整块 <｜｜ invoke name="X">…</｜｜ /invoke>（DeepSeek 对 Anthropic 文本
//     tool_use 的非标准转义）；同时也兼容标准 <invoke name="X">…</invoke>。
//   - param：块内的 <｜｜ parameter name="K"…>V</｜｜ parameter>。
var (
	reToolInvokeEcho = regexp.MustCompile(`(?s)<[｜|][｜|]?\s*invoke\b.*?</[｜|][｜|]?\s*/?invoke\s*>`)
	reToolParamEcho  = regexp.MustCompile(`(?s)<[｜|][｜|]?\s*parameter\s+name="([^"]+)"[^>]*>(.*?)</[｜|][｜|]?\s*/?parameter\s*>`)
	reToolNameInEcho = regexp.MustCompile(`name="([^"]+)"`)
)

// parseToolEcho 把 Anthropic 文本格式的工具调用块（<｜｜ invoke name="…">…</｜｜ /invoke>）
// 解析成规范 tool_calls（真正触发工具执行），并从 content 中移除这些块只留自然语言。
// 仅当被调用的工具名存在于 allowed（已注册工具集）时才转为 tool_call；未注册的保留原文，
// 避免执行到未知工具而报错。
func parseToolEcho(content string, allowed map[string]bool) (string, []ToolCall) {
	var calls []ToolCall
	clean := reToolInvokeEcho.ReplaceAllStringFunc(content, func(block string) string {
		m := reToolNameInEcho.FindStringSubmatch(block)
		if len(m) < 2 {
			return block // 解析不出 name，保留
		}
		if !allowed[strings.TrimSpace(m[1])] {
			return block // 未注册工具，保留原文，不进执行环
		}
		args := map[string]string{}
		for _, pm := range reToolParamEcho.FindAllStringSubmatch(block, -1) {
			args[strings.TrimSpace(pm[1])] = pm[2]
		}
		argsJSON, err := json.Marshal(args)
		if err != nil {
			return block
		}
		calls = append(calls, ToolCall{
			ID:   fmt.Sprintf("call_echo_%d", len(calls)),
			Type: "function",
			Function: ToolFuncCall{
				Name:      strings.TrimSpace(m[1]),
				Arguments: argsJSON,
			},
		})
		return "" // 删除该工具块，只留纯文本
	})
	// 收尾：移除残余的工具包裹标签（如 <｜｜ calls>…</｜｜ /calls>）及其它行级标签，
	// 只留下自然语言正文。
	clean = stripToolEcho(clean)
	return clean, calls
}

// stripToolEcho 从文本中剥除 Anthropic 风格的工具调用标签（DeepSeek 转义 <｜｜> 或
// 标准 <invoke>/<parameter>），只保留自然语言正文。全部为标签时返回空串。
func stripToolEcho(s string) string {
	s = reToolEchoLine.ReplaceAllString(s, "")
	s = reToolInvoke.ReplaceAllString(s, "")
	s = reToolPara.ReplaceAllString(s, "")
	s = reToolBlank.ReplaceAllString(s, "\n\n")
	return strings.TrimSpace(s)
}

// normalizeToolArguments 处理 arguments 的两种形态：
//   - 标准：{"query": "..."}，直接返回；
//   - 部分 provider 把 arguments 作为 JSON 字符串返回，即 "{\"query\":\"...\"}"，
//     此时先 Unmarshal 成 raw 字节再加解析。存在 key 形如 query 的直接对象即可。
func normalizeToolArguments(raw json.RawMessage) json.RawMessage {
	if len(raw) == 0 {
		return raw
	}
	trimmed := bytes.TrimSpace(raw)
	// 若非字符串字面量（不以 " 开头），视为已是对象
	if len(trimmed) == 0 || trimmed[0] != '"' {
		return trimmed
	}
	// arguments 是 JSON 字符串 → 解开一层
	var s string
	if err := json.Unmarshal(trimmed, &s); err == nil {
		return json.RawMessage(s)
	}
	return trimmed
}

// roundWithEmptyRecovery 发起一轮带工具的模型调用，并对「空响应」做多策略恢复：
//
//  1. 策略一：原始带工具调用。只要模型返回内容或工具调用即为合法响应，直接返回；
//  2. 策略二：仍为空则带工具重试一次（极轻退避，吸收偶发空回复）；
//  3. 策略三：仍为空则降级为「不带工具」的纯文本重试一次，兼容「带工具即空、纯文本正常」
//     的路由（如部分模型对 web_search tools 支持性差）。这是保底，保证调用方拿到非空文本。
//
// 所有策略均受 ctx 约束，ctx 取消立即返回，不阻塞上层（编排 150s 兜底）。
func (a *Agent) roundWithEmptyRecovery(ctx context.Context, history []Message) (*ChatResult, error) {
	// 策略一：原始带工具调用
	out, err := a.client.generateWithTools(ctx, history, a.toolSchemas)
	if err != nil {
		return nil, err
	}
	if out.Content != "" || len(out.ToolCalls) > 0 {
		return out, nil
	}
	lastErr := fmt.Errorf("模型返回空响应")

	// 策略二：带工具重试一次
	if len(a.toolSchemas) > 0 {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-time.After(300 * time.Millisecond):
		}
		if out2, err2 := a.client.generateWithTools(ctx, history, a.toolSchemas); err2 == nil && (out2.Content != "" || len(out2.ToolCalls) > 0) {
			return out2, nil
		} else if err2 != nil {
			lastErr = err2
		}
	}

	// 策略三：降级为不带工具的纯文本重试
	if len(a.toolSchemas) > 0 {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-time.After(300 * time.Millisecond):
		}
		if content, err3 := a.client.Generate(ctx, history); err3 == nil && strings.TrimSpace(content) != "" {
			return &ChatResult{Content: content}, nil
		} else if err3 != nil {
			lastErr = err3
		}
	}

	return nil, fmt.Errorf("模型最终回复为空（带工具重试与无工具降级均未产出内容）: %v", lastErr)
}

// generateWithTools 单轮调用：携带 tools 并解析 tool_calls。
// ctx 贯穿到 HTTP 请求，使 Agent 的多轮 tool_calls 都能被上层超时（如编排 150s 兜底）真正取消。
func (c *Client) generateWithTools(ctx context.Context, messages []Message, toolSchemas []ToolSchema) (*ChatResult, error) {
	reqBody := map[string]any{
		"model":       c.model,
		"messages":    messages,
		"stream":      false,
		"temperature": c.temperature,
		"max_tokens":  c.maxTokens,
	}
	if len(toolSchemas) > 0 {
		reqBody["tools"] = toolSchemas
	}

	resp, err := c.post(ctx, "/chat/completions", reqBody)
	if err != nil {
		return nil, fmt.Errorf("LLM 服务不可达: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("读取 LLM 响应失败: %w", err)
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("LLM 返回 %d: %s", resp.StatusCode, truncate(string(body), 200))
	}

	var raw struct {
		Model   string `json:"model"`
		Choices []struct {
			Message struct {
				Content   string     `json:"content"`
				ToolCalls []ToolCall `json:"tool_calls"`
			} `json:"message"`
			FinishReason string `json:"finish_reason"`
		} `json:"choices"`
		Usage *struct {
			PromptTokens     int `json:"prompt_tokens"`
			CompletionTokens int `json:"completion_tokens"`
			TotalTokens      int `json:"total_tokens"`
		} `json:"usage"`
	}
	if err := json.Unmarshal(body, &raw); err != nil {
		return nil, fmt.Errorf("LLM 响应解析失败: %w", err)
	}
	if len(raw.Choices) == 0 {
		return nil, fmt.Errorf("LLM 未返回任何 choice")
	}

	msg := raw.Choices[0].Message
	content := msg.Content
	toolCalls := msg.ToolCalls

	// 兜底解析「文本格式工具调用」：LMUAI/DeepSeek 等兼容 Anthropic 的 provider 常把工具
	// 调用以 <｜｜ invoke name="…">…</｜｜ /invoke> 文本而非规范 tool_calls JSON 返回。
	// 此时 ToolCalls 为空、Content 是原始标签，若直接返回就把噪声抛给用户、搜索也未真正执行。
	// 这里把文本块解析回规范 tool_calls（触发 Agent 真正运行 web_search），并清理出纯文本。
	if len(toolCalls) == 0 && content != "" && looksLikeToolEcho(content) {
		allowed := make(map[string]bool, len(toolSchemas))
		for _, ts := range toolSchemas {
			allowed[ts.Function.Name] = true
		}
		clean, parsed := parseToolEcho(content, allowed)
		if len(parsed) > 0 {
			content = clean
			toolCalls = parsed
		}
	}

	result := &ChatResult{
		Content:      content,
		Model:        raw.Model,
		FinishReason: raw.Choices[0].FinishReason,
		ToolCalls:    toolCalls,
	}
	if raw.Usage != nil {
		result.Usage.PromptTokens = raw.Usage.PromptTokens
		result.Usage.CompletionTokens = raw.Usage.CompletionTokens
		result.Usage.TotalTokens = raw.Usage.TotalTokens
	}
	return result, nil
}