import unittest

from sqlmodel import SQLModel, Session, create_engine, select

from app.agents.runtime import RuntimeContext, RuntimeTool, run_agent, set_tool_approval
from app.llm import tools_unsupported
from app.models import Agent, AgentSession, RunEvent


class AgentRuntimeTests(unittest.TestCase):
    def setUp(self):
        self.engine = create_engine("sqlite://")
        SQLModel.metadata.create_all(self.engine)
        self.db = Session(self.engine)
        self.agent = Agent(id=101, name="测试 Agent", role="production")

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

    def test_tool_call_continues_to_final_output(self):
        calls = []

        def tool(_: RuntimeContext, args: dict) -> dict:
            return {"ok": True, "value": args["query"]}

        def model_call(_db, messages, *_args, **_kwargs):
            calls.append(messages)
            if any(item.get("role") == "tool" for item in messages):
                return {"status": "ok", "text": "最终正文", "tool_calls": []}
            return {"status": "ok", "text": "", "tool_calls": [{
                "id": "call-1", "type": "function",
                "function": {"name": "knowledge_search", "arguments": '{"query":"钩子"}'},
            }]}

        result = run_agent(
            self.db, run_id=11, campaign_id=1, segment_id=2, node_id="n1", agent=self.agent,
            provider_id=1, model="test-model", phase="production",
            messages=[{"role": "user", "content": "生成一段内容"}],
            tools=[RuntimeTool("knowledge_search", "检索", {
                "type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"],
            }, tool)],
            model_call=model_call,
            prompt_meta={"template_id": 7, "template_name": "测试模板", "template_version": 3},
        )

        self.assertEqual(result["status"], "completed")
        self.assertEqual(result["text"], "最终正文")
        self.assertEqual(len(calls), 2)
        row = self.db.exec(select(AgentSession)).one()
        self.assertEqual(row.status, "completed")
        self.assertEqual(row.turn, 2)
        self.assertEqual(row.result_meta["prompt"]["template_version"], 3)
        started = self.db.exec(select(RunEvent).where(RunEvent.event_type == "model_started")).all()
        self.assertEqual(len(started), 2)
        self.assertEqual(started[0].payload["prompt"]["template_id"], 7)

    def test_approval_waits_then_resumes_same_session(self):
        def guarded(_: RuntimeContext, __: dict) -> dict:
            return {"ok": True, "approved_result": "done"}

        def model_call(_db, messages, *_args, **_kwargs):
            if any(item.get("role") == "tool" for item in messages):
                return {"status": "ok", "text": "审批后完成", "tool_calls": []}
            return {"status": "ok", "text": "", "tool_calls": [{
                "id": "call-approve", "type": "function",
                "function": {"name": "guarded_tool", "arguments": "{}"},
            }]}

        args = dict(
            run_id=12, campaign_id=1, segment_id=3, node_id="n2", agent=self.agent,
            provider_id=1, model="test-model", phase="production",
            messages=[{"role": "user", "content": "需要审批的任务"}],
            tools=[RuntimeTool("guarded_tool", "需要审批", {"type": "object", "properties": {}},
                               guarded, requires_approval=True)],
            model_call=model_call,
        )
        first = run_agent(self.db, **args)
        self.assertEqual(first["status"], "waiting_approval")
        set_tool_approval(self.db, first["session_id"], True)
        resumed = run_agent(self.db, **args)
        self.assertEqual(resumed["status"], "completed")
        self.assertEqual(resumed["text"], "审批后完成")
        self.assertEqual(resumed["session_id"], first["session_id"])

    def test_provider_without_tools_is_detected_for_fallback(self):
        self.assertTrue(tools_unsupported(Exception("400 tools are not supported by this provider")))
        self.assertFalse(tools_unsupported(Exception("401 invalid api key")))


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