"""Ingestion pipeline: parse -> chunk -> embed -> upsert, with status + retry."""

from __future__ import annotations

import uuid

from legate.db.models.workspace import Workspace
from legate.rag.embeddings.hashing import HashingEmbedder
from legate.rag.enums import DocumentStatus
from legate.rag.ingestion import ingest_document, list_document_chunks
from legate.rag.storage import build_storage
from legate.rag.vectorstore.memory import InMemoryVectorStore
from legate.repositories import documents as documents_repo
from legate.repositories import knowledge_bases as kb_repo

DOC_TEXT = (
    "Our refund policy allows customers to request a refund within 30 days.\n\n"
    "Refunds are processed to the original payment method within five business days.\n\n"
    "Digital goods are non-refundable once downloaded."
)


def _setup(db, *, filename="policy.txt", content=DOC_TEXT.encode(), store=True):
    ws = Workspace(name="W", slug=f"w-{uuid.uuid4().hex[:8]}", settings_json={})
    db.add(ws)
    db.flush()
    kb = kb_repo.create(db, workspace_id=ws.id, name="KB", embedding_provider="hashing")
    document = documents_repo.create(
        db, workspace_id=ws.id, knowledge_base_id=kb.id, source=filename
    )
    if store:
        document.storage_uri = build_storage().save(
            workspace_id=ws.id, document_id=document.id, filename=filename, data=content
        )
    db.commit()
    return ws, kb, document


def test_ingestion_produces_chunks_and_searchable_vectors(db):
    ws, kb, document = _setup(db)
    embedder = HashingEmbedder(dim=128)
    store = InMemoryVectorStore()

    result = ingest_document(
        db=db,
        document_id=document.id,
        embedder=embedder,
        vector_store=store,
        storage=build_storage(),
    )

    assert result.status is DocumentStatus.READY
    assert result.chunk_count >= 1
    chunks = list_document_chunks(db, document_id=document.id)
    assert len(chunks) == result.chunk_count
    assert all(c.vector_id is not None for c in chunks)

    qvec = embedder.embed(["how long do refunds take"])[0]
    hits = store.search(kb.vector_collection, qvec, top_k=3)
    assert hits
    assert "refund" in hits[0].payload["text"].lower()


def test_ingestion_marks_failed_on_unsupported_type(db):
    ws, kb, document = _setup(db, filename="image.bin", content=b"\x00\x01")

    result = ingest_document(
        db=db,
        document_id=document.id,
        embedder=HashingEmbedder(dim=64),
        vector_store=InMemoryVectorStore(),
        storage=build_storage(),
    )

    assert result.status is DocumentStatus.FAILED
    assert result.error and "unsupported" in result.error.lower()
    assert list_document_chunks(db, document_id=document.id) == []


def test_ingestion_marks_failed_when_no_content(db):
    ws, kb, document = _setup(db, store=False)

    result = ingest_document(
        db=db,
        document_id=document.id,
        embedder=HashingEmbedder(dim=64),
        vector_store=InMemoryVectorStore(),
        storage=build_storage(),
    )
    assert result.status is DocumentStatus.FAILED


def test_reingestion_is_idempotent(db):
    ws, kb, document = _setup(db)
    embedder = HashingEmbedder(dim=128)
    store = InMemoryVectorStore()
    storage = build_storage()

    first = ingest_document(
        db=db, document_id=document.id, embedder=embedder, vector_store=store, storage=storage
    )
    count = first.chunk_count
    second = ingest_document(
        db=db, document_id=document.id, embedder=embedder, vector_store=store, storage=storage
    )

    assert second.status is DocumentStatus.READY
    assert second.chunk_count == count
    # No duplicated chunks after re-ingestion.
    assert len(list_document_chunks(db, document_id=document.id)) == count
