package eino

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

	"github.com/cloudwego/eino/components/tool"
	"github.com/cloudwego/eino/compose"
	"github.com/cloudwego/eino/flow/agent/react"
	"github.com/cloudwego/eino/schema"

	"github.com/sundynix/sundynix-dispatcher/internal/dsl"
	"github.com/sundynix/sundynix-dispatcher/internal/harness"
	"github.com/sundynix/sundynix-shared/contract"
)

// reactMaxStep 限制 ReAct 推理-工具循环的步数上限（控成本/时延）。
// 研究型任务（搜索→抓取→推理 多轮）8 步偏紧易触顶报错；默认 12，可经 REACT_MAX_STEP 调。
func reactMaxStep() int { return envInt("REACT_MAX_STEP", 12) }

// streamHasToolCall 扫描模型整段流输出，任一片段含 tool call 即判定为工具调用。
// 比默认"只看首片段"鲁棒（兼容先吐文本/思考再给 tool call 的模型，如 deepseek）。
// 契约要求：返回前必须 Close 传入的流。
func streamHasToolCall(_ context.Context, sr *schema.StreamReader[*schema.Message]) (bool, error) {
	defer sr.Close()
	for {
		msg, err := sr.Recv()
		if err == io.EOF {
			return false, nil
		}
		if err != nil {
			return false, err
		}
		if len(msg.ToolCalls) > 0 {
			return true, nil
		}
	}
}

// mcpTool 把一个 MCP 工具（NATS 那头）适配成 Eino InvokableTool：
// 模型给的参数(JSON) + 运行时注入的绑定参数(uid/kb，不暴露给模型) 合并后经 NATS 调 MCP。
type mcpTool struct {
	info    *schema.ToolInfo
	mcpName string              // MCP 侧真实工具名（可与 info.Name 不同）
	subject func(string) string // contract.ToolSubjectGo/Py
	bind    map[string]any      // 运行时注入参数
	caller  ToolCaller
	taskID  string
	tr      *execTracer
	// timeout 该工具自报的超时预算（0=用默认 toolCallTimeout）。
	// 本地执行类工具要等用户在桌面端点确认框（人的反应时间）+ 真跑命令，
	// 远超默认 3 秒；由工具在 list_tools 里声明 timeout_sec，别在这硬编码工具名。
	timeout time.Duration
}

func (m *mcpTool) Info(_ context.Context) (*schema.ToolInfo, error) { return m.info, nil }

// InvokableRun 执行一次工具调用：合并参数 → NATS 调 MCP → 返回观察给模型；
// 调用本身落一条 ExecEvent 轨迹（"模型自己调了哪个工具"在观测里可见）。
func (m *mcpTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
	args := map[string]any{}
	if argsJSON != "" {
		_ = json.Unmarshal([]byte(argsJSON), &args)
	}
	for k, v := range m.bind {
		args[k] = v
	}
	log.Printf("[react] 模型自主调用工具 %s (mcp=%s) task=%s args=%s", m.info.Name, m.mcpName, m.taskID, truncate(argsJSON, 120))
	end := m.tr.span("tool:"+m.mcpName, "tool", "模型自主调用 "+m.info.Name)
	budget := toolCallTimeout
	if m.timeout > 0 {
		budget = m.timeout
	}
	cctx, cancel := context.WithTimeout(ctx, budget)
	defer cancel()
	res, err := m.caller.CallTool(cctx, m.subject(m.mcpName), &contract.ToolCall{Tool: m.mcpName, TaskID: m.taskID, Args: args})
	if err != nil {
		end("调用失败", err)
		return "工具调用失败：" + err.Error(), nil // 作为观察返回，不中断 ReAct
	}
	if res == nil || !res.OK {
		msg := "工具无结果"
		if res != nil && res.Error != "" {
			msg = res.Error
		}
		end(msg, nil)
		return msg, nil
	}
	end("入参 "+truncate(argsJSON, 120)+" → "+truncate(res.Content, 160), nil)
	return res.Content, nil
}

