Agent-System/internal/squad/context.go

package squad

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
)

// ContextManager 负责上下文资产的读写、索引管理与提示词注入
type ContextManager struct{}

func NewContextManager() *ContextManager {
	return &ContextManager{}
}

// LoadContextIndex 读取工作区中的 context/index.json
func (cm *ContextManager) LoadContextIndex(wsDir string) (*ContextIndex, error) {
	idxPath := filepath.Join(wsDir, "context", "index.json")
	data, err := os.ReadFile(idxPath)
	if err != nil {
		return nil, fmt.Errorf("read context/index.json: %w", err)
	}
	var idx ContextIndex
	if err := json.Unmarshal(data, &idx); err != nil {
		return nil, fmt.Errorf("unmarshal context/index.json: %w", err)
	}
	return &idx, nil
}

// FormatPromptContext 生成注入给智能体 Prompt 的上下文导航与权威等级说明
func (cm *ContextManager) FormatPromptContext(wsDir string, role string) (string, error) {
	var b strings.Builder

	b.WriteString("## 📚 任务上下文资源与权威等级\n\n")
	b.WriteString("本工作区已为你准备好完整、分级的上下文资源，请务必先查阅对应文件：\n\n")

	// 1. 读取 index.json
	idx, err := cm.LoadContextIndex(wsDir)
	if err == nil && len(idx.Items) > 0 {
		b.WriteString("| 资源路径 | 权威等级 | 用途说明 | 使用与引用指导 |\n")
		b.WriteString("|---|---|---|---|\n")
		for _, item := range idx.Items {
			b.WriteString(fmt.Sprintf("| `%s` | **%s** | %s | %s |\n",
				item.File, item.AuthorityLevel, item.Description, item.Guidance))
		}
		b.WriteString("\n")
	}

	// 2. 注入全局项目基线摘要
	projPath := filepath.Join(wsDir, "context", "project.md")
	if projData, err := os.ReadFile(projPath); err == nil {
		b.WriteString("### 📌 项目基线要点（不可推翻）\n")
		b.WriteString(string(projData) + "\n\n")
	}

	// 3. 注入材料清单提示
	matDir := filepath.Join(wsDir, "context", "materials_text")
	if entries, err := os.ReadDir(matDir); err == nil && len(entries) > 0 {
		b.WriteString("### 📄 可用输入材料（核心原始事实证据）\n")
		for _, e := range entries {
			if e.Name() == "_all.txt" {
				continue
			}
			b.WriteString(fmt.Sprintf("- `context/materials_text/%s`\n", e.Name()))
		}
		b.WriteString("> ⚠️ 引用材料技术参数或客户原话时，必须标注出处：`〔材料: 文件名 §章节 p.页码〕`。\n\n")
	}

	// 4. 注入法规知识库命中快照提示
	kbPath := filepath.Join(wsDir, "context", "knowledge", "hits.md")
	if kbData, err := os.ReadFile(kbPath); err == nil && len(strings.TrimSpace(string(kbData))) > 0 {
		b.WriteString("### ⚖️ 法规与标准快照（参考依据）\n")
		lines := strings.Split(string(kbData), "\n")
		for _, l := range lines {
			if strings.HasPrefix(l, "### ") && strings.Contains(l, "[") {
				b.WriteString("- " + strings.TrimPrefix(l, "### ") + "\n")
			}
		}
		b.WriteString("> 详情请直接阅读 `context/knowledge/hits.md`。涉及强制标准须注明标准号与条款。\n\n")
	}

	// 5. 注入前序文档追溯提示
	priorIndexPath := filepath.Join(wsDir, "context", "prior_docs", "index.md")
	if priorData, err := os.ReadFile(priorIndexPath); err == nil && len(strings.TrimSpace(string(priorData))) > 0 {
		b.WriteString("### 🔗 前序已批准文档（前后闭环追溯）\n")
		b.WriteString(string(priorData) + "\n")
		b.WriteString("> 本阶段生成内容需与前序文档结论保持自洽。\n\n")
	}

	return b.String(), nil
}

