package runtime

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

	nruntime "agent-native/runtime"
	agentllm "agent-native/runtime/llm"
	agenttools "agent-native/runtime/tools"
	"agent-system-backend/internal/prompts"
	core "agent-system-core"
)

const subagentFallback = "你是被派生的子代理。独立完成给定子任务，输出完整、结构化的结果；需要分点时分条列出，不要为了简短而省略关键内容。"

const defaultSubagentID = "agent.sub.general"

// materialSubagentID 材料整理子代理（附件即保证）：系统确定性派发，不依赖
// 提示词库配置即可工作——spec 缺省在代码内（对话同档模型 + 检索；fs 工具是
// 预留能力默认关，材料通道 = 任务正文内联 + knowledge_search）。
const materialSubagentID = "agent.sub.material"

func fallbackSubagentSpec() prompts.SubagentSpec {
	return prompts.SubagentSpec{
		Tools: []string{"knowledge_search", "regulations_search", "regulation_evidence", "web_search", "tavily_search", "time_now", "list_conversations", "read_conversation", "fs_read_file", "fs_list_dir"},
		Model: "lightweight",
	}
}

func materialSubagentSpec() prompts.SubagentSpec {
	return prompts.SubagentSpec{
		Tools: []string{"knowledge_search", "time_now"},
		Model: "chat",
	}
}

func (l *SubLLMRunner) Name() string { return "llm-subagent" }

