import json
import os
import tempfile
import unittest

from sqlalchemy import and_, select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine

from app.models import Base, KGEntity, KGRelation
from app.services.project_knowledge_graph import DDL_GUARDRAIL_POLICY_KEY, ProjectKnowledgeGraphService


class ProjectKnowledgeGraphGuardrailSyncTest(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self) -> None:
        self._old_rules_path = os.environ.get("OPENCLAW_ROUTING_RULES_PATH")

        tmp_db = tempfile.NamedTemporaryFile(prefix="pmflow-kg-guardrail-", suffix=".db", delete=False)
        tmp_db.close()
        self.db_path = tmp_db.name
        self.engine = create_async_engine(f"sqlite+aiosqlite:///{self.db_path}", future=True)
        self.session_factory = async_sessionmaker(self.engine, expire_on_commit=False, class_=AsyncSession)
        async with self.engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)

        tmp_rules = tempfile.NamedTemporaryFile(prefix="pmflow-routing-rules-", suffix=".json", delete=False)
        tmp_rules.write(
            json.dumps(
                {
                    "rules": [
                        {
                            "name": "low-intent-rule",
                            "enabled": True,
                            "priority": 300,
                            "tier": "low",
                            "provider": "glm",
                            "model": "GLM-4.7",
                            "api_env": "GLM_API_KEY",
                        },
                        {
                            "name": "high-intent-rule",
                            "enabled": True,
                            "priority": 280,
                            "tier": "high",
                            "provider": "openai",
                            "model": "GLM-5",
                            "api_env": "OPENAI_API_KEY",
                        },
                    ]
                }
            ).encode("utf-8")
        )
        tmp_rules.flush()
        tmp_rules.close()
        self.rules_path = tmp_rules.name
        os.environ["OPENCLAW_ROUTING_RULES_PATH"] = self.rules_path

    async def asyncTearDown(self) -> None:
        await self.engine.dispose()
        for path in (getattr(self, "db_path", None), getattr(self, "rules_path", None)):
            if path:
                try:
                    os.remove(path)
                except FileNotFoundError:
                    pass

        if self._old_rules_path is None:
            os.environ.pop("OPENCLAW_ROUTING_RULES_PATH", None)
        else:
            os.environ["OPENCLAW_ROUTING_RULES_PATH"] = self._old_rules_path

    async def test_sync_model_guardrails_now_builds_dual_model_policy_graph(self) -> None:
        async with self.session_factory() as db:
            service = ProjectKnowledgeGraphService(db, tenant_id=1)
            payload = await service.sync_model_guardrails_now()

            self.assertEqual(payload["models"], 2)
            self.assertGreaterEqual(payload["relations"], 4)

            policy_entity = (
                await db.execute(
                    select(KGEntity).where(
                        and_(
                            KGEntity.tenant_id == 1,
                            KGEntity.entity_type == "governance_policy",
                            KGEntity.entity_key == DDL_GUARDRAIL_POLICY_KEY,
                        )
                    )
                )
            ).scalar_one_or_none()
            self.assertIsNotNone(policy_entity)

            model_entities = (
                await db.execute(
                    select(KGEntity).where(
                        and_(
                            KGEntity.tenant_id == 1,
                            KGEntity.entity_type == "ai_model",
                        )
                    )
                )
            ).scalars().all()
            self.assertEqual(len(model_entities), 2)
            model_names = {str(row.display_name or "") for row in model_entities}
            self.assertEqual(model_names, {"GLM-4.7", "GLM-5"})

            policy_relations = (
                await db.execute(
                    select(KGRelation).where(
                        and_(
                            KGRelation.tenant_id == 1,
                            KGRelation.relation_type == "model_must_follow_policy",
                            KGRelation.object_entity_id == int(policy_entity.id),
                        )
                    )
                )
            ).scalars().all()
            self.assertEqual(len(policy_relations), 2)