package agent

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"strings"
	"sync"
	"sync/atomic"

	"git.gopher.ink/mirrors2/AutoZero/llm"
	"git.gopher.ink/mirrors2/AutoZero/session"
	"git.gopher.ink/mirrors2/AutoZero/tools"
)

// maxSubAgentTasks 限制单次 agent 调用的并行子任务数，防止一次 fan-out 打爆。
const maxSubAgentTasks = 16

// subAgentSeq 给每个子 agent 一个唯一序号，用作落盘文件名的命名空间前缀
// （并行子 agent 的 tool_call ID 会重复）。
var subAgentSeq atomic.Int64

// SubAgentTool 构造 "agent" 工具：把自包含子任务委派给全新的嵌套 agent
// （借鉴 ccb AgentTool 的"递归复用主循环"与 picoclaw SubTurn 的会话隔离/深度/并发上限）。
// 嵌套 agent 使用独立 ephemeral 会话（不落盘、不污染主会话上下文），
// 信任授权从父会话继承。
func (a *Agent) SubAgentTool() *tools.Tool {
	return newSubAgentTool(a)
}

type subAgentInput struct {
	Prompt string   `json:"prompt"` // 单任务（与历史行为逐字节一致）
	Tasks  []string `json:"tasks"`  // 多任务：并行 fan-out
	// SessionID 继续一个已有的子对话：上下文还在（同一进程内），不必重发前情。
	SessionID string `json:"session_id"`
}

// subAgentOutcome 是单个子任务的执行结果。
type subAgentOutcome struct {
	text      string
	sessionID string
	err       error
}

// 子会话表是进程级的：一次 agent 调用产生的内存会话留在这里，后续带
// session_id 的调用可以接着聊。与 bash_bg 同一已知边界——进程重启即失。
const subAgentSessionKeep = 32

var (
	subAgentSessMu    sync.Mutex
	subAgentSessions  = map[string]*session.Session{}
	subAgentSessOrder []string
)

func registerSubAgentSession(sess *session.Session) string {
	id := fmt.Sprintf("sa-%d", subAgentSeq.Add(1))
	subAgentSessMu.Lock()
	subAgentSessions[id] = sess
	subAgentSessOrder = append(subAgentSessOrder, id)
	for len(subAgentSessOrder) > subAgentSessionKeep {
		delete(subAgentSessions, subAgentSessOrder[0])
		subAgentSessOrder = subAgentSessOrder[1:]
	}
	subAgentSessMu.Unlock()
	return id
}

func lookupSubAgentSession(id string) *session.Session {
	subAgentSessMu.Lock()
	defer subAgentSessMu.Unlock()
	return subAgentSessions[id]
}