func (l *SubLLMRunner) Execute(rc RunCtx) error {
	id := defaultSubagentID
	if src, ok := rc.(interface{ SubagentID() string }); ok {
		if v := strings.TrimSpace(src.SubagentID()); v != "" {
			id = v
		}
	}
	spec := fallbackSubagentSpec()
	if id == materialSubagentID {
		spec = materialSubagentSpec()
	}
	if l.Prompts != nil {
		if got, ok := l.Prompts.SubagentSpec(rc.Profile().Owner, id); ok {
			spec = got
		}
	}
	// 系统派发的模型档覆盖（材料整理 = 对话同档）：优先于 spec 默认。
	if src, ok := rc.(interface{ Submodel() string }); ok && src.Submodel() == "chat" {
		spec.Model = "chat"
	}
	profileTurns, maxTools := loopCaps(rc.Profile().Budget)
	maxTurns := spec.MaxTurns
	if maxTurns <= 0 {
		maxTurns = math.MaxInt32
		if id == materialSubagentID {
			maxTurns = 8
		}
	}
	if profileTurns < maxTurns {
		maxTurns = profileTurns
	}
	route := ResolveLightweightModel(rc.Profile(), rc.Settings())
	if spec.Model == "chat" {
		route = ResolveChatModel(rc.Profile(), rc.Settings())
	}
	stream := l.Stream
	if stream == nil {
		if route.Model == "" || route.Base == "" {
			return fmt.Errorf("llm-subagent: 模型未配置（当前解析：model=%q base=%q）", route.Model, route.Base)
		}
		client := OpenRouteClient(route, rc.Settings())
		stream = client.Stream
	}
	stream = GuardProjectStream(rc, subagentUserText(rc), stream)
	stream = bindProviderRunID(stream, string(rc.Run().ID))
	if spec.Model == "chat" {
		stream = agentllm.WithFallback(stream, ChatFallbacks(rc.Settings()), nil)
	}
	stream = budgetedStream(stream, rc.Profile().Budget, int(rc.Settings().GetNumber("models.chat.maxOutputTokens", 0)))
	decls := subagentDecls(rc, spec)
	st := &subRun{rc: rc, localCheckpoint: checkpointEnabled(rc)}
	if id == materialSubagentID && rc.Profile().Budget.MaxTokens != nil && *rc.Profile().Budget.MaxTokens > 0 {
		st.materialLimit = *rc.Profile().Budget.MaxTokens
		st.materialOutputCap = materialOutputCap(rc)
	}
	if maxTools == 0 && id == materialSubagentID {
		maxTools = 12
	}
	budget := toolBudget(maxTools)
	reg := agenttools.NewRegistry("")
	for _, d := range decls {
		name := d.Name
		reg.Register(directTool{name: name, run: func(ctx context.Context, args json.RawMessage) agenttools.Result {
			return st.runTool(name, string(args))
		}})
	}
	wrapped := func(ctx context.Context, model string, messages []agentllm.Message, _ []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
		opts.MaxTokens = outputTokenCap(rc.Profile().Budget, rc.Settings(), "models.chat.maxOutputTokens")
		requestDecls := decls
		if st.materialLimit > 0 {
			if st.finalOnlyRejected {
				return subagentBudgetError("checkpoint requested after final-only transition")
			}
			finalOnly, stop := st.materialDecision()
			if stop {
				return subagentBudgetError("final-answer reserve unavailable")
			}
			if finalOnly {
				st.finalOnly = true
				requestDecls = make([]agentllm.ToolDecl, 0, len(decls))
				for _, d := range decls {
					if d.Name != workCheckpointTool {
						requestDecls = append(requestDecls, d)
					}
				}
				reminder := "现在用现有材料写出真实最终结论与未决项，不再发布检查点。最终答复须简明，控制在约1500字以内并一次写完，不要逐条复述材料原文。"
				if len(st.savedCheckpointIDs) > 0 {
					reminder = "阶段成果已保存；" + reminder
				}
				// Keep the original first system message in place: the controlled
				// evaluator keys a run from that message and the first user input.
				at := len(messages)
				for i, msg := range messages {
					if msg.Role == agentllm.RoleSystem {
						at = i + 1
						break
					}
				}
				withReminder := append([]agentllm.Message(nil), messages[:at]...)
				withReminder = append(withReminder, agentllm.Message{Role: agentllm.RoleSystem, Content: reminder})
				messages = append(withReminder, messages[at:]...)
			}
		}
		in := stream(ctx, model, messages, requestDecls, opts)
		out := make(chan agentllm.Event)
		go func() {
			defer close(out)
			for ev := range in {
				if ev.Kind == agentllm.KindMessageEnd {
					if ev.Message.Role == "" {
						ev.Message.Role = agentllm.RoleAssistant
					}
					if len(ev.Message.ToolCalls) > 0 && ev.Message.StopReason == "" {
						ev.Message.StopReason = agentllm.StopToolCalls
					}
				}
				select {
				case out <- ev:
				case <-ctx.Done():
					return
				}
			}
		}()
		return out
	}
	execContext := rc.Context()
	var cancel context.CancelFunc
	if id == materialSubagentID {
		execContext, cancel = context.WithTimeout(execContext, 300*time.Second)
		defer cancel()
	}
	_ = rc.Emit(core.EvAgentProgress, core.AgentProgressPayload{AgentRun: rc.Run().ID, Note: "执行子任务……", Phase: core.PhaseActing})
	user := agentllm.Message{Role: agentllm.RoleUser, Content: subagentUserText(rc)}
	window, reserve, buffer := contextWindow(rc.Settings())
	res, err := nruntime.Run(execContext, nruntime.Config{
		Model:                  route.Model,
		SystemPrompt:           l.subagentPrompt(rc, id),
		Tools:                  reg,
		Stream:                 wrapped,
		MaxTurns:               maxTurns,
		Budget:                 budget,
		RepeatGuard:            true,
		ContinueOnLength:       true,
		MaxLengthContinuations: 2,
		ContextWindow:          window,
		OutputReserve:          reserve,
		ContextBuffer:          buffer,
		OnEvent:                st.onEvent,
		PermissionAsk:          PermissionAskOf(rc),
		Trace:                  l.Trace,
		AgentName:              agentNameOf(rc.Profile()),
		RetryAttempts:          retryAttemptsOf(rc.Settings()),
	}, user)
	// Successful turns reach OnEvent(message_end); an error-bearing terminal
	// retains returned Usage in Result but does not emit that event. Reconcile
	// tokens from the final result before persisting a failed child run, while
	// retaining the tool count collected from executable turns.
	// A failed turn can omit total_tokens even after earlier turns reported it.
	// In that case Result.TotalTokens is only a partial sum, while the prompt
	// and completion totals include the failed turn's returned Usage.
	reported := max(res.Usage.TotalTokens, res.Usage.PromptTokens+res.Usage.CompletionTokens)
	if reported > st.usageTotal.TotalTokens {
		st.usageTotal.PromptTokens = res.Usage.PromptTokens
		st.usageTotal.CompletionTokens = res.Usage.CompletionTokens
		st.usageTotal.TotalTokens = reported
		st.usageTotal.LastPromptTokens = res.Usage.LastPromptTokens
		st.usageTotal.LastCompletionTokens = res.Usage.LastCompletionTokens
		st.usageTotal.CacheReadTokens = res.Usage.CacheReadTokens
		st.usageTotal.CacheWriteTokens = res.Usage.CacheWriteTokens
	}
	if sink, ok := rc.(interface{ AddUsage(core.Usage) }); ok {
		sink.AddUsage(st.usageTotal)
	}
	st.saveReasoning()
	if len(st.parts) == 0 {
		st.parts = append(st.parts, core.Part{Type: core.PartText, Text: "（子代理未返回内容）"})
	}
	_ = rc.SetAssistantParts(st.parts)
	if err != nil {
		return err
	}
	if res.Err != nil {
		return res.Err
	}
	return nil
}

