"""Document ingestion: parse -> chunk -> embed -> upsert, with status tracking.

Status transitions: pending -> processing -> ready | failed. Re-ingestion
clears the document's prior chunks and vectors first, so it is idempotent and
retryable.
"""

from __future__ import annotations

import uuid

from sqlalchemy import delete, select
from sqlalchemy.orm import Session

from legate.db.models.chunk import Chunk
from legate.db.models.document import Document
from legate.db.models.knowledge_base import KnowledgeBase
from legate.observability.logging import get_logger
from legate.rag.chunking import chunk_text
from legate.rag.embeddings.base import Embedder
from legate.rag.enums import DocumentStatus
from legate.rag.parsing import extract_text
from legate.rag.storage import LocalStorage
from legate.rag.vectorstore.base import VectorRecord, VectorStore

logger = get_logger(__name__)


def clear_document_chunks(
    db: Session, vector_store: VectorStore, kb: KnowledgeBase, document: Document
) -> None:
    """Remove a document's chunks (DB) and vectors (store)."""
    vector_store.delete_by_document(kb.vector_collection, str(document.id))
    db.execute(delete(Chunk).where(Chunk.document_id == document.id))


def ingest_document(
    *,
    db: Session,
    document_id: uuid.UUID,
    embedder: Embedder,
    vector_store: VectorStore,
    storage: LocalStorage,
) -> Document | None:
    document = db.get(Document, document_id)
    if document is None:
        logger.warning("document %s not found for ingestion", document_id)
        return None
    kb = db.get(KnowledgeBase, document.knowledge_base_id)
    if kb is None:
        document.status = DocumentStatus.FAILED
        document.error = "knowledge base not found"
        db.commit()
        return document

    document.status = DocumentStatus.PROCESSING
    document.error = None
    db.commit()

    try:
        if not document.storage_uri:
            raise ValueError("document has no stored content")
        data = storage.read(document.storage_uri)
        text = extract_text(document.source, data)
        chunks = chunk_text(text)

        clear_document_chunks(db, vector_store, kb, document)

        vectors = embedder.embed(chunks) if chunks else []
        records: list[VectorRecord] = []
        for idx, (chunk_text_value, vector) in enumerate(zip(chunks, vectors, strict=True)):
            chunk = Chunk(
                workspace_id=document.workspace_id,
                document_id=document.id,
                knowledge_base_id=kb.id,
                idx=idx,
                text=chunk_text_value,
                meta={"source": document.source},
            )
            db.add(chunk)
            db.flush()
            chunk.vector_id = str(chunk.id)
            records.append(
                VectorRecord(
                    id=str(chunk.id),
                    vector=vector,
                    payload={
                        "document_id": str(document.id),
                        "knowledge_base_id": str(kb.id),
                        "chunk_idx": idx,
                        "text": chunk_text_value,
                        "source": document.source,
                    },
                )
            )

        if records:
            vector_store.upsert(kb.vector_collection, records)

        document.chunk_count = len(records)
        document.status = DocumentStatus.READY
        document.error = None
        db.commit()
    except Exception as exc:
        logger.exception("ingestion failed for document %s", document_id)
        db.rollback()
        document = db.get(Document, document_id)
        if document is not None:
            document.status = DocumentStatus.FAILED
            document.error = str(exc)[:1000]
            db.commit()
    return document


def list_document_chunks(db: Session, *, document_id: uuid.UUID) -> list[Chunk]:
    return list(
        db.scalars(
            select(Chunk).where(Chunk.document_id == document_id).order_by(Chunk.idx.asc())
        ).all()
    )
