"""Document repository (tenant-scoped)."""

from __future__ import annotations

import uuid

from sqlalchemy import select
from sqlalchemy.orm import Session

from legate.db.models.document import Document
from legate.rag.enums import DocumentStatus


def create(
    db: Session,
    *,
    workspace_id: uuid.UUID,
    knowledge_base_id: uuid.UUID,
    source: str,
    content_type: str = "",
    checksum: str = "",
    size_bytes: int = 0,
    storage_uri: str | None = None,
    created_by: uuid.UUID | None = None,
) -> Document:
    document = Document(
        workspace_id=workspace_id,
        knowledge_base_id=knowledge_base_id,
        source=source,
        content_type=content_type,
        checksum=checksum,
        size_bytes=size_bytes,
        storage_uri=storage_uri,
        status=DocumentStatus.PENDING,
        created_by=created_by,
    )
    db.add(document)
    db.flush()
    return document


def get(db: Session, *, workspace_id: uuid.UUID, document_id: uuid.UUID) -> Document | None:
    return db.scalars(
        select(Document).where(Document.id == document_id, Document.workspace_id == workspace_id)
    ).first()


def list_for_kb(
    db: Session, *, workspace_id: uuid.UUID, knowledge_base_id: uuid.UUID
) -> list[Document]:
    return list(
        db.scalars(
            select(Document)
            .where(
                Document.workspace_id == workspace_id,
                Document.knowledge_base_id == knowledge_base_id,
            )
            .order_by(Document.created_at.asc())
        ).all()
    )


def delete(db: Session, document: Document) -> None:
    db.delete(document)
    db.flush()
