"""
Chapter-based RAG - Hierarchical approach with metadata enrichment.

Groups document sections into chapters, extracts metadata (summary, questions,
key concepts), and uses a two-stage retrieval:
1. Find relevant chapters using metadata
2. Return full chapter content
"""

# pyright: reportMissingTypeStubs=false, reportMissingImports=false, reportAttributeAccessIssue=false, reportGeneralTypeIssues=false, reportUnknownParameterType=false, reportUnknownArgumentType=false, reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownLambdaType=false, reportMissingTypeArgument=false

from pathlib import Path
from typing import Any, List, Dict
from collections import defaultdict
import json
import logging

import chromadb
from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.core.schema import Document
from llama_index.vector_stores.chroma import ChromaVectorStore
from langchain_docling import DoclingLoader

from ketju.rag.base import BaseRAG
from ketju.models import ChapterMetadata
from ketju.rag.llama_index.utils.chapter_metadata import create_chapter_metadata
from ketju.core.instrumentation import add_span_event, start_span
from ketju.core.utils.files import resolve_pdf_files
from ketju.core.utils.hash import get_file_hash
from ketju.core.utils.manifest import load_pdf_manifest, save_pdf_manifest

logger = logging.getLogger(__name__)


class ChapterRag(BaseRAG):
    """
    Chapter-based hierarchical RAG.

    - Groups sections into chapters
    - Extracts metadata (summary, key concepts, questions)
    - Embeds metadata for retrieval
    - Returns full chapter content on match
    """

    def __init__(self, name: str = "chapter_rag", **kwargs: Any):
        super().__init__(name=name, **kwargs)

        self.chroma_dir = self.index_dir / "chroma"
        self.chroma_dir.mkdir(parents=True, exist_ok=True)

        self.chapters_file = self.cache_dir / "chapters.json"

        self._setup_vector_store()
        self._load_chapters()
        self._initialized = True

    def _setup_vector_store(self):
        """Initialize Chroma vector store for chapter metadata."""
        self.chroma_client = chromadb.PersistentClient(
            path=str(self.chroma_dir)
        )

        collection = self.chroma_client.get_or_create_collection(
            name="chapters"
        )

        vector_store = ChromaVectorStore(chroma_collection=collection)
        self.storage_context = StorageContext.from_defaults(
            vector_store=vector_store
        )

        try:
            self.index = VectorStoreIndex.from_vector_store(
                vector_store=vector_store,
                storage_context=self.storage_context,
            )
            logger.info("Loaded existing chapter index")
            add_span_event("chapter_rag.index.loaded")
        except Exception:
            self.index = VectorStoreIndex(
                [],
                storage_context=self.storage_context,
            )
            logger.info("Created new chapter index")
            add_span_event("chapter_rag.index.created")

    def _load_chapters(self):
        """Load cached chapter data."""
        self.chapters = {}

        if self.chapters_file.exists():
            data = json.loads(self.chapters_file.read_text())
            for chapter in data:
                self.chapters[chapter["id"]] = chapter
            logger.info("Loaded %s cached chapters", len(self.chapters))
            add_span_event("chapter_rag.chapters.loaded", attributes={"count": len(self.chapters)})

    def _save_chapters(self):
        """Save chapters to cache."""
        chapters_list = list(self.chapters.values())
        with open(self.chapters_file, "w", encoding="utf-8") as f:
            json.dump(chapters_list, f, indent=2, ensure_ascii=False)

    def _load_pdf_sections(self, pdf_path: Path) -> List[Document]:
        """Load PDF sections using Docling."""
        loader = DoclingLoader(file_path=str(pdf_path))

        docs: List[Document] = []
        for doc in loader.load():
            meta = doc.metadata.get("dl_meta", {})
            prov = meta.get("doc_items", [{}])[0].get("prov", [{}])[0]

            docs.append(
                Document(
                    text=doc.page_content,
                    metadata={
                        "source": pdf_path.name,
                        "page": prov.get("page_no"),
                        "heading": " > ".join(meta.get("headings", [])),
                    },
                )
            )

        return docs

    def _group_by_chapter(
        self,
        docs: List[Document],
    ) -> Dict[str, dict]:
        """Group sections by chapter heading."""
        groups = defaultdict(lambda: {"docs": [], "pages": set()})

        for doc in docs:
            heading = doc.metadata.get("heading") or "Unknown"
            groups[heading]["docs"].append(doc)

            if doc.metadata.get("page"):
                groups[heading]["pages"].add(doc.metadata["page"])

        return groups

    def _extract_chapters(self, pdf_path: Path) -> List[dict]:
        """Extract and enrich chapters from PDF."""
        docs = self._load_pdf_sections(pdf_path)
        grouped = self._group_by_chapter(docs)

        chapters = []

        for heading, data in grouped.items():
            chapter_text = "\n\n".join(d.text for d in data["docs"]).strip()

            # Skip very short chapters
            if len(chapter_text) < 500:
                continue

            # Generate metadata
            logger.debug("Extracting metadata: %s", heading)
            add_span_event("chapter_rag.metadata.extract", attributes={"heading": heading})
            metadata = create_chapter_metadata(chapter_text)

            chapter_id = f"{pdf_path.stem}::{heading}"

            chapter = {
                "id": chapter_id,
                "source_pdf": pdf_path.name,
                "heading": heading,
                "text": chapter_text,
                "pages": sorted(data["pages"]),
                "metadata": metadata.model_dump(),
            }

            chapters.append(chapter)

        return chapters

    def ingest_pdf(self, pdf_path: Path) -> None:
        """Ingest a PDF by extracting chapters and metadata."""
        current_hash = get_file_hash(pdf_path)
        manifest = load_pdf_manifest(self.cache_dir, pdf_path)
        if manifest and manifest.get("hash") == current_hash:
            return

        with start_span(
            "ketju.rag.chapter_rag.ingest_pdf",
            attributes={"ketju.rag.variant": getattr(self, "name", "chapter_rag"), "pdf": pdf_path.name},
        ):
            # Extract chapters
            chapters = self._extract_chapters(pdf_path)

            logger.info("Extracted %s chapters from %s", len(chapters), pdf_path.name)
            add_span_event("chapter_rag.chapters.extracted", attributes={"count": len(chapters)})

        # Store chapters and create embeddings
        docs_to_index = []

        for chapter in chapters:
            # Store full chapter data
            self.chapters[chapter["id"]] = chapter

            # Create document for indexing (using metadata string representation)
            metadata_obj = ChapterMetadata.model_validate(chapter["metadata"])

            doc = Document(
                text=str(metadata_obj),  # Use the __str__ method for embedding
                metadata={
                    "chapter_id": chapter["id"],
                    "source": chapter["source_pdf"],
                    "heading": chapter["heading"],
                },
            )
            docs_to_index.append(doc)

        # Index metadata
        self.index.insert_nodes(docs_to_index)

        # Persist
        self.storage_context.persist()
        self._save_chapters()

        # Save manifest
        save_pdf_manifest(
            self.cache_dir,
            pdf_path,
            metadata={"num_chapters": len(chapters)},
        )

        logger.info("Indexed %s chapters", len(chapters))
        add_span_event("chapter_rag.chapters.indexed", attributes={"count": len(chapters)})

    def _ingest(
        self,
        files: List[Path | str] | None = None,
        max_files: int | None = None,
    ) -> int:
        resolved_files = resolve_pdf_files(self.docs_dir, files=files, max_files=max_files)
        for pdf_path in resolved_files:
            self.ingest_pdf(pdf_path)
        return len(resolved_files)

    def _query(
        self,
        question: str,
        similarity_top_k: int = 3,
        return_full_text: bool = True,
        **kwargs: Any,
    ) -> str:
        """
        Query using chapter metadata, return full chapter content.

        Args:
            question: User's question
            similarity_top_k: Number of chapters to retrieve
            return_full_text: If True, returns full chapter text

        Returns:
            Generated answer or retrieved chapter text
        """
        with start_span(
            "ketju.rag.chapter_rag.query",
            attributes={
                "ketju.rag.variant": getattr(self, "name", "chapter_rag"),
                "ketju.rag.question_len": len(question),
                "ketju.rag.top_k": similarity_top_k,
            },
        ):
            retriever = self.index.as_retriever(similarity_top_k=similarity_top_k)

            nodes = retriever.retrieve(question)

        if not nodes:
            return "No relevant information found."

        # Get full chapter content
        results = []
        for node in nodes:
            chapter_id = node.metadata.get("chapter_id")
            if chapter_id and chapter_id in self.chapters:
                chapter = self.chapters[chapter_id]

                if return_full_text:
                    results.append(
                        f"## {chapter['heading']}\n"
                        f"Source: {chapter['source_pdf']} (pages {chapter['pages']})\n\n"
                        f"{chapter['text']}"
                    )
                else:
                    results.append(
                        f"## {chapter['heading']}\n"
                        f"Summary: {chapter['metadata']['summary']}\n"
                        f"Key Concepts: {', '.join(chapter['metadata']['key_concepts'])}"
                    )

        return "\n\n" + "="*60 + "\n\n".join(results)

    def retrieve_chapters(
        self,
        question: str,
        similarity_top_k: int = 3,
    ) -> List[dict]:
        """
        Retrieve relevant chapters.

        Returns:
            List of chapter dictionaries
        """
        retriever = self.index.as_retriever(
            similarity_top_k=similarity_top_k
        )

        nodes = retriever.retrieve(question)

        chapters = []
        for node in nodes:
            chapter_id = node.metadata.get("chapter_id")
            if chapter_id and chapter_id in self.chapters:
                chapters.append(self.chapters[chapter_id])

        return chapters

    def get_stats(self) -> dict:
        """Get statistics about this RAG instance."""
        stats = super().get_stats()
        stats["num_chapters"] = len(self.chapters)

        try:
            collection = self.chroma_client.get_collection("chapters")
            stats["num_indexed"] = collection.count()
        except Exception:
            stats["num_indexed"] = 0

        return stats


# Backward-compatible alias used in docs/examples.
ChapterRAG = ChapterRag