// InjectKnowledgeHits 向现有工作区动态注入或追加法规知识库条目
func (cm *ContextManager) InjectKnowledgeHits(wsDir string, hits []KnowledgeHit) error {
	if len(hits) == 0 {
		return nil
	}
	kbDir := filepath.Join(wsDir, "context", "knowledge")
	if err := os.MkdirAll(kbDir, 0o755); err != nil {
		return err
	}
	hitsPath := filepath.Join(kbDir, "hits.md")

	var b strings.Builder
	existing, _ := os.ReadFile(hitsPath)
	if len(existing) > 0 {
		b.Write(existing)
		b.WriteString("\n\n")
	} else {
		b.WriteString("# 法规与标准检索条文摘录\n\n")
	}

	for _, hit := range hits {
		b.WriteString(fmt.Sprintf("### [%s] %s %s\n\n%s\n\n> 来源: %s\n\n---\n\n",
			strings.ToUpper(hit.Category), hit.DocCode, hit.Title, hit.Content, hit.Source))
	}

	if err := os.WriteFile(hitsPath, []byte(b.String()), 0o644); err != nil {
		return err
	}

	return cm.addContextItem(wsDir, ContextItem{
		File:           "context/knowledge/hits.md",
		Description:    "法规知识库中命中的医疗器械分类、国家标准与指导原则",
		AuthorityLevel: "reference_standard",
		Guidance:       "合规与技术论据源。引用必须带标准号；涉及审批结论需标注'需向药监局审评中心确认'",
	})
}

// InjectMaterials 向现有工作区注入上传的新材料
func (cm *ContextManager) InjectMaterials(wsDir string, materials []MaterialItem) error {
	if len(materials) == 0 {
		return nil
	}
	matDir := filepath.Join(wsDir, "context", "materials_text")
	if err := os.MkdirAll(matDir, 0o755); err != nil {
		return err
	}

	for _, m := range materials {
		name := filepath.Base(m.Filename)
		if name == "" || name == "." {
			name = "material.txt"
		}

		textContent, err := ExtractTextWithAnchors(name, m.Content)
		if err != nil || strings.TrimSpace(textContent) == "" {
			textContent = string(m.Content)
		}
		textBytes := []byte(textContent)

		p := filepath.Join(matDir, name)
		if err := os.WriteFile(p, textBytes, 0o644); err != nil {
			return err
		}
		if !strings.HasSuffix(name, ".txt") && !strings.HasSuffix(name, ".md") {
			_ = os.WriteFile(filepath.Join(matDir, name+".txt"), textBytes, 0o644)
		}

		// Also write to root materials_text/ for backward compatibility
		rootMatDir := filepath.Join(wsDir, "materials_text")
		_ = os.MkdirAll(rootMatDir, 0o755)
		_ = os.WriteFile(filepath.Join(rootMatDir, name), textBytes, 0o644)
		if !strings.HasSuffix(name, ".txt") && !strings.HasSuffix(name, ".md") {
			_ = os.WriteFile(filepath.Join(rootMatDir, name+".txt"), textBytes, 0o644)
		}

		_ = cm.addContextItem(wsDir, ContextItem{
			File:           filepath.ToSlash(filepath.Join("context", "materials_text", name)),
			Description:    fmt.Sprintf("输入材料文本: %s", name),
			AuthorityLevel: "primary_evidence",
			Guidance:       "核心原始证据。提取参数需标明出处；材料内部矛盾需在矛盾表中并列呈现",
		})
	}

	return nil
}

func (cm *ContextManager) addContextItem(wsDir string, item ContextItem) error {
	idxPath := filepath.Join(wsDir, "context", "index.json")
	var idx ContextIndex
	if data, err := os.ReadFile(idxPath); err == nil {
		_ = json.Unmarshal(data, &idx)
	}

	found := false
	for i, existing := range idx.Items {
		if existing.File == item.File {
			idx.Items[i] = item
			found = true
			break
		}
	}
	if !found {
		idx.Items = append(idx.Items, item)
	}

	data, err := json.MarshalIndent(idx, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(idxPath, data, 0o644)
}