from __future__ import annotations

# pyright: reportMissingTypeStubs=false

import hashlib
from pathlib import Path
from collections.abc import Sequence
from typing import Any, TYPE_CHECKING

from ketju.core.utils.hash import get_file_hash
from ketju.core.utils.manifest import load_pdf_manifest, save_pdf_manifest
from llama_index.core.ingestion import IngestionPipeline
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.node_parser.file.simple_file import SimpleFileNodeParser
from llama_index.core.readers.file.base import SimpleDirectoryReader
from llama_index.core.schema import NodeRelationship, RelatedNodeInfo
from llama_index.readers.file import PDFReader

from ketju.rag.llama_index.contracts import LlamaIndexIngestionHandler

if TYPE_CHECKING:
    from ketju.rag.llama_index.variants.simple import LlamaIndexRag

PIPELINE_NAME = "simple_rag_file_ingestion"
PIPELINE_PROJECT = "simple_rag"
DEFAULT_CHUNK_SIZE = 1000
DEFAULT_CHUNK_OVERLAP = 100


class SimplePdfIngestionHandler(LlamaIndexIngestionHandler):
    """Default ingestion strategy for PDF chunk-based Simple RAG."""

    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 = sum(self._ingest_pdf(rag, pdf_path) for pdf_path in resolved_files)
        rag.save()
        return ingested_nodes

    def _ingest_pdf(self, rag: LlamaIndexRag, pdf_path: Path) -> int:
        current_hash = get_file_hash(pdf_path)
        manifest = load_pdf_manifest(rag.cache_dir, pdf_path)
        if manifest and manifest.get("hash") == current_hash:
            return 0

        documents = self._load_documents(pdf_path)
        if not documents:
            return 0

        nodes = self._build_nodes(documents)
        ref_doc_id = self._build_ref_doc_id(rag, pdf_path)
        self._assign_stable_node_ids(nodes=nodes, ref_doc_id=ref_doc_id, pdf_path=pdf_path)
        self._replace_existing_ref_doc_if_needed(rag=rag, ref_doc_id=ref_doc_id, had_manifest=manifest is not None)

        if nodes:
            rag.vector_index.insert_nodes(nodes)
        save_pdf_manifest(
            rag.cache_dir,
            pdf_path,
            metadata={"ref_doc_id": ref_doc_id, "node_count": len(nodes)},
        )
        return len(nodes)

    def _load_documents(self, pdf_path: Path) -> list[Any]:
        reader = SimpleDirectoryReader(
            input_files=[str(pdf_path)],
            file_extractor={".pdf": PDFReader(return_full_document=True)},
            filename_as_id=True,
            recursive=False,
        )
        return list(reader.load_data())

    def _build_nodes(self, documents: list[Any]) -> list[Any]:
        pipeline = IngestionPipeline(
            name=PIPELINE_NAME,
            project_name=PIPELINE_PROJECT,
            transformations=[
                SimpleFileNodeParser(include_metadata=True, include_prev_next_rel=True),
                SentenceSplitter(
                    chunk_size=DEFAULT_CHUNK_SIZE,
                    chunk_overlap=DEFAULT_CHUNK_OVERLAP,
                ),
            ],
        )
        return list(pipeline.run(documents=documents))

    def _build_ref_doc_id(self, rag: LlamaIndexRag, pdf_path: Path) -> str:
        return f"{rag.name}:{pdf_path.resolve()}"

    def _assign_stable_node_ids(self, *, nodes: list[Any], ref_doc_id: str, pdf_path: Path) -> None:
        # Keep chunk IDs deterministic so the same input yields stable node identity.
        for idx, node in enumerate(nodes):
            text = getattr(node, "text", "") or ""
            digest = hashlib.sha256(f"{pdf_path.resolve()}::{text}".encode("utf-8")).hexdigest()[:16]
            node.id_ = f"{ref_doc_id}:{idx}:{digest}"
            node.relationships[NodeRelationship.SOURCE] = RelatedNodeInfo(node_id=ref_doc_id)

    def _replace_existing_ref_doc_if_needed(
        self,
        *,
        rag: LlamaIndexRag,
        ref_doc_id: str,
        had_manifest: bool,
    ) -> None:
        if not had_manifest:
            return
        try:
            rag.vector_index.delete_ref_doc(ref_doc_id, delete_from_docstore=True)
        except Exception:
            pass
