// Package context - restorer.go 定义压缩后恢复器接口和内置恢复器实现.
//
// 升华改进(ELEVATED): 恢复器是可插拔的.编程场景恢复"最近读的文件",
// 仓储场景可以注册"当前订单状态恢复器".引擎不关心恢复什么,
// 只负责按优先级调用恢复器并控制总 token 预算.
// 替代方案:硬编码 5 种恢复逻辑(原始设计,简单但无法扩展).
package context

import (
	"context"
	"fmt"
	"sort"
	"strings"
)

// PostCompactRestorer 压缩后恢复器接口.
// 每个恢复器负责一类上下文的恢复,有独立的 token 预算.
type PostCompactRestorer interface {
	Name() string
	Priority() int // 越小越先执行
	TokenBudget() int
	Restore(ctx context.Context, state *CompactState) ([]RestoreItem, error)
}

// RestoreItem 是单个恢复条目.
type RestoreItem struct {
	Type    string // "file" / "skill" / "state" / 自定义
	Name    string // 资源名(文件路径,技能名等)
	Content string // 恢复的内容
	Tokens  int    // 估算 token 数
}

// CompactState 保存压缩前的状态快照,供恢复器使用.
type CompactState struct {
	// RecentFiles 压缩前的文件缓存快照
	RecentFiles []string
	// ActiveSkills 活跃技能列表
	ActiveSkills []string
	// MCPServers MCP 服务器状态
	MCPServers []string
	// ActiveAgents 异步代理状态
	ActiveAgents []string
	// CustomData 通用:场景自定义数据
	CustomData map[string]any
}

// FileStateCacheReader 文件缓存只读接口,避免循环依赖.
type FileStateCacheReader interface {
	RecentFiles(limit int) []string
}

// SkillProvider 技能提供者接口,避免循环依赖.
type SkillProvider interface {
	ActiveSkills() []SkillInfo
}

// SkillInfo 技能基本信息.
type SkillInfo struct {
	Name    string
	Content string
}

// RestoreManager 管理所有恢复器,按优先级执行,控制总预算.
type RestoreManager struct {
	restorers   []PostCompactRestorer
	totalBudget int // 默认 80K tokens
}

// NewRestoreManager 创建恢复管理器.
// totalBudget 为 0 时使用默认值 80000.
func NewRestoreManager(totalBudget int) *RestoreManager {
	if totalBudget <= 0 {
		totalBudget = 80000
	}
	return &RestoreManager{
		totalBudget: totalBudget,
	}
}

// Register 注册一个恢复器.
func (m *RestoreManager) Register(r PostCompactRestorer) {
	m.restorers = append(m.restorers, r)
}

// Execute 按优先级执行所有恢复器,控制总 token 预算.
// 精妙之处(CLEVER): 按优先级排序后逐个执行,每个恢复器的结果立即计入已用预算.
// 如果某个恢复器的结果超过剩余预算,则截断该恢复器的结果.
// 这保证了高优先级恢复器始终能获得足够的预算.
func (m *RestoreManager) Execute(ctx context.Context, state *CompactState) ([]RestoreItem, error) {
	if len(m.restorers) == 0 {
		return nil, nil
	}

	// 按优先级排序(稳定排序保持注册顺序)
	sorted := make([]PostCompactRestorer, len(m.restorers))
	copy(sorted, m.restorers)
	sort.SliceStable(sorted, func(i, j int) bool {
		return sorted[i].Priority() < sorted[j].Priority()
	})

	var allItems []RestoreItem
	usedTokens := 0

	for _, r := range sorted {
		// 检查上下文是否已取消
		select {
		case <-ctx.Done():
			return allItems, ctx.Err()
		default:
		}

		remaining := m.totalBudget - usedTokens
		if remaining <= 0 {
			break
		}

		// 单个恢复器预算不超过其声明的上限和剩余预算
		budget := r.TokenBudget()
		if budget > remaining {
			budget = remaining
		}

		items, err := r.Restore(ctx, state)
		if err != nil {
			// 单个恢复器失败不影响其他恢复器
			continue
		}

		// 按预算截取恢复结果
		for _, item := range items {
			if usedTokens+item.Tokens > m.totalBudget {
				break
			}
			if item.Tokens > budget {
				break
			}
			allItems = append(allItems, item)
			usedTokens += item.Tokens
			budget -= item.Tokens
		}
	}

	return allItems, nil
}

// --- 5 个内置恢复器(编程场景) ---

// FileRestorer 恢复最近读过的文件.
// 最多 5 个文件,每个 5K token,总预算 50K.
type FileRestorer struct {
	fileCache FileStateCacheReader
}

// NewFileRestorer 创建文件恢复器.
func NewFileRestorer(fc FileStateCacheReader) *FileRestorer {
	return &FileRestorer{fileCache: fc}
}

func (r *FileRestorer) Name() string     { return "file-restorer" }
func (r *FileRestorer) Priority() int    { return 10 }
func (r *FileRestorer) TokenBudget() int { return 50000 }

