"""KB 持久化存储 — PostgreSQL 后端，替换内存 KnowledgeSourceStore。

遵循 memory/episodic.py 的 async session 模式。
ACL 条目与 KB 元数据共享事务边界（同一 PG DB）。
"""

from __future__ import annotations

import logging

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker

from agentkit.rag_platform.models import (
    Document,
    DocumentStatus,
    KBStatus,
    KnowledgeBase,
    KBAclModel,
    KBModel,
    DocumentModel,
    ORMBase,
)

logger = logging.getLogger(__name__)


def create_session_factory(database_url: str):
    """创建 async session factory（遵循 memory/models.py 模式）。"""
    engine = create_async_engine(database_url, echo=False)
    return sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)


async def ensure_tables(database_url: str) -> None:
    """创建 rag_platform 表（幂等，IF NOT EXISTS）。"""
    engine = create_async_engine(database_url, echo=False)
    async with engine.begin() as conn:
        await conn.run_sync(ORMBase.metadata.create_all)
    await engine.dispose()
    logger.info("rag_platform tables ensured")


class KBStore:
    """KB/Document 持久化存储（PostgreSQL async）。

    替换内存 KnowledgeSourceStore — KB 元数据重启不丢失。
    """

    def __init__(self, session_factory) -> None:
        self._sf = session_factory

    async def create_kb(
        self,
        name: str,
        owner: str,
        description: str = "",
        default_query_mode: str = "blend",
        default_hit_processing: str = "model_opt",
        caching_disabled: bool = False,
    ) -> KnowledgeBase:
        """创建知识库 + owner ACL 条目（同一事务）。"""
        async with self._sf() as db:
            kb = KBModel(
                name=name,
                owner=owner,
                description=description,
                default_query_mode=default_query_mode,
                default_hit_processing=default_hit_processing,
                caching_disabled=caching_disabled,
            )
            db.add(kb)
            await db.flush()  # 获取 kb.id

            acl = KBAclModel(kb_id=kb.id, user_id=owner, role="owner")
            db.add(acl)
            await db.commit()
            await db.refresh(kb)
            return KnowledgeBase.model_validate(kb)

    async def get_kb(self, kb_id: str) -> KnowledgeBase | None:
        async with self._sf() as db:
            stmt = select(KBModel).where(KBModel.id == kb_id)
            result = await db.execute(stmt)
            row = result.scalars().first()
            return KnowledgeBase.model_validate(row) if row else None

    async def list_kbs(self, owner: str | None = None) -> list[KnowledgeBase]:
        async with self._sf() as db:
            stmt = select(KBModel).where(KBModel.status == KBStatus.active.value)
            if owner:
                stmt = stmt.where(KBModel.owner == owner)
            stmt = stmt.order_by(KBModel.created_at.desc())
            result = await db.execute(stmt)
            return [KnowledgeBase.model_validate(r) for r in result.scalars().all()]

    async def delete_kb(self, kb_id: str) -> bool:
        """删除 KB — CASCADE 自动删除关联 documents 和 ACL 条目。"""
        async with self._sf() as db:
            stmt = select(KBModel).where(KBModel.id == kb_id)
            result = await db.execute(stmt)
            kb = result.scalars().first()
            if kb is None:
                return False
            await db.delete(kb)
            await db.commit()
            return True

    async def add_document(
        self,
        kb_id: str,
        filename: str,
        file_type: str,
        file_size: int,
    ) -> Document:
        async with self._sf() as db:
            doc = DocumentModel(
                kb_id=kb_id,
                filename=filename,
                file_type=file_type,
                file_size=file_size,
            )
            db.add(doc)
            await db.commit()
            await db.refresh(doc)
            return Document.model_validate(doc)

    async def get_document(self, document_id: str) -> Document | None:
        async with self._sf() as db:
            stmt = select(DocumentModel).where(DocumentModel.id == document_id)
            result = await db.execute(stmt)
            row = result.scalars().first()
            return Document.model_validate(row) if row else None

    async def list_documents(self, kb_id: str) -> list[Document]:
        async with self._sf() as db:
            stmt = (
                select(DocumentModel)
                .where(DocumentModel.kb_id == kb_id)
                .order_by(DocumentModel.created_at.desc())
            )
            result = await db.execute(stmt)
            return [Document.model_validate(r) for r in result.scalars().all()]

    async def update_document_status(
        self,
        document_id: str,
        status: DocumentStatus,
        error_message: str | None = None,
    ) -> Document | None:
        async with self._sf() as db:
            stmt = select(DocumentModel).where(DocumentModel.id == document_id)
            result = await db.execute(stmt)
            doc = result.scalars().first()
            if doc is None:
                return None
            doc.status = status.value
            doc.error_message = error_message
            await db.commit()
            await db.refresh(doc)
            return Document.model_validate(doc)

    async def delete_document(self, document_id: str) -> bool:
        async with self._sf() as db:
            stmt = select(DocumentModel).where(DocumentModel.id == document_id)
            result = await db.execute(stmt)
            doc = result.scalars().first()
            if doc is None:
                return False
            await db.delete(doc)
            await db.commit()
            return True