func newSubAgentTool(parent *Agent) *tools.Tool {
	desc := "Delegate a self-contained subtask (research, exploration, bulk refactoring) " +
		"to a fresh sub-agent with the same tools. The sub-agent runs in an isolated context; " +
		"only its final answer comes back. Give it a complete, self-sufficient prompt. " +
		"The result carries a sessionID; pass it back as session_id to continue that exact " +
		"sub-agent conversation without re-sending the earlier context (same process only). " +
		"Use `tasks` (array) to delegate several independent subtasks at once — they run in " +
		"parallel, share the same workspace, and must not touch the same files."
	if parent.CanSpawnSubAgent() {
		desc += " Nested sub-agents may delegate further."
	} else {
		desc += " Sub-agents spawned here cannot delegate further."
	}
	return &tools.Tool{
		Name:        "agent",
		Description: desc,
		SideEffect:  true, // 子 agent 可能执行副作用工具：整体过一次信任门控
		Parameters: map[string]any{
			"type": "object",
			"properties": map[string]any{
				"prompt": map[string]any{
					"type":        "string",
					"description": "Complete, self-sufficient task description with all needed context",
				},
				"tasks": map[string]any{
					"type":        "array",
					"description": "Several independent self-contained subtasks to run in parallel",
					"items":       map[string]any{"type": "string"},
				},
				"session_id": map[string]any{
					"type":        "string",
					"description": "Continue a specific previous sub-agent conversation (the sessionID from its result). Omit to start a new one.",
				},
			},
		},
		Run: func(ctx context.Context, input string) (tools.Result, error) {
			if !parent.CanSpawnSubAgent() {
				return tools.Result{}, fmt.Errorf("max sub-agent depth (%d) reached", parent.MaxDepth)
			}
			// 在 fan-out 之前把 stderr 包成并发安全的 writer——必须在并发之前做，
			// 因为它会改写 parent.Opt.Stderr（放在 goroutine 里就是 data race）。
			parent.ensureLockedStderr()
			var p subAgentInput
			if err := json.Unmarshal([]byte(input), &p); err != nil {
				return tools.Result{}, fmt.Errorf("invalid agent tool input: %w", err)
			}
			if strings.TrimSpace(p.SessionID) != "" && len(p.Tasks) > 0 {
				return tools.Result{}, fmt.Errorf("session_id continues one conversation; it cannot be combined with tasks")
			}
			tasks, err := normaliseSubAgentTasks(p)
			if err != nil {
				return tools.Result{}, err
			}
			if len(tasks) == 1 {
				// 单任务：返回子 agent 最终回复与可续的 sessionID。
				out := parent.runSubAgent(ctx, tasks[0], strings.TrimSpace(p.SessionID))
				if out.err != nil {
					if out.sessionID != "" {
						return tools.Result{}, fmt.Errorf("%w (sessionID: %s)", out.err, out.sessionID)
					}
					return tools.Result{}, out.err
				}
				return tools.TextResult(fmt.Sprintf(
					"<subagent sessionID=%q state=\"completed\">\n%s\n</subagent>",
					out.sessionID, out.text)), nil
			}
			return parent.runSubAgentBatch(ctx, tasks), nil
		},
	}
}

// normaliseSubAgentTasks 归一化 prompt / tasks 两种入参，返回非空任务列表。
func normaliseSubAgentTasks(p subAgentInput) ([]string, error) {
	raw := p.Tasks
	if strings.TrimSpace(p.Prompt) != "" {
		if len(raw) > 0 {
			return nil, fmt.Errorf("agent tool takes either prompt or tasks, not both")
		}
		raw = []string{p.Prompt}
	}
	if len(raw) == 0 {
		return nil, fmt.Errorf("agent tool requires a non-empty prompt or tasks")
	}
	if len(raw) > maxSubAgentTasks {
		return nil, fmt.Errorf("too many tasks (%d > %d)", len(raw), maxSubAgentTasks)
	}
	out := make([]string, 0, len(raw))
	for i, t := range raw {
		t = strings.TrimSpace(t)
		if t == "" {
			return nil, fmt.Errorf("task %d is empty", i+1)
		}
		out = append(out, t)
	}
	return out, nil
}