func (l *SubLLMRunner) subagentPrompt(rc RunCtx, id string) string {
	if l.Prompts == nil {
		return subagentFallback
	}
	if body, err := l.Prompts.AgentPrompt(rc.Profile().Owner, id); err == nil && strings.TrimSpace(body) != "" {
		return body
	}
	return subagentFallback
}

func subagentUserText(rc RunCtx) string {
	text := rc.Subtask()
	if src, ok := rc.(interface{ Subhint() string }); ok {
		if hint := strings.TrimSpace(src.Subhint()); hint != "" {
			text += "\n\n执行提示：" + hint
		}
	}
	return text
}

func subagentDecls(rc RunCtx, spec prompts.SubagentSpec) []agentllm.ToolDecl {
	allow := map[string]bool{}
	for _, name := range spec.Tools {
		allow[name] = true
	}
	var decls []agentllm.ToolDecl
	for _, tool := range rc.Tools() {
		if !allow[tool.Name] || !rc.Profile().Tools.ToolAllowed(tool.Name) {
			continue
		}
		decls = append(decls, agentllm.ToolDecl{Name: tool.Name, Description: tool.Description, Parameters: tool.Parameters})
	}
	if checkpointEnabled(rc) {
		decls = append(decls, checkpointDecl())
	}
	return decls
}

type subSlot struct {
	call     agentllm.ToolCall
	start    time.Time
	summary  string
	status   string
	artifact *core.ArtifactRef
}

type subRun struct {
	localCheckpoint    bool
	finalOnly          bool
	finalOnlyRejected  bool
	materialLimit      int
	materialOutputCap  int
	savedCheckpointIDs map[string]struct{}
	responses          int
	usageKnown         bool
	rc                 RunCtx
	parts              []core.Part
	queued             []agentllm.ToolCall
	inflight           *subSlot
	turnText           strings.Builder
	turnThink          strings.Builder
	thinkStart         time.Time
	thinkEnd           time.Time
	usageTotal         core.Usage
}

func (st *subRun) onEvent(ev nruntime.Event) {
	switch ev.Type {
	case "turn_start":
		st.turnText.Reset()
		st.turnThink.Reset()
		st.thinkStart = time.Time{}
		st.thinkEnd = time.Time{}
	case "text_delta":
		st.turnText.WriteString(ev.Text)
		_ = st.rc.Emit(core.EvAgentProgress, core.AgentProgressPayload{AgentRun: st.rc.Run().ID, Note: truncateRun(st.turnText.String(), 80)})
	case "reasoning_delta":
		now := time.Now()
		if st.thinkStart.IsZero() {
			st.thinkStart = now
		}
		st.thinkEnd = now
		st.turnThink.WriteString(ev.Text)
		_ = st.rc.Emit(core.EvThinkingDelta, core.ThinkingDeltaPayload{Index: len(st.parts), Text: ev.Text})
	case "message_end":
		st.responses++
		reportedTotal := ev.TotalTokens
		if reportedTotal == 0 {
			reportedTotal = ev.PromptTokens + ev.CompletionTokens
		}
		if st.responses == 1 {
			st.usageKnown = reportedTotal > 0
		} else {
			st.usageKnown = st.usageKnown && reportedTotal > 0
		}
		st.usageTotal.PromptTokens += ev.PromptTokens
		st.usageTotal.CompletionTokens += ev.CompletionTokens
		st.usageTotal.TotalTokens += reportedTotal
		st.usageTotal.ToolCalls += len(ev.Message.ToolCalls)
		st.usageTotal.LastPromptTokens = ev.PromptTokens
		st.usageTotal.LastCompletionTokens = ev.CompletionTokens
		st.usageTotal.CacheReadTokens += ev.CacheReadTokens
		st.usageTotal.CacheWriteTokens += ev.CacheWriteTokens
		st.queued = append([]agentllm.ToolCall(nil), ev.Message.ToolCalls...)
		hasCalls := len(ev.Message.ToolCalls) > 0
		st.saveReasoning()
		if body := strings.TrimSpace(st.turnText.String()); body != "" {
			st.parts = append(st.parts, core.Part{Type: core.PartText, Text: body, Process: hasCalls})
		}
	case "tool_execution_start":
		call := agentllm.ToolCall{Name: ev.Tool, Arguments: ev.Detail}
		if len(st.queued) > 0 {
			call = st.queued[0]
			st.queued = st.queued[1:]
		}
		st.inflight = &subSlot{call: call, start: time.Now()}
		_ = st.rc.Emit(core.EvToolCall, core.ToolCallPayload{CallID: call.ID, Name: call.Name, Arguments: call.Arguments})
		_ = st.rc.Emit(core.EvAgentProgress, core.AgentProgressPayload{AgentRun: st.rc.Run().ID, Note: "调用 " + call.Name})
	case "tool_execution_end":
		slot := st.inflight
		if slot == nil {
			return
		}
		end := time.Now()
		_ = st.rc.Emit(core.EvToolResult, core.ToolResultPayload{CallID: slot.call.ID, Status: slot.status, Summary: slot.summary})
		st.parts = append(st.parts,
			core.Part{Type: core.PartToolCall, Process: true, ToolCall: &core.ToolCallPart{ID: slot.call.ID, Name: slot.call.Name, Arguments: slot.call.Arguments, Duration: durationMs(slot.start, end)}},
			core.Part{Type: core.PartToolResult, Process: true, ToolResult: &core.ToolResultPart{CallID: slot.call.ID, Status: slot.status, Summary: slot.summary}},
		)
		if slot.artifact != nil {
			st.parts = append(st.parts, core.Part{Type: core.PartArtifactRef, Artifact: slot.artifact})
		}
		st.inflight = nil
	}
}