// toolCatalogEntry 是 MCP list_tools 上报的一条工具元信息（与 mcp-go listTools 输出对齐）。
type toolCatalogEntry struct {
	Name      string `json:"name"`
	CN        string `json:"cn"`
	Desc      string `json:"desc"`
	Agent     bool   `json:"agent_exposed"`
	AgentName string `json:"agent_name"`
	Params    []struct {
		Name     string `json:"name"`
		Type     string `json:"type"`
		Desc     string `json:"desc"`
		Required bool   `json:"required"`
	} `json:"params"`
	Inject []string `json:"inject"`
	// TimeoutSec 工具自报的超时预算（秒，0=用默认）。需要人工确认或长耗时的工具靠它突破默认 3 秒。
	TimeoutSec int `json:"timeout_sec"`
}

// agentTools 动态构建 ReAct 可用的工具集：分别向 mcp-go / mcp-py 探 list_tools 自描述目录，
// 取 agent_exposed 的工具按上报参数 schema 建 InvokableTool；inject 参数（user_id/session_id/
// kb/task_id）服务端运行时绑定、不暴露给模型。某台 MCP 离线即跳过（降级）。
// 新增工具只需在对应 MCP 注册表标 agent，无需改这里（杜绝硬编码）。
func (o *Orchestrator) agentTools(b *board, taskID string, tr *execTracer) []tool.BaseTool {
	if o.tools == nil {
		return nil
	}
	var out []tool.BaseTool
	out = append(out, o.discoverTools(contract.ToolSubjectGo, b, taskID, tr)...)
	out = append(out, o.discoverTools(contract.ToolSubjectPy, b, taskID, tr)...)
	// 平台工具族（gateway 提供，JARVIS 大脑中枢）：查任务/派报告等平台操作，同协议动态发现。
	out = append(out, o.discoverTools(contract.ToolSubjectPlatform, b, taskID, tr)...)
	return out
}

// discoverTools 向某台 MCP（subject 前缀决定 go/py）探 list_tools，把 agent_exposed 的工具
// 转成 Eino InvokableTool。该 MCP 不可用 / 无应答时返回空（不阻断）。
func (o *Orchestrator) discoverTools(subject func(string) string, b *board, taskID string, tr *execTracer) []tool.BaseTool {
	cctx, cancel := context.WithTimeout(context.Background(), toolCallTimeout)
	defer cancel()
	res, err := o.tools.CallTool(cctx, subject("list_tools"), &contract.ToolCall{Tool: "list_tools"})
	if err != nil || res == nil || !res.OK {
		return nil
	}
	var cat struct {
		Tools []toolCatalogEntry `json:"tools"`
	}
	if json.Unmarshal([]byte(res.Content), &cat) != nil {
		return nil
	}
	injectVal := map[string]any{"user_id": b.uid, "session_id": b.sid, "task_id": taskID, "kb": b.kb, "tenant_id": b.tenant}

	var out []tool.BaseTool
	for _, e := range cat.Tools {
		if !e.Agent {
			continue
		}
		params := map[string]*schema.ParameterInfo{}
		for _, p := range e.Params {
			params[p.Name] = &schema.ParameterInfo{Type: schema.DataType(p.Type), Desc: p.Desc, Required: p.Required}
		}
		bind := map[string]any{}
		for _, inj := range e.Inject {
			if v, ok := injectVal[inj]; ok && v != "" {
				bind[inj] = v
			}
		}
		name := e.AgentName
		if name == "" {
			name = e.Name
		}
		out = append(out, &mcpTool{
			mcpName: e.Name,
			subject: subject,
			caller:  o.tools, taskID: taskID, tr: tr,
			bind:    bind,
			timeout: time.Duration(e.TimeoutSec) * time.Second,
			info: &schema.ToolInfo{
				Name: name, Desc: e.Desc,
				ParamsOneOf: schema.NewParamsOneOfByParams(params),
			},
		})
	}
	return out
}

