heicode-mananger/heicode/controller/agent_template_library.go

package controller

import (
	"embed"
	"strconv"
	"strings"
	"sync"

	"github.com/gin-gonic/gin"

	"github.com/heicode/manager/common"
	"github.com/heicode/manager/model"
)

// Preset agent definitions (Claude-Code style subagent .md), embedded into the
// binary and seeded into the agent_templates table on first use. Derived from
// oh-my-claudecode (MIT) — see agent_template_presets/NOTICE.md.
//
//go:embed agent_template_presets/*.md
var presetAgentFS embed.FS

// presetZh maps a preset key to its Chinese display name + description (the web
// console shows these; the underlying .md stays English for the agent runtime).
var presetZh = map[string]struct{ Name, Desc string }{
	"analyst":            {"需求分析师", "梳理并澄清需求，输出清晰的需求定义"},
	"architect":          {"架构顾问", "只读地分析代码、定位缺陷、给出架构与调试建议（不改代码）"},
	"code-reviewer":      {"代码审查员", "审查代码改动，找出 bug 与质量问题"},
	"code-simplifier":    {"代码精简师", "在不改变行为的前提下简化、去重、提升可读性"},
	"critic":             {"方案评审员", "评审计划/方案，挑出风险、漏洞与缺陷"},
	"debugger":           {"调试专家", "定位并复现缺陷的根本原因"},
	"designer":           {"界面设计师", "负责 UI/界面与交互体验设计"},
	"document-specialist": {"文档专家", "编写与整理技术文档"},
	"executor":           {"执行工程师", "按计划落地实现代码改动"},
	"explore":            {"代码探索员", "在代码库中广度搜索、快速定位相关代码"},
	"git-master":         {"Git 专家", "负责分支、提交、合并等 git 操作"},
	"planner":            {"规划师", "把目标拆解成可执行的实施计划"},
	"qa-tester":          {"QA 测试员", "手动验证功能是否按预期工作"},
	"scientist":          {"实验科学家", "提出假设并通过实验验证"},
	"security-reviewer":  {"安全审查员", "审查代码中的安全漏洞与风险"},
	"test-engineer":      {"测试工程师", "编写自动化测试"},
	"tracer":             {"链路追踪员", "追踪代码执行路径与数据流"},
	"verifier":           {"验收员", "运行并验证改动是否真正生效"},
	"writer":             {"写作员", "撰写面向用户的文案与内容"},
}

// parseAgentFrontmatter extracts simple `key: value` pairs from the leading
// `---` YAML frontmatter block of an agent .md.
func parseAgentFrontmatter(md string) map[string]string {
	out := map[string]string{}
	s := strings.TrimLeft(md, "\uFEFF \t\r\n")
	if !strings.HasPrefix(s, "---") {
		return out
	}
	rest := s[3:]
	end := strings.Index(rest, "\n---")
	if end < 0 {
		return out
	}
	for _, line := range strings.Split(rest[:end], "\n") {
		line = strings.TrimSpace(line)
		idx := strings.Index(line, ":")
		if idx <= 0 {
			continue
		}
		key := strings.TrimSpace(line[:idx])
		val := strings.TrimSpace(line[idx+1:])
		out[strings.ToLower(key)] = val
	}
	return out
}

var seedAgentTemplatesOnce sync.Once

// ensureAgentTemplatesSeeded idempotently inserts any missing preset templates.
func ensureAgentTemplatesSeeded() {
	// Guard the DB check OUTSIDE the Once so we don't consume it before the DB
	// is ready (which would permanently skip seeding).
	if model.DB == nil {
		return
	}
	seedAgentTemplatesOnce.Do(func() {
		entries, err := presetAgentFS.ReadDir("agent_template_presets")
		if err != nil {
			common.SysLog("agent template preset read: " + err.Error())
			return
		}
		order := 0
		for _, e := range entries {
			name := e.Name()
			if !strings.HasSuffix(name, ".md") || name == "NOTICE.md" {
				continue
			}
			key := strings.TrimSuffix(name, ".md")
			order++

			var count int64
			model.DB.Model(&model.AgentTemplate{}).Where("template_key = ?", key).Count(&count)
			if count > 0 {
				continue // already present (do not overwrite admin edits)
			}
			raw, err := presetAgentFS.ReadFile("agent_template_presets/" + name)
			if err != nil {
				continue
			}
			fm := parseAgentFrontmatter(string(raw))
			zh := presetZh[key]
			nameZh := zh.Name
			if nameZh == "" {
				nameZh = key
			}
			descZh := zh.Desc
			if descZh == "" {
				descZh = fm["description"]
			}
			row := model.AgentTemplate{
				TemplateKey:   key,
				NameZh:        nameZh,
				DescriptionZh: descZh,
				Model:         fm["model"],
				Definition:    string(raw),
				Source:        "preset",
				Status:        "active",
				SortOrder:     order,
			}
			if err := model.DB.Create(&row).Error; err != nil {
				common.SysLog("agent template seed " + key + ": " + err.Error())
			}
		}
	})
}