// Provider-supplied public reasoning uses the same retained process part as
// the foreground. No request prompt or hidden transport payload is copied.
func (st *subRun) saveReasoning() {
	if st.turnThink.Len() == 0 {
		return
	}
	st.parts = append(st.parts, core.Part{Type: core.PartThinking, Text: st.turnThink.String(), Process: true, DurationMs: durationMs(st.thinkStart, st.thinkEnd)})
	st.turnThink.Reset()
}

func (st *subRun) runTool(name, args string) agenttools.Result {
	if name == workCheckpointTool && st.localCheckpoint {
		if st.finalOnly {
			st.finalOnlyRejected = true
			res := agenttools.Result{Content: "最终答复阶段不再接受检查点", IsError: true}
			if st.inflight != nil {
				st.inflight.summary = res.Content
				st.inflight.status = "error"
			}
			return res
		}
		var submitted struct {
			CheckpointID string `json:"checkpointId"`
		}
		_ = json.Unmarshal([]byte(args), &submitted)
		res := publishCheckpoint(st.rc, args)
		if !res.IsError && submitted.CheckpointID != "" {
			if st.savedCheckpointIDs == nil {
				st.savedCheckpointIDs = make(map[string]struct{})
			}
			st.savedCheckpointIDs[submitted.CheckpointID] = struct{}{}
		}
		res = st.withFinalOnlyNotice(res)
		if st.inflight != nil {
			st.inflight.summary = res.Content
			st.inflight.status = "ok"
			if res.IsError {
				st.inflight.status = "error"
			}
		}
		return res
	}
	call := agentllm.ToolCall{Name: name, Arguments: args}
	if st.inflight != nil && st.inflight.call.Name != "" {
		call = st.inflight.call
	}
	res := st.rc.ExecuteTool(core.ToolCallRequest{CallID: call.ID, Name: call.Name, Arguments: json.RawMessage(call.Arguments)})
	summary := res.Output
	status := res.Status
	if status == "" {
		status = "ok"
	}
	var ref *core.ArtifactRef
	if res.Artifact != nil {
		ref = &core.ArtifactRef{ArtifactID: res.Artifact.ID, Name: res.Artifact.Name, Kind: res.Artifact.Kind}
	}
	out := agenttools.Result{Content: summary, IsError: status != "ok"}
	out = st.withFinalOnlyNotice(out)
	if st.inflight != nil {
		st.inflight.summary = out.Content
		st.inflight.status = status
		st.inflight.artifact = ref
	}
	return out
}

func subagentBudgetError(reason string) <-chan agentllm.Event {
	out := make(chan agentllm.Event, 1)
	out <- agentllm.Event{Kind: agentllm.KindError, Err: fmt.Errorf("budget_exceeded: %s", reason), ErrorMessage: reason}
	close(out)
	return out
}

// delegateMaterialsNote 模型自派后台任务的材料指引（R2）：子代理看不到父回合
// 内联的附件全文，指明 knowledge_search 检索通道（上传时解析文本已进索引）。
func delegateMaterialsNote(docs []core.ParsedDocument) string {
	names := make([]string, 0, len(docs))
	for _, d := range docs {
		if n := strings.TrimSpace(d.Meta.Filename); n != "" {
			names = append(names, n)
		}
	}
	if len(names) == 0 {
		return ""
	}
	return "\n\n[系统补充] 本会话材料文件：" + strings.Join(names, "、") + "。解析全文可用 knowledge_search 以文件名为关键词检索获取。"
}

func subagentProfileID(agent string) string {
	switch agent {
	case "explore", "research", "plan", "verify", "general", "writer":
		return "agent.sub." + agent
	default:
		return defaultSubagentID
	}
}