package squad

// ============================================================================
// 长任务运行时治理内核 (TaskGov)
// 设计规范: .doc/archive/design/superpowers/specs/2026-09-09-longtask-governance-router-and-adhoc-design.md
// 覆盖: StepFSM 状态机 / BudgetLedger 预算账本 / Trace 执行轨迹 /
//       LoopDetector 重复识别 / StepManifest 归档锁定
// ============================================================================

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

// ----------------------------------------------------------------------------
// 1. StepFSM 步骤有限状态机
// ----------------------------------------------------------------------------

// 步骤状态常量（与既有 workflow.json 的 legacy 状态字符串完全兼容）
const (
	StepStatePending   = "pending"
	StepStateExecuting = "running"   // legacy 值保持不变
	StepStateVerifying = "verifying" // 新增：校验态
	StepStateRepairing = "repairing" // 新增：修复态
	StepStateCompleted = "completed"
	StepStateFailed    = "failed"
	StepStateAborted   = "aborted" // 新增：异常终止（用户/超预算/超时）
	StepStateSkipped   = "skipped"
)

// stepTransitions 合法状态迁移表：状态机收口一切裸字符串状态赋值
var stepTransitions = map[string]map[string]bool{
	StepStatePending:   {StepStateExecuting: true, StepStateSkipped: true, StepStateFailed: true, StepStateAborted: true},
	StepStateExecuting: {StepStateVerifying: true, StepStateCompleted: true, StepStateFailed: true, StepStateAborted: true, StepStateRepairing: true},
	StepStateVerifying: {StepStateCompleted: true, StepStateRepairing: true, StepStateFailed: true, StepStateAborted: true},
	StepStateRepairing: {StepStateVerifying: true, StepStateCompleted: true, StepStateFailed: true, StepStateAborted: true},
	StepStateFailed:    {StepStateExecuting: true, StepStateAborted: true, StepStateSkipped: true}, // 断点续跑：失败可重试
	StepStateAborted:   {StepStateExecuting: true, StepStateAborted: true, StepStateSkipped: true}, // 中止可恢复
	StepStateSkipped:   {},
	StepStateCompleted: {StepStateRepairing: true}, // 完成冻结后仅可经显式重开进入修复
}

// ValidStepTransition 校验状态迁移合法性
func ValidStepTransition(from, to string) bool {
	if from == "" {
		from = StepStatePending
	}
	allowed, ok := stepTransitions[from]
	return ok && allowed[to]
}

// TransitionStep 对工作流指定步骤执行受 FSM 守卫的状态迁移
func TransitionStep(wf *WorkflowPlan, stepID, to string) error {
	if wf == nil {
		return fmt.Errorf("workflow is nil")
	}
	for i := range wf.Steps {
		if wf.Steps[i].ID == stepID {
			from := wf.Steps[i].Status
			if !ValidStepTransition(from, to) {
				return fmt.Errorf("非法状态迁移: step %s 不允许 %s → %s", stepID, from, to)
			}
			wf.Steps[i].Status = to
			return nil
		}
	}
	return fmt.Errorf("step %s not found in workflow", stepID)
}

// ----------------------------------------------------------------------------
// 2. BudgetLedger 预算账本（多级阈值 + 软硬两级）
// ----------------------------------------------------------------------------

const (
	DefaultMaxSteps        = 12  // 单任务最大执行步骤数（观测 P95 × 2 余量）
	DefaultMaxRepairRounds = 2   // 质量修复轮上限（收编既有 maxRetries）
	DefaultMaxAgentCalls   = 16  // 底层 agent 调用总数上限（步骤 + 修复）
	SoftBudgetRatio        = 0.8 // 软线：80% 触发收敛提示
	BudgetModeApprox       = "approximate"
)

// BudgetPolicy 预算策略
type BudgetPolicy struct {
	MaxSteps        int `json:"max_steps"`
	MaxRepairRounds int `json:"max_repair_rounds"`
	MaxAgentCalls   int `json:"max_agent_calls"`
}

// DefaultBudgetPolicy 返回默认预算策略
func DefaultBudgetPolicy() BudgetPolicy {
	return BudgetPolicy{MaxSteps: DefaultMaxSteps, MaxRepairRounds: DefaultMaxRepairRounds, MaxAgentCalls: DefaultMaxAgentCalls}
}

// BudgetLedger 预算账本（持久化 ledger.json；Execute 前检查、后入账）
type BudgetLedger struct {
	Policy         BudgetPolicy `json:"policy"`
	StepExecutions int          `json:"step_executions"`
	AgentCalls     int          `json:"agent_calls"`
	RepairRounds   int          `json:"repair_rounds"`
	SoftWarned     bool         `json:"soft_warned,omitempty"`
	Aborted        bool         `json:"aborted,omitempty"`
	AbortReason    string       `json:"abort_reason,omitempty"`
	BudgetMode     string       `json:"budget_mode"`
}