// loadAgentTemplate fetches an active template by its key.
func loadAgentTemplate(key string) (model.AgentTemplate, bool) {
	ensureAgentTemplatesSeeded()
	var row model.AgentTemplate
	if model.DB == nil || strings.TrimSpace(key) == "" {
		return row, false
	}
	if err := model.DB.Where("template_key = ?", strings.TrimSpace(key)).First(&row).Error; err != nil {
		return row, false
	}
	return row, true
}

// HeicodeListAgentTemplates: GET /api/heicode/agent-templates
// Client-facing list — Chinese name/description for display; no definition body.
func HeicodeListAgentTemplates(c *gin.Context) {
	ensureAgentTemplatesSeeded()
	var rows []model.AgentTemplate
	if model.DB != nil {
		_ = model.DB.Where("status = ?", "active").Order("sort_order asc, id asc").Find(&rows).Error
	}
	items := make([]gin.H, 0, len(rows))
	for _, row := range rows {
		items = append(items, gin.H{
			"template_id": row.TemplateKey,
			"name":        row.NameZh,
			"description": row.DescriptionZh,
			"model":       row.Model,
			"status":      row.Status,
		})
	}
	common.ApiSuccess(c, gin.H{"items": items, "total": len(items)})
}

// ── Admin maintenance (AdminAuth) ───────────────────────────────────────────

// AdminListAgentTemplates: GET /api/agent-templates — full rows incl definition.
func AdminListAgentTemplates(c *gin.Context) {
	ensureAgentTemplatesSeeded()
	var rows []model.AgentTemplate
	if model.DB != nil {
		_ = model.DB.Order("sort_order asc, id asc").Find(&rows).Error
	}
	common.ApiSuccess(c, gin.H{"items": rows, "total": len(rows)})
}

// AdminCreateAgentTemplate: POST /api/agent-templates
func AdminCreateAgentTemplate(c *gin.Context) {
	var req model.AgentTemplate
	if err := common.UnmarshalBodyReusable(c, &req); err != nil {
		agentError(c, "POLICY_REJECTED", "invalid request body")
		return
	}
	req.TemplateKey = strings.TrimSpace(req.TemplateKey)
	if req.TemplateKey == "" || strings.TrimSpace(req.Definition) == "" {
		agentError(c, "POLICY_REJECTED", "template_key and definition are required")
		return
	}
	if req.NameZh == "" {
		req.NameZh = req.TemplateKey
	}
	req.Source = "custom"
	if req.Status == "" {
		req.Status = "active"
	}
	req.Id = 0
	if model.DB == nil {
		agentError(c, "DEPLOYMENT_PERSIST_FAILED", "database not initialised")
		return
	}
	if err := model.DB.Create(&req).Error; err != nil {
		agentError(c, "DEPLOYMENT_PERSIST_FAILED", "failed to create template: "+err.Error())
		return
	}
	common.ApiSuccess(c, req)
}

// AdminUpdateAgentTemplate: PUT /api/agent-templates/:id
func AdminUpdateAgentTemplate(c *gin.Context) {
	id, _ := strconv.Atoi(strings.TrimSpace(c.Param("id")))
	if id <= 0 || model.DB == nil {
		agentError(c, "POLICY_REJECTED", "valid id required")
		return
	}
	var row model.AgentTemplate
	if err := model.DB.Where("id = ?", id).First(&row).Error; err != nil {
		agentError(c, "DEPLOYMENT_CONFLICT", "template not found")
		return
	}
	var req model.AgentTemplate
	if err := common.UnmarshalBodyReusable(c, &req); err != nil {
		agentError(c, "POLICY_REJECTED", "invalid request body")
		return
	}
	// Update editable fields only; key/source stay. Guard the critical fields
	// (name_zh, definition) so an omitted/empty value can't wipe a working
	// template; description/model/sort_order may be cleared intentionally.
	updates := map[string]any{
		"description_zh": req.DescriptionZh,
		"model":          req.Model,
		"sort_order":     req.SortOrder,
	}
	if strings.TrimSpace(req.NameZh) != "" {
		updates["name_zh"] = req.NameZh
	}
	if strings.TrimSpace(req.Definition) != "" {
		updates["definition"] = req.Definition
	}
	if strings.TrimSpace(req.Status) != "" {
		updates["status"] = req.Status
	}
	if err := model.DB.Model(&row).Updates(updates).Error; err != nil {
		agentError(c, "DEPLOYMENT_PERSIST_FAILED", "failed to update template")
		return
	}
	_ = model.DB.Where("id = ?", id).First(&row).Error
	common.ApiSuccess(c, row)
}

// AdminDeleteAgentTemplate: DELETE /api/agent-templates/:id
func AdminDeleteAgentTemplate(c *gin.Context) {
	id, _ := strconv.Atoi(strings.TrimSpace(c.Param("id")))
	if id <= 0 || model.DB == nil {
		agentError(c, "POLICY_REJECTED", "valid id required")
		return
	}
	if err := model.DB.Where("id = ?", id).Delete(&model.AgentTemplate{}).Error; err != nil {
		agentError(c, "DEPLOYMENT_PERSIST_FAILED", "failed to delete template")
		return
	}
	common.ApiSuccess(c, gin.H{"id": id, "status": "deleted"})
}