package squad

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

	"agent-native/runtime/llm"
)

// LLMPlanner LLM 规划器：用模型理解目标 → 拆解任务 → 选/派角色 → 生成步骤。
// 模型客户端与默认模型在构造时注入（角色级模型在派生角色时由模型自行指定）。
// 输出为严格 JSON；解析失败时重试一次，仍失败则 fail-closed 返回错误。
type LLMPlanner struct {
	Client    *llm.Client
	Model     string
	APIBase   string
	APIKey    string
	MaxTokens int
	// Fallback 在模型调用或解析失败时接手。空则保持 fail-closed。
	Fallback Planner
	// MaxSteps 计划步骤硬上限（默认 32）。
	MaxSteps int
}

// NewLLMPlanner 创建 LLM 规划器。
func NewLLMPlanner(client *llm.Client, model string) *LLMPlanner {
	return &LLMPlanner{Client: client, Model: model}
}

func (p *LLMPlanner) Name() string { return "llm" }

// llmPlanSchema 期望模型输出的 JSON 结构（宽松解析：仅取必要字段）。
type llmPlanSchema struct {
	Clarify  []llmQuestion `json:"clarify,omitempty"`
	Confirm  bool          `json:"confirm,omitempty"`
	Reason   string        `json:"confirm_reason,omitempty"`
	Agents   []llmAgent    `json:"agents"`
	Steps    []llmStep     `json:"steps"`
	Template string        `json:"template,omitempty"`
}

type llmQuestion struct {
	ID      string        `json:"id,omitempty"`
	Text    string        `json:"text"`
	Options []HumanOption `json:"options,omitempty"`
}

type llmAgent struct {
	Action           string   `json:"action"` // use | derive
	Role             string   `json:"role"`
	Name             string   `json:"name,omitempty"`
	BaseRole         string   `json:"base_role,omitempty"`
	Model            string   `json:"model,omitempty"`
	Backend          string   `json:"backend,omitempty"`
	ToolProfile      string   `json:"tool_profile,omitempty"`    // workspace | reasoning
	Tools            []string `json:"tools,omitempty"`           // 工具白名单（可选）
	Skills           []string `json:"skills,omitempty"`          // 技能标识（须已安装）
	BashPermission   string   `json:"bash_permission,omitempty"` // allow | ask | deny
	DomainTopic      string   `json:"domain_topic,omitempty"`
	Responsibilities []string `json:"responsibilities,omitempty"`
}

type llmStep struct {
	ID               string            `json:"id"`
	Role             string            `json:"role"`
	Title            string            `json:"title"`
	Parallel         bool              `json:"parallel,omitempty"`
	Targets          []string          `json:"targets,omitempty"`
	RequiresApproval bool              `json:"requires_approval,omitempty"`
	IntentID         string            `json:"intent_id,omitempty"`
	IntentSlots      map[string]string `json:"intent_slots,omitempty"`
}

const llmPlannerSystem = `你是智能体小队的规划器。给定用户目标、已注册角色清单和可选的任务模板，输出一个可执行的小队计划。
规则：
1. 只输出一个 JSON 对象，不要输出任何解释文字或代码围栏。
2. 目标信息不足时必须输出 {"clarify":[{"id":"q1","text":"问题","options":[{"value":"continue","label":"…"}]}]} 并停止，不要替用户假设。判定标准（任一命中即 clarify）：目标少于 15 个字；没有明确的对象/主题（如"帮我写个东西"）；没有可辨识的交付物；关键约束互相矛盾。
3. 正常情况输出 {"agents":[...],"steps":[...],"confirm":bool,"confirm_reason":"...","template":"markdown 骨架"}。
4. agents: 优先复用已注册角色（{"action":"use","role":"已注册key"}）；确无合适角色才派生（{"action":"derive","role":"新key","base_role":"基线key","name":"中文名","domain_topic":"领域","responsibilities":[...]}）。派生数量尽量少。派生角色可附执行配置：tool_profile（workspace|reasoning，默认 workspace）、tools（工具白名单，如 ["read_file","grep"] 或命名空间通配 ["mcp__github__*"]）、skills（技能标识，须已安装）、bash_permission（allow|ask|deny，默认 deny=禁用命令执行；ask=先问人）。仅在角色确需该能力时声明，默认不展开。
5. steps: 每步 {"id":"step_01_x","role":"角色key","title":"步骤职责"}；可并行的相邻步骤 parallel=true；高风险/破坏性/对外发布类步骤 requires_approval=true；步骤通常 2~6 个，含梳理、专业执行、汇总成文、终审把关。
6. 涉及写入、删除、对外提交或超出常规预算的计划 confirm=true。
7. 通用守则：缺失的信息若会实质改变结果、执行路径、权限或安全，必须 clarify 问人；否则自行决定并在计划中说明你决定了什么。严禁假设交付完成：无法覆盖的目标范围要如实体现在步骤或假设中，不许静默缩水。
8. 目标下方若给出"系统能力目录"：子任务命中目录能力时，该步输出 {"role":"intent","intent_id":"目录id","intent_slots":{...},"title":"..."}，槽位必须按目录填齐；必填槽位信息不足时禁止编造，改为 clarify 并使用目录的提问建议。未命中目录的子任务照常派给角色。没有目录时不得输出 intent_id 步骤。`