// Restore 从 CompactState 中恢复最近读过的文件信息.
func (r *FileRestorer) Restore(_ context.Context, state *CompactState) ([]RestoreItem, error) {
	files := state.RecentFiles
	if len(files) == 0 && r.fileCache != nil {
		files = r.fileCache.RecentFiles(5)
	}

	maxFiles := 5
	if len(files) > maxFiles {
		files = files[:maxFiles]
	}

	var items []RestoreItem
	for _, f := range files {
		content := fmt.Sprintf("[Recently accessed file]: %s", f)
		tokens := estimateTextTokens(content)
		items = append(items, RestoreItem{
			Type:    "file",
			Name:    f,
			Content: content,
			Tokens:  tokens,
		})
	}
	return items, nil
}

// SkillRestorer 恢复活跃技能.
// 每个 5K token,总预算 25K.
type SkillRestorer struct {
	skillLoader SkillProvider
}

// NewSkillRestorer 创建技能恢复器.
func NewSkillRestorer(sp SkillProvider) *SkillRestorer {
	return &SkillRestorer{skillLoader: sp}
}

func (r *SkillRestorer) Name() string     { return "skill-restorer" }
func (r *SkillRestorer) Priority() int    { return 20 }
func (r *SkillRestorer) TokenBudget() int { return 25000 }

// Restore 恢复活跃技能.
func (r *SkillRestorer) Restore(_ context.Context, state *CompactState) ([]RestoreItem, error) {
	skillNames := state.ActiveSkills
	if len(skillNames) == 0 {
		return nil, nil
	}

	var items []RestoreItem
	// 如果有 SkillProvider,查询实际内容
	if r.skillLoader != nil {
		skills := r.skillLoader.ActiveSkills()
		seen := make(map[string]bool)
		for _, name := range skillNames {
			seen[name] = true
		}
		for _, s := range skills {
			if !seen[s.Name] {
				continue
			}
			content := s.Content
			if len(content) > 5000 {
				content = content[:5000] + "\n... [skill content truncated]"
			}
			tokens := estimateTextTokens(content)
			items = append(items, RestoreItem{
				Type:    "skill",
				Name:    s.Name,
				Content: fmt.Sprintf("[Active skill: %s]\n%s", s.Name, content),
				Tokens:  tokens,
			})
		}
	} else {
		// 没有 provider,只恢复技能名称
		for _, name := range skillNames {
			content := fmt.Sprintf("[Active skill]: %s", name)
			items = append(items, RestoreItem{
				Type:    "skill",
				Name:    name,
				Content: content,
				Tokens:  estimateTextTokens(content),
			})
		}
	}
	return items, nil
}

// PlanRestorer 恢复当前计划.
type PlanRestorer struct{}

func (r *PlanRestorer) Name() string     { return "plan-restorer" }
func (r *PlanRestorer) Priority() int    { return 30 }
func (r *PlanRestorer) TokenBudget() int { return 5000 }

// Restore 恢复当前计划信息.
func (r *PlanRestorer) Restore(_ context.Context, state *CompactState) ([]RestoreItem, error) {
	if state.CustomData == nil {
		return nil, nil
	}
	plan, ok := state.CustomData["current_plan"].(string)
	if !ok || plan == "" {
		return nil, nil
	}
	if len(plan) > 5000 {
		plan = plan[:5000] + "\n... [plan truncated]"
	}
	tokens := estimateTextTokens(plan)
	return []RestoreItem{{
		Type:    "state",
		Name:    "current-plan",
		Content: fmt.Sprintf("[Current plan]\n%s", plan),
		Tokens:  tokens,
	}}, nil
}

// MCPRestorer MCP 指令增量恢复.
type MCPRestorer struct{}

func (r *MCPRestorer) Name() string     { return "mcp-restorer" }
func (r *MCPRestorer) Priority() int    { return 40 }
func (r *MCPRestorer) TokenBudget() int { return 5000 }

// Restore 恢复 MCP 服务器状态.
func (r *MCPRestorer) Restore(_ context.Context, state *CompactState) ([]RestoreItem, error) {
	if len(state.MCPServers) == 0 {
		return nil, nil
	}
	content := "[Active MCP servers]: " + strings.Join(state.MCPServers, ", ")
	tokens := estimateTextTokens(content)
	return []RestoreItem{{
		Type:    "state",
		Name:    "mcp-servers",
		Content: content,
		Tokens:  tokens,
	}}, nil
}

// AgentRestorer 异步代理状态恢复.
type AgentRestorer struct{}

func (r *AgentRestorer) Name() string     { return "agent-restorer" }
func (r *AgentRestorer) Priority() int    { return 50 }
func (r *AgentRestorer) TokenBudget() int { return 5000 }

// Restore 恢复异步代理状态.
func (r *AgentRestorer) Restore(_ context.Context, state *CompactState) ([]RestoreItem, error) {
	if len(state.ActiveAgents) == 0 {
		return nil, nil
	}
	content := "[Active agents]: " + strings.Join(state.ActiveAgents, ", ")
	tokens := estimateTextTokens(content)
	return []RestoreItem{{
		Type:    "state",
		Name:    "active-agents",
		Content: content,
		Tokens:  tokens,
	}}, nil
}