// LoadOrCreateLedger 读取或初始化工作区预算账本
func LoadOrCreateLedger(wsDir string, policy BudgetPolicy) *BudgetLedger {
	l := &BudgetLedger{Policy: policy, BudgetMode: BudgetModeApprox}
	if data, err := os.ReadFile(filepath.Join(wsDir, "ledger.json")); err == nil {
		var loaded BudgetLedger
		if json.Unmarshal(data, &loaded) == nil && loaded.Policy.MaxSteps > 0 {
			l = &loaded
			// 策略以最新配置为准（允许运维收紧）
			l.Policy = policy
		}
	}
	return l
}

// Save 持久化账本
func (l *BudgetLedger) Save(wsDir string) error {
	if l == nil {
		return nil
	}
	return writeJSON(filepath.Join(wsDir, "ledger.json"), l)
}

// ReserveStep 预扣一步预算（执行前检查）
func (l *BudgetLedger) ReserveStep() error {
	if l.Aborted {
		return fmt.Errorf("预算已强制终止: %s", l.AbortReason)
	}
	if l.StepExecutions+1 > l.Policy.MaxSteps {
		l.abort(fmt.Sprintf("步骤数超限 (%d/%d)", l.StepExecutions+1, l.Policy.MaxSteps))
		return fmt.Errorf("预算硬停: 步骤数超限 (已执行 %d，上限 %d)", l.StepExecutions, l.Policy.MaxSteps)
	}
	l.StepExecutions++
	return nil
}

// RecordAgentCall 入账一次底层 agent 调用
func (l *BudgetLedger) RecordAgentCall() {
	if l == nil {
		return
	}
	l.AgentCalls++
}

// ReserveRepair 预扣一轮质量修复预算
func (l *BudgetLedger) ReserveRepair() error {
	if l == nil {
		return nil
	}
	if l.Aborted {
		return fmt.Errorf("预算已强制终止: %s", l.AbortReason)
	}
	if l.RepairRounds+1 > l.Policy.MaxRepairRounds {
		l.abort(fmt.Sprintf("修复轮数超限 (%d/%d)", l.RepairRounds+1, l.Policy.MaxRepairRounds))
		return fmt.Errorf("预算硬停: 质量修复轮数超限 (已 %d，上限 %d)", l.RepairRounds, l.Policy.MaxRepairRounds)
	}
	l.RepairRounds++
	return nil
}

func (l *BudgetLedger) abort(reason string) {
	l.Aborted = true
	l.AbortReason = reason
}

// CheckSoftLimit 检查是否触及软线（80%），返回提示语；幂等（只提示一次）
func (l *BudgetLedger) CheckSoftLimit() string {
	if l == nil || l.SoftWarned || l.Aborted {
		return ""
	}
	hit := false
	if float64(l.StepExecutions) >= SoftBudgetRatio*float64(l.Policy.MaxSteps) {
		hit = true
	}
	if float64(l.AgentCalls) >= SoftBudgetRatio*float64(l.Policy.MaxAgentCalls) {
		hit = true
	}
	if !hit {
		return ""
	}
	l.SoftWarned = true
	return fmt.Sprintf("预算软线提醒: 步骤 %d/%d、agent 调用 %d/%d，请收敛探索、直接产出",
		l.StepExecutions, l.Policy.MaxSteps, l.AgentCalls, l.Policy.MaxAgentCalls)
}

// UsageSummary 人类可读的预算用量摘要
func (l *BudgetLedger) UsageSummary() string {
	if l == nil {
		return "预算账本未启用"
	}
	state := "正常"
	if l.Aborted {
		state = "已强制终止(" + l.AbortReason + ")"
	}
	return fmt.Sprintf("步骤 %d/%d | agent调用 %d/%d | 修复轮 %d/%d | %s",
		l.StepExecutions, l.Policy.MaxSteps, l.AgentCalls, l.Policy.MaxAgentCalls, l.RepairRounds, l.Policy.MaxRepairRounds, state)
}

// ----------------------------------------------------------------------------
// 3. Trace 执行轨迹（trace.jsonl）
// ----------------------------------------------------------------------------

// TraceEntry 单次底层 agent 调用的轨迹记录
type TraceEntry struct {
	Seq                 int       `json:"seq"`
	StepID              string    `json:"step_id"`
	Role                string    `json:"role"`
	Timestamp           time.Time `json:"ts"`
	PromptDigest        string    `json:"prompt_digest"`
	OutputDigest        string    `json:"output_digest"`
	SimilarityToPrev    float64   `json:"similarity_to_prev"`
	DurationMs          int64     `json:"duration_ms"`
	Success             bool      `json:"success"`
	ExecErr             string    `json:"exec_err,omitempty"`
	ResultSuccess       bool      `json:"res_success,omitempty"`
	ResultError         string    `json:"res_error_reason,omitempty"`
	HasTrustedPatch     bool      `json:"has_trusted_patch,omitempty"`
	WroteThisInvocation bool      `json:"wrote_this_invocation,omitempty"`
	MatErr              string    `json:"mat_err,omitempty"`
	RecoveredFromDisk   bool      `json:"recovered_from_disk,omitempty"`
	CancelReason        string    `json:"cancel_reason,omitempty"`
}

