"""Repositories for the Knowledge Base.

Retrieval is tenant-scoped by construction: every search joins
``document_chunks -> documents -> knowledge_sources`` and filters on
``knowledge_sources.organization_id``. There is no search entry point that takes
a chunk id directly, so a caller cannot reach across tenants.
"""

from __future__ import annotations

from sqlalchemy import func, or_, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.domain.knowledge.models import (
    Document,
    DocumentChunk,
    EmbeddingJob,
    EmbeddingJobStatus,
    KnowledgeSource,
    RetrievalPolicy,
    SourceKind,
)
from app.infrastructure.db.models import (
    DocumentChunkRow,
    DocumentRow,
    EmbeddingJobRow,
    KnowledgeSourceRow,
    RetrievalPolicyRow,
)


def _source(row: KnowledgeSourceRow) -> KnowledgeSource:
    return KnowledgeSource(
        id=row.id, organization_id=row.organization_id, project_id=row.project_id,
        name=row.name, kind=SourceKind(row.kind), uri=row.uri, created_at=row.created_at,
    )


def _document(row: DocumentRow) -> Document:
    return Document(
        id=row.id, source_id=row.source_id, title=row.title, content=row.content,
        created_at=row.created_at,
    )


class KnowledgeSourceRepository:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def get(self, source_id: str) -> KnowledgeSource | None:
        row = await self.session.get(KnowledgeSourceRow, source_id)
        return _source(row) if row else None

    async def get_for_org(self, organization_id: str, source_id: str) -> KnowledgeSource | None:
        stmt = select(KnowledgeSourceRow).where(
            KnowledgeSourceRow.id == source_id,
            KnowledgeSourceRow.organization_id == organization_id,
        )
        row = (await self.session.execute(stmt)).scalar_one_or_none()
        return _source(row) if row else None

    async def list_for_org(
        self, organization_id: str, *, limit: int = 100, offset: int = 0
    ) -> list[KnowledgeSource]:
        stmt = (
            select(KnowledgeSourceRow)
            .where(KnowledgeSourceRow.organization_id == organization_id)
            .order_by(KnowledgeSourceRow.created_at.desc())
            .limit(min(limit, 500))
            .offset(offset)
        )
        return [_source(r) for r in (await self.session.execute(stmt)).scalars()]

    async def create(self, source: KnowledgeSource) -> KnowledgeSource:
        row = KnowledgeSourceRow(
            id=source.id, organization_id=source.organization_id,
            project_id=source.project_id, name=source.name, kind=str(source.kind),
            uri=source.uri, created_at=source.created_at,
        )
        self.session.add(row)
        await self.session.flush()
        return _source(row)

    async def delete(self, source_id: str) -> None:
        row = await self.session.get(KnowledgeSourceRow, source_id)
        if row:
            await self.session.delete(row)
            await self.session.flush()


class DocumentRepository:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def list_for_source(self, source_id: str) -> list[Document]:
        stmt = (
            select(DocumentRow)
            .where(DocumentRow.source_id == source_id)
            .order_by(DocumentRow.created_at)
        )
        return [_document(r) for r in (await self.session.execute(stmt)).scalars()]

    async def create(self, document: Document) -> Document:
        row = DocumentRow(
            id=document.id, source_id=document.source_id, title=document.title,
            content=document.content, created_at=document.created_at,
        )
        self.session.add(row)
        await self.session.flush()
        return _document(row)


