prompt-generator/internal/db/db.go

package db

import (
	"database/sql"
	"fmt"
	"log"

	_ "github.com/go-sql-driver/mysql"

	"prompt-generator/internal/config"
)

var DB *sql.DB

func Init(cfg *config.Config) error {
	dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8mb4&parseTime=true&loc=Local",
		cfg.DBUser, cfg.DBPassword, cfg.DBHost, cfg.DBPort, cfg.DBName)

	var err error
	DB, err = sql.Open("mysql", dsn)
	if err != nil {
		return fmt.Errorf("failed to open database: %w", err)
	}

	if err = DB.Ping(); err != nil {
		return fmt.Errorf("failed to ping database: %w", err)
	}

	// 确保 MySQL session 时区为 UTC+8
	if _, err = DB.Exec("SET time_zone = '+08:00'"); err != nil {
		log.Printf("Warning: failed to set MySQL time_zone to +08:00: %v", err)
	}

	DB.SetMaxOpenConns(25)
	DB.SetMaxIdleConns(5)

	log.Println("Database connected successfully")
	return nil
}

func AutoMigrate() error {
	tables := []string{
		`CREATE TABLE IF NOT EXISTS tags (
			id bigint NOT NULL AUTO_INCREMENT,
			name varchar(64) NOT NULL,
			description varchar(255) NOT NULL DEFAULT '',
			scope enum('system','personal') NOT NULL DEFAULT 'system',
			created_at datetime NOT NULL DEFAULT CURRENT_TIMESTAMP,
			PRIMARY KEY (id),
			KEY idx_scope (scope)
		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,

		`CREATE TABLE IF NOT EXISTS tag_options (
			id bigint NOT NULL AUTO_INCREMENT,
			tag_id bigint NOT NULL,
			label varchar(128) NOT NULL,
			constraint_text text NOT NULL,
			sort_order int NOT NULL DEFAULT 0,
			created_at datetime NOT NULL DEFAULT CURRENT_TIMESTAMP,
			PRIMARY KEY (id),
			KEY idx_tag_id (tag_id)
		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,

		`CREATE TABLE IF NOT EXISTS snippets (
			id bigint NOT NULL AUTO_INCREMENT,
			name varchar(128) NOT NULL,
			content text NOT NULL,
			category varchar(64) NOT NULL DEFAULT '',
			scope enum('system','personal') NOT NULL DEFAULT 'system',
			created_at datetime NOT NULL DEFAULT CURRENT_TIMESTAMP,
			PRIMARY KEY (id),
			KEY idx_scope (scope)
		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,

		`CREATE TABLE IF NOT EXISTS builder_sessions (
			id bigint NOT NULL AUTO_INCREMENT,
			title varchar(255) NOT NULL DEFAULT '',
			project_name varchar(255) NOT NULL DEFAULT '',
			final_prompt text,
			claude_session_id varchar(128) DEFAULT NULL,
			created_at datetime NOT NULL DEFAULT CURRENT_TIMESTAMP,
			updated_at datetime NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
			PRIMARY KEY (id),
			KEY idx_project_name (project_name),
			KEY idx_claude_session_id (claude_session_id)
		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,

		`CREATE TABLE IF NOT EXISTS builder_session_tags (
			id bigint NOT NULL AUTO_INCREMENT,
			builder_session_id bigint NOT NULL,
			tag_id bigint NOT NULL,
			tag_option_id bigint NOT NULL,
			PRIMARY KEY (id),
			KEY idx_session_id (builder_session_id)
		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,

		`CREATE TABLE IF NOT EXISTS builder_session_snippets (
			id bigint NOT NULL AUTO_INCREMENT,
			builder_session_id bigint NOT NULL,
			snippet_id bigint NOT NULL,
			sort_order int NOT NULL DEFAULT 0,
			PRIMARY KEY (id),
			KEY idx_session_id (builder_session_id)
		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,

		`CREATE TABLE IF NOT EXISTS builder_session_materials (
			id bigint NOT NULL AUTO_INCREMENT,
			builder_session_id bigint NOT NULL,
			title varchar(255) NOT NULL DEFAULT '',
			content text NOT NULL,
			sort_order int NOT NULL DEFAULT 0,
			PRIMARY KEY (id),
			KEY idx_session_id (builder_session_id)
		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,
	}

	for _, ddl := range tables {
		if _, err := DB.Exec(ddl); err != nil {
			return fmt.Errorf("failed to create table: %w", err)
		}
	}

	log.Println("Database migration completed")
	return nil
}

// LoadSettings loads persisted settings from the database into config
func LoadSettings(cfg *config.Config) {
	// Create settings table if not exists
	DB.Exec(`CREATE TABLE IF NOT EXISTS settings (
		key_name varchar(64) NOT NULL PRIMARY KEY,
		value text NOT NULL
	) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`)

	rows, err := DB.Query("SELECT key_name, value FROM settings")
	if err != nil {
		return
	}
	defer rows.Close()

	for rows.Next() {
		var key, value string
		if err := rows.Scan(&key, &value); err != nil {
			continue
		}
		switch key {
		case "llm_api_base_url":
			cfg.LLMAPIBaseURL = value
		case "llm_api_key":
			cfg.LLMAPIKey = value
		case "llm_model_name":
			cfg.LLMModelName = value
		}
	}
}

// EnsurePromptsTable creates the prompts table if it doesn't exist (for dev environments)
func EnsurePromptsTable() {
	_, err := DB.Exec(`CREATE TABLE IF NOT EXISTS prompts (
		id bigint NOT NULL AUTO_INCREMENT,
		session_id varchar(128) NOT NULL,
		project_name varchar(255) NOT NULL DEFAULT '',
		prompt text NOT NULL,
		created_at datetime NOT NULL DEFAULT CURRENT_TIMESTAMP,
		PRIMARY KEY (id),
		KEY idx_session_id (session_id),
		KEY idx_project_name (project_name),
		KEY idx_created_at (created_at)
	) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`)
	if err != nil {
		log.Printf("Warning: failed to ensure prompts table: %v", err)
	}
}

func SeedData() error {
	// Check if system tags already exist
	var count int
	err := DB.QueryRow("SELECT COUNT(*) FROM tags WHERE scope='system'").Scan(&count)
	if err != nil {
		return err
	}
	if count > 0 {
		log.Println("Seed data already exists, skipping")
		return nil
	}

	tx, err := DB.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	// Insert system tags
	type tagSeed struct {
		name    string
		desc    string
		options []struct {
			label      string
			constraint string
		}
	}

	seeds := []tagSeed{
		{
			name: "Git 规范",
			desc: "控制 Git 操作范围和提交规范",
			options: []struct{ label, constraint string }{
				{"标准 commit，禁止 push", "所有变更必须先 commit，禁止直接 push 到远程仓库。每次 commit 需要清晰的提交信息。"},
				{"标准 commit，允许 push", "变更完成后可 commit 并 push 到远程仓库。"},
				{"仅允许 commit，禁止 merge/rebase", "只允许 commit 操作，禁止 merge、rebase、force push 等危险操作。"},
			},
		},
		{
			name: "环境背景",
			desc: "提供运行环境上下文信息",
			options: []struct{ label, constraint string }{
				{"服务端环境（Linux 生产）", "当前工作环境为 Linux 生产服务器，注意路径分隔符和命令兼容性。"},
				{"本地开发环境", "当前为本地开发环境，可自由使用开发工具和调试手段。"},
				{"CI/CD 环境", "当前在 CI/CD 流水线中运行，注意环境变量和权限限制。"},
				{"测试环境", "当前为测试环境，可以进行破坏性测试。"},
			},
		},
		{
			name: "工具约束",
			desc: "限制可用的工具和操作方式",
			options: []struct{ label, constraint string }{
				{"必须使用 AskUserQuestion 提问", "遇到需要用户确认的问题时，必须使用 AskUserQuestion 工具，不要自行判断。"},
				{"优先使用专用工具而非 shell", "优先使用 Read/Write/Edit 等专用工具，避免直接使用 shell 命令处理文件。"},
				{"禁止使用 Agent 工具", "不要使用 Agent 工具派发子任务，在当前会话中完成所有工作。"},
			},
		},
		{
			name: "提交策略",
			desc: "控制 Git 提交的粒度和时机",
			options: []struct{ label, constraint string }{
				{"每完成一个小任务 commit 一次", "每完成一个独立的子任务就立即 commit，保持提交粒度细小。"},
				{"全部完成后一次性 commit", "所有任务完成后统一 commit 一个大的变更。"},
				{"按功能模块分别 commit", "按功能模块划分，每个模块完成后分别 commit。"},
			},
		},
		{
			name: "阅读策略",
			desc: "信息获取的顺序和方式",
			options: []struct{ label, constraint string }{
				{"先读 docs/ 了解架构再读代码", "先阅读项目文档和架构说明，理解整体设计后再阅读具体代码实现。"},
				{"直接读代码", "直接阅读源代码，通过代码理解项目结构。"},
				{"先读 README 和 CLAUDE.md", "先阅读 README.md 和 CLAUDE.md 了解项目规范和约束。"},
			},
		},
		{
			name: "输出格式",
			desc: "约束输出内容的格式",
			options: []struct{ label, constraint string }{
				{"Markdown", "输出内容使用 Markdown 格式，包含适当的标题、列表和代码块。"},
				{"JSON", "输出内容使用 JSON 格式，确保数据结构清晰。"},
				{"纯文本", "输出纯文本内容，不使用任何格式标记。"},
				{"代码块优先", "优先使用代码块展示内容，减少文字说明。"},
			},
		},
		{
			name: "安全约束",
			desc: "操作安全边界",
			options: []struct{ label, constraint string }{
				{"禁止访问外部网络", "不要发起任何网络请求，不访问外部 API 或下载资源。"},
				{"禁止修改系统文件", "不要修改 /etc、/usr 等系统目录下的文件。"},
				{"禁止执行危险命令", "不要执行 rm -rf、chmod 777、dd 等危险命令。"},
			},
		},
		{
			name: "代码风格",
			desc: "代码编写规范",
			options: []struct{ label, constraint string }{
				{"遵循项目现有风格", "严格遵循项目中已有的代码风格和命名约定，保持一致性。"},
				{"严格 ESLint/Prettier", "代码必须符合 ESLint 和 Prettier 规则，不得有 lint 错误。"},
				{"自由风格", "不限制代码风格，以可读性为优先。"},
			},
		},
		{
			name: "测试要求",
			desc: "测试策略和要求",
			options: []struct{ label, constraint string }{
				{"必须编写单元测试", "所有新增功能必须配套单元测试，确保测试覆盖关键逻辑。"},
				{"仅手动验证", "通过手动测试验证功能，不强制要求自动化测试。"},
				{"TDD 方式", "采用测试驱动开发，先写测试再写实现。"},
			},
		},
		{
			name: "错误处理",
			desc: "遇到异常情况的行为",
			options: []struct{ label, constraint string }{
				{"遇到不确定先问用户", "遇到不确定的情况时，使用 AskUserQuestion 向用户确认后再继续。"},
				{"尽可能自行判断", "尽量自行分析和判断，减少对用户的打扰。"},
				{"遇错停止等待指示", "遇到任何错误立即停止，等待用户进一步指示。"},
			},
		},
		{
			name: "上下文策略",
			desc: "Token 使用和输出策略",
			options: []struct{ label, constraint string }{
				{"节约上下文，精简输出", "输出尽量精简，只包含必要信息，减少 token 消耗。"},
				{"详细输出，不省略", "输出完整详细的信息，不省略任何内容。"},
				{"平衡模式", "在信息完整性和 token 效率之间取得平衡。"},
			},
		},
		{
			name: "工作模式",
			desc: "工作执行方式",
			options: []struct{ label, constraint string }{
				{"规划优先，先出方案再执行", "先制定详细方案，获得用户确认后再动手执行。"},
				{"直接执行，边做边调", "直接开始执行，根据反馈及时调整方向。"},
				{"探索模式，先研究再动手", "先深入研究和探索，充分理解后再开始执行。"},
			},
		},
	}

	for _, s := range seeds {
		res, err := tx.Exec("INSERT INTO tags (name, description, scope) VALUES (?, ?, 'system')", s.name, s.desc)
		if err != nil {
			return err
		}
		tagID, _ := res.LastInsertId()

		for i, opt := range s.options {
			_, err := tx.Exec("INSERT INTO tag_options (tag_id, label, constraint_text, sort_order) VALUES (?, ?, ?, ?)",
				tagID, opt.label, opt.constraint, i)
			if err != nil {
				return err
			}
		}
	}

	// Insert system snippets
	snippets := []struct {
		name     string
		content  string
		category string
	}{
		{
			name:     "项目架构概述",
			content:  "## 项目架构\n\n请在开始工作前了解项目整体架构：\n- 后端框架和技术栈\n- 前端框架和技术栈\n- 数据库和缓存方案\n- 部署和运维方式\n\n确保你的修改符合项目整体架构设计。",
			category: "架构",
		},
		{
			name:     "技术栈说明",
			content:  "## 技术栈\n\n请在此处填写项目使用的主要技术栈：\n- 语言：\n- 框架：\n- 数据库：\n- 缓存：\n- 消息队列：\n- 部署：",
			category: "架构",
		},
		{
			name:     "目录结构描述",
			content:  "## 目录结构\n\n项目主要目录结构：\n```\n/\n├── src/          # 源代码\n├── tests/        # 测试文件\n├── docs/         # 文档\n├── config/       # 配置文件\n└── scripts/      # 脚本工具\n```\n\n请确保新增文件放在正确的目录下。",
			category: "架构",
		},
		{
			name:     "编码规范要求",
			content:  "## 编码规范\n\n1. 变量和函数命名使用 camelCase\n2. 常量使用 UPPER_SNAKE_CASE\n3. 类名使用 PascalCase\n4. 每个函数不超过 50 行\n5. 复杂逻辑必须添加注释\n6. 错误处理不能忽略",
			category: "规范",
		},
		{
			name:     "测试框架配置",
			content:  "## 测试要求\n\n- 使用项目现有的测试框架\n- 测试文件命名：xxx_test.go / xxx.test.ts\n- 测试覆盖率要求：核心逻辑 > 80%\n- 每个测试用例独立，不依赖执行顺序\n- Mock 外部依赖，不依赖真实服务",
			category: "测试",
		},
		{
			name:     "部署流程说明",
			content:  "## 部署流程\n\n1. 代码合并到 main 分支\n2. CI 自动运行测试\n3. 构建 Docker 镜像\n4. 部署到 staging 环境验证\n5. 手动确认后部署到生产环境\n\n注意：不要修改 CI/CD 配置文件，除非明确要求。",
			category: "运维",
		},
		{
			name:     "PR 规范",
			content:  "## PR 规范\n\n- PR 标题简洁明了，说明变更内容\n- 描述中包含：变更原因、主要改动、测试方式\n- 单个 PR 不超过 500 行变更\n- 必须通过 CI 检查\n- 至少一人 Code Review 后合并",
			category: "规范",
		},
		{
			name:     "Code Review 要点",
			content:  "## Code Review 要点\n\n- 代码逻辑是否正确\n- 是否有边界条件未处理\n- 是否有安全隐患\n- 命名是否清晰\n- 是否有冗余代码\n- 测试是否充分\n- 性能是否有问题",
			category: "规范",
		},
	}

	for _, s := range snippets {
		_, err := tx.Exec("INSERT INTO snippets (name, content, category, scope) VALUES (?, ?, ?, 'system')",
			s.name, s.content, s.category)
		if err != nil {
			return err
		}
	}

	if err := tx.Commit(); err != nil {
		return err
	}

	log.Println("Seed data inserted successfully")
	return nil
}