// runSubAgentBatch 执行多个子任务。只有**问不到人**时才并行——判据是"有没有审批
// 通道"（回调或 stdin），不是 ApproveAll：危险门不看 ApproveAll（见 execOne），
// 而宿主端同一时刻只有一个待裁决槽位（WebUI 的 run.pending 会被后一条直接盖掉），
// 于是「完全访问」档的 WebUI 回合既是 ApproveAll 又有回调，并行 fan-out 会让前一条
// 询问挂到超时。真无人值守（-p / cron / 渠道）两个通道都没有，才谈得上并行。
func (a *Agent) runSubAgentBatch(ctx context.Context, tasks []string) tools.Result {
	conc := 1
	if a.Opt.Approval == nil && a.Opt.Stdin == nil && a.SubAgentConcurrency > 1 {
		conc = a.SubAgentConcurrency
	} else if a.SubAgentConcurrency > 1 {
		fmt.Fprintln(a.Opt.Stderr, "[agent] running sub-tasks serially (an approver is reachable; parallel would contend for the single approval slot)")
	}
	fmt.Fprintf(a.Opt.Stderr, "[agent] delegating %d sub-tasks (concurrency %d)\n", len(tasks), conc)

	outs := make([]subAgentOutcome, len(tasks))
	if conc == 1 {
		for i := range tasks {
			outs[i] = a.runSubAgent(ctx, tasks[i], "")
		}
	} else {
		sem := make(chan struct{}, conc)
		var wg sync.WaitGroup
		for i := range tasks {
			wg.Add(1)
			go func(i int) {
				defer wg.Done()
				select {
				case sem <- struct{}{}:
				case <-ctx.Done():
					outs[i] = subAgentOutcome{err: ctx.Err()}
					return
				}
				defer func() { <-sem }()
				outs[i] = a.runSubAgent(ctx, tasks[i], "")
			}(i)
		}
		wg.Wait()
	}

	var b strings.Builder
	okCount := 0
	for _, o := range outs {
		if o.err == nil {
			okCount++
		}
	}
	fmt.Fprintf(&b, "Delegated %d sub-tasks (concurrency %d; %d ok, %d failed):\n",
		len(tasks), conc, okCount, len(tasks)-okCount)
	for i, o := range outs {
		if o.err != nil {
			fmt.Fprintf(&b, "\n===== [%d/%d] failed =====\n%v\n", i+1, len(tasks), o.err)
			continue
		}
		fmt.Fprintf(&b, "\n===== [%d/%d] ok =====\n%s\n", i+1, len(tasks), o.text)
	}
	text := strings.TrimRight(b.String(), "\n")
	if okCount == 0 {
		// 全军覆没：以失败返回，让父级的失败熔断计数生效。
		return tools.Result{Text: text, Err: fmt.Errorf("all %d sub-tasks failed", len(tasks))}
	}
	return tools.TextResult(text)
}