class DocumentChunkRepository:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def create_many(self, chunks: list[DocumentChunk]) -> int:
        for chunk in chunks:
            self.session.add(
                DocumentChunkRow(
                    id=chunk.id, document_id=chunk.document_id, index=chunk.index,
                    content=chunk.content, embedding_json=list(chunk.embedding_json),
                )
            )
        await self.session.flush()
        return len(chunks)

    async def count_for_source(self, source_id: str) -> int:
        stmt = (
            select(func.count())
            .select_from(DocumentChunkRow)
            .join(DocumentRow, DocumentRow.id == DocumentChunkRow.document_id)
            .where(DocumentRow.source_id == source_id)
        )
        return int((await self.session.execute(stmt)).scalar_one())

    async def count_for_sources(self, source_ids: list[str]) -> dict[str, int]:
        """Chunk counts for many sources in one grouped query.

        Listing sources needs a count per row; doing that one source at a time
        is an N+1. ``GROUP BY`` collapses it into a single round trip, and
        sources with no chunks are simply absent from the result.
        """
        if not source_ids:
            return {}
        stmt = (
            select(DocumentRow.source_id, func.count(DocumentChunkRow.id))
            .select_from(DocumentChunkRow)
            .join(DocumentRow, DocumentRow.id == DocumentChunkRow.document_id)
            .where(DocumentRow.source_id.in_(source_ids))
            .group_by(DocumentRow.source_id)
        )
        rows = (await self.session.execute(stmt)).all()
        return {str(source_id): int(count) for source_id, count in rows}

    async def keyword_search(
        self,
        *,
        organization_id: str,
        query: str,
        project_id: str | None = None,
        limit: int = 5,
    ) -> list[dict]:
        """MVP retrieval: case-insensitive substring match over chunk text.

        A real deployment swaps this for a vector search through
        ``VectorStorePort``; the tenant scoping stays identical.
        """
        needle = f"%{query.strip().lower()}%"
        if not query.strip():
            return []

        stmt = (
            select(
                DocumentChunkRow.id,
                DocumentChunkRow.content,
                DocumentChunkRow.index,
                DocumentRow.title,
                KnowledgeSourceRow.id,
                KnowledgeSourceRow.name,
                KnowledgeSourceRow.project_id,
            )
            .join(DocumentRow, DocumentRow.id == DocumentChunkRow.document_id)
            .join(KnowledgeSourceRow, KnowledgeSourceRow.id == DocumentRow.source_id)
            .where(
                KnowledgeSourceRow.organization_id == organization_id,
                func.lower(DocumentChunkRow.content).like(needle),
            )
            .order_by(DocumentChunkRow.document_id, DocumentChunkRow.index)
            .limit(min(limit, 50))
        )
        if project_id is not None:
            # An org-wide source (project_id is NULL) is visible to every project.
            stmt = stmt.where(
                or_(
                    KnowledgeSourceRow.project_id.is_(None),
                    KnowledgeSourceRow.project_id == project_id,
                )
            )

        rows = (await self.session.execute(stmt)).all()
        return [
            {
                "chunk_id": chunk_id,
                "content": content,
                "index": index,
                "document_title": title,
                "source_id": source_id,
                "source_name": source_name,
                "project_id": source_project_id,
            }
            for chunk_id, content, index, title, source_id, source_name, source_project_id in rows
        ]


class EmbeddingJobRepository:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def create(self, job: EmbeddingJob) -> EmbeddingJob:
        row = EmbeddingJobRow(
            id=job.id, source_id=job.source_id, status=str(job.status),
            error=job.error, created_at=job.created_at,
        )
        self.session.add(row)
        await self.session.flush()
        return EmbeddingJob(
            id=row.id, source_id=row.source_id,
            status=EmbeddingJobStatus(row.status), error=row.error,
            created_at=row.created_at,
        )

    async def latest_for_source(self, source_id: str) -> EmbeddingJob | None:
        stmt = (
            select(EmbeddingJobRow)
            .where(EmbeddingJobRow.source_id == source_id)
            .order_by(EmbeddingJobRow.created_at.desc())
        )
        row = (await self.session.execute(stmt)).scalars().first()
        if row is None:
            return None
        return EmbeddingJob(
            id=row.id, source_id=row.source_id,
            status=EmbeddingJobStatus(row.status), error=row.error,
            created_at=row.created_at,
        )

    async def set_status(
        self, job_id: str, status: EmbeddingJobStatus, *, error: str | None = None
    ) -> None:
        row = await self.session.get(EmbeddingJobRow, job_id)
        if row:
            row.status = str(status)
            if error is not None:
                row.error = error
            await self.session.flush()


class RetrievalPolicyRepository:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def list_for_org(self, organization_id: str) -> list[RetrievalPolicy]:
        stmt = select(RetrievalPolicyRow).where(
            RetrievalPolicyRow.organization_id == organization_id
        )
        return [
            RetrievalPolicy(
                id=r.id, organization_id=r.organization_id, name=r.name,
                scope_json=dict(r.scope_json or {}), created_at=r.created_at,
            )
            for r in (await self.session.execute(stmt)).scalars()
        ]

    async def create(self, policy: RetrievalPolicy) -> RetrievalPolicy:
        row = RetrievalPolicyRow(
            id=policy.id, organization_id=policy.organization_id, name=policy.name,
            scope_json=dict(policy.scope_json), created_at=policy.created_at,
        )
        self.session.add(row)
        await self.session.flush()
        return RetrievalPolicy(
            id=row.id, organization_id=row.organization_id, name=row.name,
            scope_json=dict(row.scope_json or {}), created_at=row.created_at,
        )