import unittest

from sqlmodel import SQLModel, Session, create_engine

from app.agents.prompting import (
    DEFAULT_OUTPUT_SCHEMA,
    render_prompt,
    validate_schema,
)
from app.main import (create_agent, create_prompt_template, delete_prompt_template,
                      list_prompt_templates, update_prompt_template)
from app.models import Agent, AgentPromptTemplate


class PromptingTests(unittest.TestCase):
    def setUp(self):
        self.engine = create_engine("sqlite://")
        SQLModel.metadata.create_all(self.engine)
        self.db = Session(self.engine)

    def tearDown(self):
        self.db.close()

    def test_render_prompt_keeps_sections_and_replaces_known_slots(self):
        prompt = render_prompt(
            {
                "name": "测试模板",
                "role_prompt": "你是{agent_name}，专门负责{agent_role}。",
                "task_template": "标题：{title}\n标签：{content_tag}",
                "constraints": ["只输出正文"],
                "examples": [{"input": "旧输入", "output": "旧输出"}],
                "output_schema": DEFAULT_OUTPUT_SCHEMA,
                "style": "大白话",
                "version": 3,
            },
            {"agent_name": "心法 Agent", "agent_role": "production", "title": "关系沟通",
             "content_tag": "心法", "time_range": "00:00-05:00"},
            rules="只引用已提供材料",
            knowledge="<source>真实素材</source>",
            runtime_rules="Runtime 硬规则",
        )
        self.assertIn("心法 Agent", prompt)
        self.assertIn("关系沟通", prompt)
        self.assertIn("旧输入", prompt)
        self.assertIn("只引用已提供材料", prompt)
        self.assertIn("Runtime 硬规则", prompt)
        self.assertNotIn("{agent_name}", prompt)

    def test_validate_schema_requires_fields_and_types(self):
        ok, error = validate_schema(
            {"core_concepts": "沟通", "techniques": "倾听", "products_ref": "无",
             "golden_lines": "先听再说", "emotion_tone": "温和", "ending_hook": "继续练习"},
            DEFAULT_OUTPUT_SCHEMA,
        )
        self.assertTrue(ok)
        self.assertEqual(error, "")
        invalid, reason = validate_schema({"core_concepts": ""}, DEFAULT_OUTPUT_SCHEMA)
        self.assertFalse(invalid)
        self.assertIn("缺少必填字段", reason)
        invalid_type, reason = validate_schema(
            {"core_concepts": [], "techniques": "倾听", "products_ref": "无", "golden_lines": "先听再说",
             "emotion_tone": "温和", "ending_hook": "继续练习"}, DEFAULT_OUTPUT_SCHEMA)
        self.assertFalse(invalid_type)
        self.assertIn("不是字符串", reason)

    def test_prompt_template_api_versions_and_delete_protection(self):
        created = create_prompt_template({
            "name": "可测试模板", "role_prompt": "你是{agent_name}",
            "task_template": "标题：{title}", "constraints": ["只输出正文"],
            "examples": [{"input": "A", "output": "B"}],
            "output_schema": {"type": "object", "required": ["result"]},
        }, self.db)
        template = self.db.get(AgentPromptTemplate, created["id"])
        self.assertIsNotNone(template)
        self.assertEqual(template.version, 1)
        updated = update_prompt_template(template.id, {"style": "大白话"}, self.db)
        self.assertEqual(updated["version"], 2)
        listed = list_prompt_templates(self.db)
        self.assertEqual(listed[0]["agent_count"], 0)
        agent = Agent(name="绑定模板", prompt_template_id=template.id)
        self.db.add(agent); self.db.commit()
        with self.assertRaises(Exception):
            delete_prompt_template(template.id, self.db)

    def test_new_agent_is_bound_to_enabled_default_template(self):
        default = AgentPromptTemplate(name="默认绑定模板", builtin=True, enabled=True)
        self.db.add(default); self.db.commit(); self.db.refresh(default)
        created = create_agent({"name": "自动绑定 Agent", "role": "qc"}, self.db)
        agent = self.db.get(Agent, created["id"])
        self.assertIsNotNone(agent)
        self.assertEqual(agent.prompt_template_id, default.id)


if __name__ == "__main__":
    unittest.main()