import unittest
from types import SimpleNamespace

from fastapi import HTTPException
from sqlmodel import SQLModel, Session, create_engine, select

from app.access import can_access, require_kb, require_memory
from app.main import (create_kb, grant_knowledge_base, grant_memory, list_kbs,
                      list_memories, list_admin_knowledge_bases, list_admin_memories)
from app.models import AppUser, KnowledgeBase, Memory


def request_for(user: AppUser):
    return SimpleNamespace(state=SimpleNamespace(user=user))


class ResourceAccessTests(unittest.TestCase):
    def setUp(self):
        self.engine = create_engine("sqlite://")
        SQLModel.metadata.create_all(self.engine)
        self.db = Session(self.engine)
        self.admin = AppUser(username="admin", display_name="管理员", role="admin")
        self.user1 = AppUser(username="user_one", display_name="用户一", role="editor")
        self.user2 = AppUser(username="user_two", display_name="用户二", role="editor")
        self.db.add_all([self.admin, self.user1, self.user2])
        self.db.commit()
        for user in (self.admin, self.user1, self.user2):
            self.db.refresh(user)

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

    def test_scope_filters_private_global_and_shared_resources(self):
        private = create_kb({"name": "用户一私有库", "visibility": "global"},
                            self.db, request_for(self.user1))
        global_kb = create_kb({"name": "管理员全局库", "visibility": "global"},
                              self.db, request_for(self.admin))
        shared = create_kb({"name": "授权库", "visibility": "private"},
                            self.db, request_for(self.admin))
        grant_knowledge_base(shared["id"], {"user_id": self.user2.id, "permission": "use"}, self.db)

        user1_ids = {row["id"] for row in list_kbs(self.db, request_for(self.user1))}
        user2_ids = {row["id"] for row in list_kbs(self.db, request_for(self.user2))}
        self.assertIn(private["id"], user1_ids)
        self.assertNotIn(private["id"], user2_ids)
        self.assertIn(global_kb["id"], user2_ids)
        self.assertIn(shared["id"], user2_ids)
        self.assertEqual(require_kb(self.db, shared["id"], request_for(self.user2), "use").id,
                         shared["id"])
        with self.assertRaises(HTTPException):
            require_kb(self.db, shared["id"], request_for(self.user1), "use")

    def test_grant_permissions_are_hierarchical_and_admin_views_include_grants(self):
        kb = KnowledgeBase(name="授权校验", organization_id=self.user1.organization_id,
                           owner_user_id=self.user1.id, visibility="shared")
        memory = Memory(agent_id=0, content="共享记忆", organization_id=self.user1.organization_id,
                        owner_user_id=self.user1.id, visibility="shared")
        self.db.add(kb); self.db.add(memory); self.db.commit()
        self.db.refresh(kb); self.db.refresh(memory)
        grant_knowledge_base(kb.id, {"user_id": self.user2.id, "permission": "use"}, self.db)
        grant_memory(memory.id, {"user_id": self.user2.id, "permission": "use"}, self.db)
        self.assertTrue(can_access(kb, self.user2, "read", "use"))
        self.assertTrue(can_access(kb, self.user2, "use", "use"))
        self.assertFalse(can_access(kb, self.user2, "manage", "use"))
        self.assertEqual(len(list_admin_knowledge_bases(self.db)[0]["grants"]), 1)
        self.assertEqual(len(list_admin_memories(self.db)[0]["grants"]), 1)
        self.assertEqual(require_memory(self.db, memory.id, request_for(self.user2), "use").id,
                         memory.id)
        rows = list_memories(s=self.db, request=request_for(self.user2))
        self.assertEqual([row["id"] for row in rows], [memory.id])


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