from __future__ import annotations

# pyright: reportMissingTypeStubs=false, reportMissingImports=false, reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false, reportUnknownParameterType=false

from pathlib import Path
from collections.abc import Sequence

from llama_index.core.ingestion import IngestionPipeline
from llama_index.node_parser.docling import DoclingNodeParser
from llama_index.readers.docling import DoclingReader

from ketju.rag.llama_index.contracts import LlamaIndexIngestionHandler
from ketju.rag.llama_index.variants.simple import LlamaIndexRag


class SimpleDoclingIngestionHandler(LlamaIndexIngestionHandler):
    """Docling ingestion strategy that inserts parsed nodes into the current index."""

    def ingest(
        self,
        rag: LlamaIndexRag,
        *,
        files: Sequence[Path | str] | None = None,
        max_files: int | None = None,
    ) -> int:
        resolved_files = rag.resolve_pdf_files(files=files, max_files=max_files)
        ingested_nodes = 0
        for file in resolved_files:
            ingested_nodes += self._ingest_file(rag, file)
        rag.save()
        return ingested_nodes

    def _ingest_file(self, rag: LlamaIndexRag, file: Path) -> int:
        reader = DoclingReader(export_type=DoclingReader.ExportType.JSON)
        pipeline = IngestionPipeline(
            name="simple_docling_rag_ingestion",
            project_name="simple_docling_rag",
            transformations=[DoclingNodeParser(include_metadata=True, include_prev_next_rel=True)],
            documents=reader.load_data(file),
            docstore=rag.docstore,
            vector_store=rag.vector_store,
        )
        nodes = list(pipeline.run())
        if nodes:
            rag.vector_index.insert_nodes(nodes)
        return len(nodes)