// runSubAgent 跑一个子任务：新建嵌套 agent（独立会话、深度+1），返回其最终回复。
// sessionID 非空时复用进程内的子会话，让模型能接着上一轮追问，而不必重发前情。
func (a *Agent) runSubAgent(ctx context.Context, prompt, sessionID string) subAgentOutcome {
	trusted := make(map[string]bool, len(a.trusted))
	for k, v := range a.trusted {
		trusted[k] = v
	}
	var sess *session.Session
	if sessionID != "" {
		sess = lookupSubAgentSession(sessionID)
		if sess == nil {
			return subAgentOutcome{
				sessionID: sessionID,
				err:       fmt.Errorf("no such subagent session %q (sessions live in this a2go process)", sessionID),
			}
		}
	} else {
		sess = &session.Session{} // 内存会话：不落盘，但留在子会话表里供后续续聊
		sessionID = registerSubAgentSession(sess)
	}
	fmt.Fprintln(a.Opt.Stderr, "[agent] sub-agent started")

	// 每层重建工具表：必须去掉父级的 agent 闭包（它绑的是父级，会让深度失效），
	// 再按本层深度重新注册指向自己的 agent 工具。plan/memory 一律不继承——它们由
	// 拥有者维护；子 agent 经 System（sysNow 的冻结快照）**只读**看到父级记忆与计划。
	//
	// ask_user 也必须去掉：那张表里的工具闭包绑的是**父 agent**，子 agent 调用它
	// 等于经父级的通道去问人、并把父级的回合收尾标志设上。ZCode 能给子 agent 用
	// 是因为它专门做了"路由回父会话"，那是另一期工程；在这里挂住等一个没人能答的
	// 问题才是真的糟。
	reg := a.Tools.CloneWithout("agent", "plan", "memory", "ask_user")

	nested := &Agent{
		Provider: a.Provider,
		Tools:    reg,
		Session:  sess, // 内存会话：不落盘；续聊靠进程内的子会话表
		// 父计划的只读快照：子 agent 能看到当前进展，但不能修改（无 plan 工具）。
		System:    a.sysNow(),
		Thinking:  a.Thinking,  // 继承父级思考深度
		ResultDir: a.ResultDir, // 子 agent 的超长结果同样落盘可回捞
		// 深度与并发：子 agent 继承父级上限，深度 +1。
		Depth:               a.Depth + 1,
		MaxDepth:            a.MaxDepth,
		SubAgentConcurrency: a.SubAgentConcurrency,
		spillNS:             fmt.Sprintf("sa%d-", subAgentSeq.Add(1)),
		Opt: Options{
			ApproveAll: a.Opt.ApproveAll,
			Stdin:      a.Opt.Stdin,
			Stdout:     io.Discard, // 子过程输出不与主会话混流
			// 继承同一个（已加锁的）writer 实例：并行子 agent 共享父级 stderr。
			// 保证由本 agent 自己的 Run 提前完成（见工具 Run 开头）。
			Stderr: a.Opt.Stderr,
			// 审批通道必须继承：它是**宿主**的回调（TUI 面板 / WebUI 弹框），不是绑在
			// 某个 agent 上的闭包。此前只带了 Stdin，于是 WebUI 里子 agent 撞上危险门
			// 时既没有回调也没有 stdin（那边 stdin 传的就是 nil），询问直接落到
			// "无人可问→拒绝"，弹框根本不会出现，模型只收到一句"没有审批人"。
			Approval: a.Opt.Approval,
		},
		trusted: trusted,
		// 继承父级已解析的预算与策略：否则用户改过的上下文预算/审批策略
		// 在子 agent 里静默失效（applyDefaults 会覆盖成默认值）。
		ContextWindow:      a.ContextWindow,
		CompactAt:          a.CompactAt,
		MinCompactTokens:   a.MinCompactTokens,
		MaxToolResultChars: a.MaxToolResultChars,
		MaxImages:          a.MaxImages,
		ImageBudgetBytes:   a.ImageBudgetBytes,
		StreamRetryBudget:  a.StreamRetryBudget,
		ApprovalPolicy:     a.ApprovalPolicy,
		// 审批档位继承：「完全访问」是人对本回合说的"全部自动放行"，委派给下级
		// 不该悄悄收回这个授权（否则模型把 push 交给子 agent 又会撞上危险门）。
		// DangerousPolicy 刻意**不**继承：子 agent 一律留在 fail-closed（ask），
		// 父级设了 dangerous=allow 也不该把下级的拦截一起关掉。
		Mode: a.Mode,
		// 密钥脱敏必须继承：否则子 agent 读到的一份密钥会原样写进它的结果，
		// 再作为文本回到父级 transcript 与外发通道。
		Secrets: a.Secrets,
		Skills:  a.Skills,
	}
	nested.applyDefaults()
	if nested.CanSpawnSubAgent() {
		_ = reg.Register(nested.SubAgentTool()) // 闭包绑 nested：深度才能正确递增
	}
	if err := nested.Run(ctx, prompt); err != nil {
		return subAgentOutcome{sessionID: sessionID, err: fmt.Errorf("sub-agent: %w", err)}
	}
	fmt.Fprintln(a.Opt.Stderr, "[agent] sub-agent finished")

	for i := len(nested.Session.Messages) - 1; i >= 0; i-- {
		if nested.Session.Messages[i].Role == llm.RoleAssistant {
			final := strings.TrimSpace(nested.Session.Messages[i].Content)
			if final == "" {
				break
			}
			return subAgentOutcome{text: final, sessionID: sessionID}
		}
	}
	return subAgentOutcome{sessionID: sessionID, err: fmt.Errorf("sub-agent produced no final reply")}
}

// ensureLockedStderr 保证 stderr 是并发安全的：首次调用时包一层 lockedWriter 并
// 缓存，子孙 agent 直接**继承同一个实例**（不再重复包装——那样会各自持锁而失去
// 串行化，甚至同锁重入自锁）。
func (a *Agent) ensureLockedStderr() io.Writer {
	if lw, ok := a.Opt.Stderr.(*lockedWriter); ok {
		return lw
	}
	lw := &lockedWriter{w: a.Opt.Stderr}
	a.Opt.Stderr = lw
	return lw
}

// lockedWriter 串行化对底层 writer 的写入：并行子 agent 共享父级 stderr。
type lockedWriter struct {
	mu sync.Mutex
	w  io.Writer
}

func (l *lockedWriter) Write(p []byte) (int, error) {
	l.mu.Lock()
	defer l.mu.Unlock()
	return l.w.Write(p)
}