// runReactAgent 执行带"自主工具"的 agent 节点：模型在 ReAct 循环里自行决定调哪些 MCP 工具。
// 模型不支持函数调用 / 无工具时降级回普通 runAgent。最终答复流式回流；工具调用由适配器落轨迹。
func (o *Orchestrator) runReactAgent(ctx context.Context, taskID string, b *board, system string, n dsl.Node, tr *execTracer, node string) {
	tcm := o.agentPool(b).ToolCallingModel()
	tools := o.agentTools(b, taskID, tr)
	if tcm == nil || len(tools) == 0 {
		tr.info(node, "system", "ReAct 降级", "模型不支持函数调用或无可用工具，退回普通对话")
		o.runAgent(ctx, taskID, b, system, tr, node, labelOf(n, "ReAct 智能体"))
		return
	}

	ag, err := react.NewAgent(ctx, &react.AgentConfig{
		ToolCallingModel: tcm,
		ToolsConfig:      compose.ToolsNodeConfig{Tools: tools},
		MaxStep:          reactMaxStep(),
		// 默认检查器只看首个流片段；deepseek 等常先吐文本再给 tool call → 漏判。
		// 改成扫描整段流，任一片段含 tool call 即判定为工具调用。
		StreamToolCallChecker: streamHasToolCall,
	})
	if err != nil {
		tr.emit(node, "model", "error", "构建 ReAct 智能体", err.Error(), 0)
		o.runAgent(ctx, taskID, b, system, tr, node, labelOf(n, "ReAct 智能体"))
		return
	}

	rc := &RunCtx{
		UserID: b.uid, SessionID: b.sid,
		System: firstNonEmpty(system, defaultAgentSystem),
		Query:  b.query,
		// 自主 agent 不预注入画像：让它经 recall_user_memory 工具按需自取（否则模型直接答、不调工具）。
		Profile:  "",
		History:  b.history,
		ToolOut:  append(append([]string{}, b.toolOut...), b.refs...),
		Upstream: append([]string{}, b.agentOut...), // 前序协作 agent 产出 → 接力
	}
	msgs, _ := buildMessages(ctx, rc)
	// 成本护栏：计入输入 token；触顶则中止整图。
	if bud := harness.BudgetFrom(ctx); bud != nil {
		for _, m := range msgs {
			bud.AddPrompt(m.Content)
		}
		if bud.Exceeded() {
			if b.fatalErr == nil {
				b.fatalErr = errBudget
			}
			tr.emit(node, "system", "error", "token 预算", "已达单任务预算上限，中止", 0)
			return
		}
	}

	tr.emit(node, "model", "start", "ReAct 智能体（自主调工具）", fmt.Sprintf("%d 个工具可用", len(tools)), 0)
	t0 := time.Now()
	sr, err := ag.Stream(ctx, msgs)
	if err != nil {
		tr.emit(node, "model", "error", "ReAct 智能体", err.Error(), time.Since(t0).Milliseconds())
		return
	}
	defer sr.Close()
	o.streamAgentReply(ctx, taskID, b, sr, tr, node, "ReAct 智能体", t0)
}

// streamAgentReply 把一个 agent / 协调者的流式输出回流到 sink：跨分片脱敏 + 计输出 token +
// 落产出(供下游接力) + 轨迹收尾。runReactAgent 与 runCoordinator 共用，杜绝两路尾段漂移。
func (o *Orchestrator) streamAgentReply(ctx context.Context, taskID string, b *board, sr *schema.StreamReader[*schema.Message], tr *execTracer, node, label string, t0 time.Time) {
	chunks := 0
	var produced strings.Builder       // 本节点产出（供下游 agent 接力）
	red := harness.NewStreamRedactor() // 输出护栏：跨分片脱敏，杜绝密钥被切断而漏检
	emit := func(safe string) {
		if safe == "" {
			return
		}
		_ = o.sink.PublishToken(taskID, []byte(safe))
		produced.WriteString(safe)
		chunks++
	}
	for {
		chunk, rerr := sr.Recv()
		if rerr == io.EOF {
			break
		}
		if rerr != nil {
			tr.emit(node, "model", "error", label, rerr.Error(), time.Since(t0).Milliseconds())
			return
		}
		if chunk.Content == "" {
			continue // 工具调用片段无正文，跳过；正文只来自模型答复
		}
		emit(red.Push(chunk.Content))
	}
	emit(red.Flush()) // 吐出暂留尾部
	if bud := harness.BudgetFrom(ctx); bud != nil {
		bud.AddComplete(produced.String()) // 成本护栏：计入输出 token
	}
	o.recordAgentOutput(b, produced.String())
	tr.emit(node, "model", "end", label,
		fmt.Sprintf("%d 段输出 / %d 字", chunks, len([]rune(produced.String()))), time.Since(t0).Milliseconds())
}