func (p *LLMPlanner) Plan(ctx context.Context, input PlanInput) (*PlanOutput, error) {
	if p.Client == nil || p.Model == "" {
		if p.Fallback != nil {
			return p.Fallback.Plan(ctx, input)
		}
		return nil, fmt.Errorf("llm planner requires client and model")
	}
	userMsg := p.buildPrompt(input)

	var parsed *llmPlanSchema
	var lastErr error
	for attempt := 0; attempt < 2; attempt++ {
		sys := llmPlannerSystem
		if attempt > 0 {
			sys += "\n注意：上一次输出无法解析为 JSON。这次只输出 JSON 对象本身。"
		}
		raw, err := p.complete(ctx, sys, userMsg)
		if err != nil {
			lastErr = err
			continue
		}
		schema, err := parseLLMPlan(raw)
		if err != nil {
			lastErr = err
			continue
		}
		parsed = schema
		break
	}
	if parsed == nil {
		if p.Fallback != nil {
			return p.Fallback.Plan(ctx, input)
		}
		return nil, fmt.Errorf("llm planner output unparsable after retry: %w", lastErr)
	}

	if len(parsed.Clarify) > 0 && len(parsed.Agents) == 0 {
		cl := &Clarification{Reason: "模型判定目标信息不足"}
		for _, q := range parsed.Clarify {
			cl.Questions = append(cl.Questions, Question{ID: q.ID, Text: q.Text, Options: q.Options})
		}
		return &PlanOutput{Clarification: cl}, nil
	}
	if len(parsed.Steps) == 0 {
		return nil, fmt.Errorf("llm planner produced no steps")
	}
	return p.materialize(input, parsed)
}

// materialize 把模型输出落成 RoleSpec/StepSpec：use 校验存在，derive 走注册表派生。
func (p *LLMPlanner) materialize(input PlanInput, parsed *llmPlanSchema) (*PlanOutput, error) {
	var agents []RoleSpec
	addAgent := func(spec RoleSpec) {
		for _, a := range agents {
			if a.Role == spec.Role {
				return
			}
		}
		agents = append(agents, spec)
	}
	for _, la := range parsed.Agents {
		switch strings.ToLower(la.Action) {
		case "derive":
			derived, err := input.Registry.Derive(DeriveRequest{
				RoleKey:          la.Role,
				Name:             la.Name,
				BaseRole:         la.BaseRole,
				Backend:          la.Backend,
				Model:            la.Model,
				ToolProfile:      la.ToolProfile,
				Tools:            la.Tools,
				Skills:           la.Skills,
				BashPermission:   la.BashPermission,
				DomainTopic:      la.DomainTopic,
				Responsibilities: la.Responsibilities,
			})
			if err != nil {
				return nil, fmt.Errorf("derive role %s: %w", la.Role, err)
			}
			addAgent(derived)
		default: // use
			spec, ok := input.Registry.Get(la.Role)
			if !ok {
				return nil, fmt.Errorf("plan references unregistered role %q (use derive action for new roles)", la.Role)
			}
			addAgent(spec)
		}
	}

	steps := make([]StepSpec, 0, len(parsed.Steps))
	for i, ls := range parsed.Steps {
		if ls.IntentID == "" && ls.Role == "" {
			return nil, fmt.Errorf("step %d missing role", i)
		}
		if ls.IntentID == "" {
			if _, ok := input.Registry.Get(ls.Role); !ok {
				return nil, fmt.Errorf("step %d references unregistered role %q", i, ls.Role)
			}
		} else if ls.Role == "" {
			ls.Role = "intent"
		}
		if _, inSquad := findAgent(agents, ls.Role); !inSquad {
			if spec, ok := input.Registry.Get(ls.Role); ok {
				addAgent(spec)
			}
		}
		id := ls.ID
		if id == "" {
			id = fmt.Sprintf("step_%02d", i+1)
		}
		targets := ls.Targets
		if len(targets) == 0 {
			targets = []string{"output.md"}
		}
		steps = append(steps, StepSpec{
			ID: id, Role: ls.Role, Title: ls.Title, Parallel: ls.Parallel,
			Targets: targets, RequiresApproval: ls.RequiresApproval, Status: "pending",
			IntentID: ls.IntentID, IntentSlots: ls.IntentSlots,
		})
	}
	maxSteps := p.MaxSteps
	if maxSteps <= 0 {
		maxSteps = 32
	}
	if len(steps) > maxSteps {
		return nil, fmt.Errorf("plan has %d steps, max %d", len(steps), maxSteps)
	}

	plan := &SquadPlan{
		Squad: SquadConfig{SquadName: truncate(input.Goal.Text, 40), Version: "1", Agents: agents},
		Workflow: WorkflowPlan{
			WorkflowID:   NewPlanID(),
			Title:        truncate(input.Goal.Text, 80),
			PlanningMode: PlanModeAutonomous,
			Status:       "pending",
			Steps:        steps,
		},
		Template: parsed.Template,
	}
	return &PlanOutput{Plan: plan, NeedsConfirmation: parsed.Confirm, ConfirmReason: parsed.Reason}, nil
}