// AppendTrace 追加一条执行轨迹（jsonl，天然即历史轨迹审计源）
func AppendTrace(wsDir string, e TraceEntry) error {
	if wsDir == "" {
		return nil
	}
	data, _ := json.Marshal(e)
	f, err := os.OpenFile(filepath.Join(wsDir, "trace.jsonl"), os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = f.Write(append(data, '\n'))
	return err
}

// LoadTrace 读取全部轨迹
func LoadTrace(wsDir string) []TraceEntry {
	data, err := os.ReadFile(filepath.Join(wsDir, "trace.jsonl"))
	if err != nil {
		return nil
	}
	var entries []TraceEntry
	for _, line := range strings.Split(string(data), "\n") {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		var e TraceEntry
		if json.Unmarshal([]byte(line), &e) == nil {
			entries = append(entries, e)
		}
	}
	return entries
}

func digestString(s string) string {
	sum := sha256.Sum256([]byte(s))
	return hex.EncodeToString(sum[:])[:16]
}

// OutputSimilarity 计算两段产物的相似度（rune 三元组 shingle Jaccard，无需分词）
func OutputSimilarity(a, b string) float64 {
	if a == "" && b == "" {
		return 1
	}
	if a == "" || b == "" {
		return 0
	}
	ra, rb := []rune(a), []rune(b)
	if len(ra) < 3 || len(rb) < 3 {
		if a == b {
			return 1
		}
		return 0
	}
	setA := make(map[string]struct{}, len(ra))
	for i := 0; i+3 <= len(ra); i++ {
		setA[string(ra[i:i+3])] = struct{}{}
	}
	setB := make(map[string]struct{}, len(rb))
	for i := 0; i+3 <= len(rb); i++ {
		setB[string(rb[i:i+3])] = struct{}{}
	}
	inter, union := 0, 0
	for k := range setA {
		union++
		if _, ok := setB[k]; ok {
			inter++
		}
	}
	for k := range setB {
		if _, ok := setA[k]; !ok {
			union++
		}
	}
	if union == 0 {
		return 1
	}
	return float64(inter) / float64(union)
}

// ----------------------------------------------------------------------------
// 4. LoopDetector 重复识别（原地打转 / 收敛停滞 / 振荡）
// ----------------------------------------------------------------------------

// LoopVerdict 重复检测结论
type LoopVerdict struct {
	Detected bool
	Kind     string // repeat | stall | oscillation
	Reason   string
	Hint     string // 注入下一轮 prompt 的反重复提示
}

// DetectRepetition 基于轨迹检测指定步骤的重复/停滞/振荡（纯函数，易测试）
func DetectRepetition(entries []TraceEntry, stepID string) LoopVerdict {
	var seq []TraceEntry
	for _, e := range entries {
		if e.StepID == stepID {
			seq = append(seq, e)
		}
	}
	n := len(seq)
	if n < 2 {
		return LoopVerdict{}
	}
	last, prev := seq[n-1], seq[n-2]

	// 原地打转：连续两次 prompt+output 完全一致
	if last.PromptDigest == prev.PromptDigest && last.OutputDigest == prev.OutputDigest && last.Success && prev.Success {
		return LoopVerdict{
			Detected: true,
			Kind:     "repeat",
			Reason:   fmt.Sprintf("步骤 %s 连续两次执行产出完全一致（digest=%s），判定原地打转", stepID, last.OutputDigest),
			Hint:     "⚠️ 反重复指令：你上一轮已经产出过与本次完全相同的内容且未通过验收。禁止复读原文；必须换一种实质不同的方式重写或深化，否则任务将被强制终止。",
		}
	}

	// 收敛停滞：连续两次产出相似度 ≥ 0.98（微调式空转）
	if last.SimilarityToPrev >= 0.98 && last.Success && prev.Success {
		return LoopVerdict{
			Detected: true,
			Kind:     "stall",
			Reason:   fmt.Sprintf("步骤 %s 连续两轮产物相似度 %.3f，判定收敛停滞", stepID, last.SimilarityToPrev),
			Hint:     fmt.Sprintf("⚠️ 反重复指令：最近两轮产物相似度高达 %.1f%%，属于无进度微调空转。请直接针对未通过的验收项重写对应章节，不要再小幅润色全文。", last.SimilarityToPrev*100),
		}
	}

	// 振荡：最近三次在两个 digest 间来回切换 (A→B→A)
	if n >= 3 {
		a, b, c := seq[n-3], seq[n-2], seq[n-1]
		if a.OutputDigest == c.OutputDigest && a.OutputDigest != b.OutputDigest && a.Success && c.Success {
			return LoopVerdict{
				Detected: true,
				Kind:     "oscillation",
				Reason:   fmt.Sprintf("步骤 %s 产出在两个版本间振荡（%s ↔ %s）", stepID, a.OutputDigest, b.OutputDigest),
				Hint:     "⚠️ 反重复指令：检测到你在两个版本间来回摇摆。请停下来基于验收标准直接给出收敛版本，禁止再回到先前任一版本。",
			}
		}
	}
	return LoopVerdict{}
}

// ----------------------------------------------------------------------------
// 5. StepManifest 步骤归档锁定（完成即冻结）
// ----------------------------------------------------------------------------

// StepManifest 步骤级防篡改指纹（复用 SHA-256 清单机制）
type StepManifest struct {
	StepID   string            `json:"step_id"`
	Role     string            `json:"role"`
	Title    string            `json:"title,omitempty"`
	Files    map[string]string `json:"files"` // 子工作区文件名 -> sha256
	FrozenAt time.Time         `json:"frozen_at"`
}

// FreezeStep 步骤完成后归档锁定：扫描子工作区生成 SHA-256 指纹清单
func FreezeStep(wsDir string, step StepSpec) (*StepManifest, error) {
	if step.SubWorkspace == "" {
		return nil, fmt.Errorf("step %s 缺少 sub_workspace，无法归档", step.ID)
	}
	subDir := filepath.Join(wsDir, filepath.FromSlash(step.SubWorkspace))
	entries, err := os.ReadDir(subDir)
	if err != nil {
		return nil, fmt.Errorf("read sub workspace: %w", err)
	}
	m := &StepManifest{StepID: step.ID, Role: step.Role, Title: step.Title, Files: make(map[string]string), FrozenAt: time.Now()}
	for _, e := range entries {
		if e.IsDir() || e.Name() == "step_manifest.json" {
			continue
		}
		if h, err := computeSHA256(filepath.Join(subDir, e.Name())); err == nil {
			m.Files[e.Name()] = h
		}
	}
	if len(m.Files) == 0 {
		return nil, fmt.Errorf("step %s 子工作区为空，拒绝空归档", step.ID)
	}
	if err := writeJSON(filepath.Join(subDir, "step_manifest.json"), m); err != nil {
		return nil, err
	}
	appendGovAudit(wsDir, "step_frozen", map[string]string{"step_id": step.ID, "files": fmt.Sprintf("%d", len(m.Files))})
	return m, nil
}

// ReopenStep 显式重开归档步骤（带审计）；重开后进入 repairing 态可重新执行
func ReopenStep(wsDir string, wf *WorkflowPlan, stepID, reason string) error {
	if wf == nil {
		return fmt.Errorf("workflow is nil")
	}
	for i := range wf.Steps {
		if wf.Steps[i].ID == stepID {
			if !wf.Steps[i].Frozen {
				return fmt.Errorf("step %s 未处于归档锁定状态", stepID)
			}
			if err := TransitionStep(wf, stepID, StepStateRepairing); err != nil {
				return err
			}
			wf.Steps[i].Frozen = false
			wf.Steps[i].ReopenCount++
			appendGovAudit(wsDir, "step_reopened", map[string]string{"step_id": stepID, "reason": reason})
			return SaveWorkflowJSON(wsDir, wf)
		}
	}
	return fmt.Errorf("step %s not found", stepID)
}

// SaveWorkflowJSON 持久化工作流（供治理内核复用）
func SaveWorkflowJSON(wsDir string, wf *WorkflowPlan) error {
	return writeJSON(filepath.Join(wsDir, "workflow.json"), wf)
}

// appendGovAudit 追加治理审计日志（gov_audit.jsonl）
func appendGovAudit(wsDir, action string, detail map[string]string) {
	if wsDir == "" {
		return
	}
	entry := map[string]any{"action": action, "detail": detail, "ts": time.Now()}
	data, _ := json.Marshal(entry)
	f, err := os.OpenFile(filepath.Join(wsDir, "gov_audit.jsonl"), os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644)
	if err != nil {
		return
	}
	defer f.Close()
	_, _ = f.Write(append(data, '\n'))
}

// SortedStepIDs 返回排序后的步骤 ID（测试辅助）
func SortedStepIDs(wf *WorkflowPlan) []string {
	if wf == nil {
		return nil
	}
	ids := make([]string, 0, len(wf.Steps))
	for _, s := range wf.Steps {
		ids = append(ids, s.ID)
	}
	sort.Strings(ids)
	return ids
}