func (p *LLMPlanner) buildPrompt(input PlanInput) string {
	var b strings.Builder
	b.WriteString("## 用户目标\n")
	b.WriteString(segmentGoal(input.Goal.Text))
	b.WriteString("\n")
	if len(input.Goal.Constraints) > 0 {
		b.WriteString("\n## 约束\n")
		for k, v := range input.Goal.Constraints {
			fmt.Fprintf(&b, "- %s: %s\n", k, v)
		}
	}
	if len(input.Goal.Materials) > 0 {
		b.WriteString("\n## 输入材料（文件名）\n")
		for _, m := range input.Goal.Materials {
			fmt.Fprintf(&b, "- %s\n", m.Filename)
		}
	}
	b.WriteString("\n## 已注册角色\n")
	for _, r := range input.Registry.List() {
		fmt.Fprintf(&b, "- key=%s name=%s backend=%s keywords=%v responsibilities=%v\n",
			r.Role, r.Name, r.Backend, r.Keywords, r.Responsibilities)
	}
	if input.Blueprint != nil {
		b.WriteString("\n## 命中任务模板（可作为基线微调）\n")
		enc, _ := json.Marshal(input.Blueprint)
		b.Write(enc)
		b.WriteString("\n")
	}
	if input.IntentCatalog != nil && input.IntentCatalog.Markdown != "" {
		b.WriteString("\n## 系统能力目录（可直接调用，不必派智能体）\n")
		b.WriteString(input.IntentCatalog.Markdown)
		b.WriteString("\n")
	}
	return b.String()
}

const goalSegmentRunes = 2000

// segmentGoal 把长目标按段标出，避免整段原文无界注入提示词。
func segmentGoal(text string) string {
	runes := []rune(text)
	if len(runes) <= goalSegmentRunes {
		return text + "\n"
	}
	var b strings.Builder
	for i := 0; i < len(runes); i += goalSegmentRunes {
		end := i + goalSegmentRunes
		if end > len(runes) {
			end = len(runes)
		}
		fmt.Fprintf(&b, "\n### 目标第 %d 段\n%s\n", i/goalSegmentRunes+1, string(runes[i:end]))
	}
	return b.String()
}

// complete 调一次模型并拼接完整文本。
func (p *LLMPlanner) complete(ctx context.Context, system, user string) (string, error) {
	msgs := []llm.Message{
		{Role: llm.RoleSystem, Content: system},
		{Role: llm.RoleUser, Content: user},
	}
	opts := llm.StreamOptions{MaxTokens: p.MaxTokens, APIBase: p.APIBase, APIKey: p.APIKey}
	var sb strings.Builder
	for ev := range p.Client.Stream(ctx, p.Model, msgs, nil, opts) {
		if ev.Kind == llm.KindError {
			return "", fmt.Errorf("%s", ev.ErrorMessage)
		}
		if ev.Kind == llm.KindTextDelta {
			sb.WriteString(ev.Text)
		}
		if ev.Kind == llm.KindMessageEnd {
			if ev.Message.IsError {
				return "", fmt.Errorf("assistant error: %s", ev.Message.Content)
			}
			if sb.Len() == 0 {
				sb.WriteString(ev.Message.Content)
			}
			break
		}
	}
	if ctx.Err() != nil {
		return "", ctx.Err()
	}
	return sb.String(), nil
}

// parseLLMPlan 从模型输出稳健解析 JSON（剥离 think 块/代码围栏/前后杂文）。
func parseLLMPlan(raw string) (*llmPlanSchema, error) {
	s := stripThinkBlocks(strings.TrimSpace(raw))
	s = strings.TrimPrefix(s, "```json")
	s = strings.TrimPrefix(s, "```")
	s = strings.TrimSuffix(s, "```")
	s = strings.TrimSpace(s)
	// 截取首个 { 到最后一个 }，容忍模型前后夹杂说明文字。
	if i := strings.IndexByte(s, '{'); i >= 0 {
		if j := strings.LastIndexByte(s, '}'); j > i {
			s = s[i : j+1]
		}
	}
	var out llmPlanSchema
	if err := json.Unmarshal([]byte(s), &out); err != nil {
		return nil, err
	}
	return &out, nil
}

// stripThinkBlocks 移除 MiniMax 等模型内嵌在 content 中的 <think>…</think> 块。
func stripThinkBlocks(s string) string {
	for {
		i := strings.Index(s, "<think>")
		if i < 0 {
			return s
		}
		j := strings.Index(s[i:], "</think>")
		if j < 0 {
			return s[:i] // 未闭合：丢弃思考余段
		}
		s = s[:i] + s[i+j+len("</think>"):]
	}
}