"""Core orchestrator — the main ResSum class that ties all components together."""

from __future__ import annotations

import json
import logging
import uuid
from pathlib import Path
from typing import Optional, Union

from langchain_text_splitters import RecursiveCharacterTextSplitter

from res_sum.schemas import DomainConfig, StructuredSummary, PaperMetadata
from res_sum.domains.config import load_domain
from res_sum.parsers.pdf_parser import extract_text_from_pdf, extract_sections
from res_sum.parsers.grobid_parser import extract_with_grobid, is_grobid_available
from res_sum.models.providers import get_llm
from res_sum.prompts.templates import build_summarization_prompt
from res_sum.prompts.loader import load_prompt
from res_sum.retrieval.vector_store import VectorStore
from res_sum.retrieval.hybrid import HybridRetriever
from res_sum.graph.knowledge_graph import KnowledgeGraph
from res_sum.graph.entities import (
    extract_from_full_paper,
    extract_from_chunks_with_consolidation,
    _FULL_PAPER_CHAR_LIMIT,
)
from res_sum.graph.community import (
    detect_communities,
    generate_community_summaries,
    save_communities,
    load_communities,
)
from res_sum.output.exporters import save_as_docx, save_as_json, save_as_csv

logger = logging.getLogger("res_sum")


class ResSum:
    """Research paper summarization engine with GraphRAG-enhanced retrieval.

    Usage:
        rs = ResSum(
            llm_provider="ollama",
            domain="ecology",
            data_dir="./my_project",
        )
        rs.ingest_papers("./pdfs/")
        summary = rs.summarize("Summarize key findings on pollinator decline")
    """

    def __init__(
        self,
        llm_provider: str = "ollama",
        model: Optional[str] = None,
        api_key: Optional[str] = None,
        domain: Union[str, Path, DomainConfig] = "ecology",
        data_dir: str | Path = "./knowledge_base",
        chunk_size: int = 4000,
        chunk_overlap: int = 200,
        temperature: float = 0.02,
        start_section: str = "methodology",
        prompt: Optional[Union[str, Path]] = None,
        use_grobid: bool = False,
        grobid_url: str = "http://localhost:8070",
        build_graph: bool = True,
    ) -> None:
        """Initialize ResSum.

        Args:
            llm_provider: LLM provider ("ollama", "ollama_cloud", "groq", "openai", "anthropic").
            model: Model name. None uses provider default.
            api_key: API key. Falls back to environment variables.
            domain: Domain config ("ecology", "general", path to YAML, or DomainConfig).
            data_dir: Directory for persistent storage (ChromaDB, graphs, metadata).
            chunk_size: Text chunk size in characters.
            chunk_overlap: Overlap between chunks in characters.
            temperature: LLM temperature.
            start_section: Where to start extracting text for summarization
                ("introduction" or "methodology"). This controls the vector store
                only — what text is available for retrieval and summary generation.
                The knowledge graph always processes the full paper (introduction
                through conclusion) regardless of this setting. References,
                acknowledgments, and other non-research sections are always excluded.
            prompt: Custom prompt for extraction and summarization. Can be:
                - An inline string with your prompt text
                - A path to a .txt, .py, or .json file containing the prompt
                - None to use the default domain-specific prompt
                This prompt drives what gets extracted into the knowledge graph
                and how papers are summarized. Tailor it to your research question
                for more accurate results.
            use_grobid: If True, try GROBID for PDF parsing (falls back to PyMuPDF if unavailable).
            grobid_url: URL of the GROBID server (default: http://localhost:8070).
            build_graph: If True, build knowledge graph during ingestion (enables GraphRAG retrieval).
        """
        self.llm_provider = llm_provider
        self.model = model
        self.api_key = api_key
        self.temperature = temperature
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.start_section = start_section
        self.data_dir = Path(data_dir)
        self.use_grobid = use_grobid
        self.grobid_url = grobid_url
        self.build_graph = build_graph

        # Load custom prompt (if provided)
        self._user_prompt = load_prompt(prompt)
        if self._user_prompt:
            logger.info(f"Custom prompt loaded ({len(self._user_prompt)} chars)")

        # Check GROBID availability if requested
        if use_grobid:
            if is_grobid_available(grobid_url):
                logger.info(f"GROBID server available at {grobid_url}")
            else:
                logger.warning(
                    f"GROBID server not reachable at {grobid_url}. "
                    f"Falling back to PyMuPDF parser."
                )
                self.use_grobid = False

        # Load domain config
        self.domain_config: DomainConfig = load_domain(domain)
        logger.info(f"Loaded domain: {self.domain_config.name}")

        # Initialize LLM
        self._llm = get_llm(
            provider=llm_provider,
            model=model,
            api_key=api_key,
            temperature=temperature,
        )

        # Initialize vector store
        self._vector_store = VectorStore(data_dir=self.data_dir)

        # Initialize knowledge graph
        self._knowledge_graph = KnowledgeGraph()
        graph_path = self.data_dir / "knowledge_graph.graphml"
        if graph_path.exists():
            self._knowledge_graph.load(graph_path)

        # Load communities if they exist
        communities_path = self.data_dir / "communities.json"
        self._communities: dict[str, int] = {}
        self._community_summaries: dict[int, str] = {}
        if communities_path.exists():
            self._communities, self._community_summaries = load_communities(communities_path)

        # Text splitter
        self._splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
        )

        # Track ingested papers
        self._ingested_papers: list[PaperMetadata] = []

    # ------------------------------------------------------------------
    # Ingestion
    # ------------------------------------------------------------------

    def ingest_papers(
        self,
        pdf_directory: str | Path,
        skip_existing: bool = True,
        rebuild: bool = False,
    ) -> list[PaperMetadata]:
        """Ingest all PDFs from a directory into the vector store and knowledge graph.

        Args:
            pdf_directory: Directory containing PDF files.
            skip_existing: If True, skip papers already in the vector store.
            rebuild: If True, wipe the existing knowledge base and rebuild from scratch.
                Useful when you've updated prompts, models, or domain configs.

        Returns:
            List of PaperMetadata for each ingested paper.
        """
        if rebuild:
            self._rebuild_knowledge_base()

        pdf_dir = Path(pdf_directory)
        if not pdf_dir.is_dir():
            raise NotADirectoryError(f"Not a directory: {pdf_dir}")

        pdf_files = sorted(pdf_dir.glob("*.pdf"))
        if not pdf_files:
            logger.warning(f"No PDF files found in {pdf_dir}")
            return []

        results: list[PaperMetadata] = []
        total = len(pdf_files)
        new_papers_ingested = False

        for idx, pdf_path in enumerate(pdf_files, 1):
            filename = pdf_path.name
            logger.info(f"[{idx}/{total}] Processing: {filename}")

            # Skip if already ingested
            if skip_existing and self._vector_store.has_paper(filename):
                logger.info(f"  Skipping (already ingested): {filename}")
                results.append(PaperMetadata(filename=filename, ingested=True))
                continue

            try:
                metadata = self._ingest_single_paper(pdf_path)
                results.append(metadata)
                new_papers_ingested = True
                logger.info(
                    f"  Ingested: {metadata.num_chunks} chunks, "
                    f"{metadata.num_entities} entities, "
                    f"{metadata.num_relationships} relationships"
                )
            except Exception as e:
                logger.error(f"  Error processing {filename}: {e}")
                results.append(PaperMetadata(filename=filename, ingested=False))

        # Post-ingestion: entity resolution, community detection, persistence
        if new_papers_ingested and self.build_graph and self._knowledge_graph.num_nodes() > 0:
            self._post_ingestion()

        self._ingested_papers.extend(results)
        return results

    def _ingest_single_paper(self, pdf_path: Path) -> PaperMetadata:
        """Ingest a single PDF into the vector store and knowledge graph.

        Two separate text flows:
        - Knowledge graph: uses FULL paper text (intro → conclusion) for complete
          intellectual extraction
        - Vector store: uses section-filtered text (respects start_section) for
          targeted retrieval during summarization
        """
        filename = pdf_path.name

        # Step 1: Extract full text from PDF
        full_text = extract_text_from_pdf(pdf_path)
        if not full_text.strip():
            raise ValueError(f"No text extracted from {filename}")

        # Step 2: Get section-filtered text for vector store (respects start_section)
        if self.use_grobid:
            try:
                summary_text, sections = extract_with_grobid(
                    pdf_path, grobid_url=self.grobid_url
                )
                logger.info(f"  Parsed with GROBID: {filename}")
            except Exception as e:
                logger.warning(f"  GROBID failed for {filename}: {e}. Falling back to PyMuPDF.")
                summary_text, sections = extract_sections(
                    full_text, start_section=self.start_section,
                    domain_config=self.domain_config,
                )
        else:
            summary_text, sections = extract_sections(
                full_text, start_section=self.start_section,
                domain_config=self.domain_config,
            )

        # Fall back to full text if section extraction found nothing
        if not summary_text.strip():
            logger.warning(f"  No sections found, using full text for {filename}")
            summary_text = full_text

        # Step 3: Chunk section-filtered text → vector store (for summarization)
        summary_chunks = self._splitter.split_text(summary_text)
        if not summary_chunks:
            raise ValueError(f"No chunks produced from {filename}")

        metadatas = [
            {
                "paper_filename": filename,
                "chunk_index": i,
                "section": self._guess_section(summary_chunks[i], sections),
            }
            for i in range(len(summary_chunks))
        ]
        chunk_ids = self._vector_store.add_chunks(summary_chunks, metadatas=metadatas)

        # Step 4: Extract knowledge graph from full paper (intro → conclusion)
        num_entities = 0
        num_relationships = 0
        if self.build_graph and (self.domain_config.entity_types or self._user_prompt):
            # Get clean text: intro → conclusion, excluding refs/ack
            kg_text, _ = extract_sections(
                full_text, start_section="introduction",
                domain_config=self.domain_config,
            )
            if not kg_text.strip():
                kg_text = full_text

            # Choose extraction strategy based on paper length
            if len(kg_text) <= _FULL_PAPER_CHAR_LIMIT:
                # Full-paper extraction (single LLM call — coherent graph)
                logger.info(f"  Extracting knowledge graph (full paper, {len(kg_text)} chars)")
                extraction_result = extract_from_full_paper(
                    full_text=kg_text,
                    llm=self._llm,
                    domain_config=self.domain_config,
                    paper_filename=filename,
                    custom_prompt=self._user_prompt,
                )
            else:
                # Long paper fallback: per-chunk + LLM consolidation
                full_chunks = self._splitter.split_text(kg_text)
                kg_chunk_ids = [str(uuid.uuid4()) for _ in full_chunks]
                logger.info(
                    f"  Extracting knowledge graph (long paper: {len(full_chunks)} chunks, "
                    f"will consolidate)"
                )
                extraction_result = extract_from_chunks_with_consolidation(
                    chunks=full_chunks,
                    chunk_ids=kg_chunk_ids,
                    llm=self._llm,
                    domain_config=self.domain_config,
                    custom_prompt=self._user_prompt,
                )

            # Add PAPER node at the center and connect extracted nodes to it
            from res_sum.schemas import ExtractedEntity, ExtractedRelationship
            import re as _re

            # Clean up filename into a short paper label
            # Try to extract "Author et al YEAR" pattern, fall back to truncated name
            raw_name = filename.replace(".pdf", "").replace(".PDF", "")
            raw_name = raw_name.replace("-", " ").replace("_", " ")
            raw_name = _re.sub(r"\s+", " ", raw_name).strip()

            # Look for "Author et al YEAR" or "Author YEAR" at the start
            author_year = _re.match(r"^(.+?(?:et al\.?)?\s+\d{4})", raw_name)
            if author_year:
                clean_name = author_year.group(1).strip()
            elif len(raw_name) > 40:
                clean_name = raw_name[:40].rsplit(" ", 1)[0] + "..."
            else:
                clean_name = raw_name

            paper_node = ExtractedEntity(
                name=clean_name,
                node_type="PAPER",
                entity_type="PAPER",
                description=f"Source paper: {filename}",
            )
            extraction_result.entities.insert(0, paper_node)

            # Connect research problems, methods, study systems to the paper node
            paper_name = paper_node.name
            for entity in extraction_result.entities:
                if entity.node_type in ("RESEARCH_PROBLEM", "METHOD", "STUDY_SYSTEM", "LIMITATION"):
                    edge_type = {
                        "RESEARCH_PROBLEM": "ADDRESSES",
                        "METHOD": "USES_METHOD",
                        "STUDY_SYSTEM": "STUDIED_IN",
                        "LIMITATION": "HAS_LIMITATION",
                    }.get(entity.node_type, "ADDRESSES")
                    extraction_result.relationships.append(
                        ExtractedRelationship(
                            source=paper_name,
                            target=entity.name,
                            relation_type=edge_type,
                        )
                    )

            self._knowledge_graph.add_extraction_result(extraction_result, paper_filename=filename)
            num_entities = len(extraction_result.entities)
            num_relationships = len(extraction_result.relationships)

        found_sections = [s for s, text in sections.items() if text.strip()]

        return PaperMetadata(
            filename=filename,
            filepath=str(pdf_path),
            num_chunks=len(summary_chunks),
            num_entities=num_entities,
            num_relationships=num_relationships,
            sections_found=found_sections,
            ingested=True,
        )

    def _rebuild_knowledge_base(self) -> None:
        """Wipe existing knowledge base and start fresh."""
        logger.info("Rebuilding knowledge base from scratch...")

        # Clear ChromaDB — use the API, don't delete files (avoids locked DB issues)
        try:
            self._vector_store._client.delete_collection("paper_chunks")
            self._vector_store._collection = (
                self._vector_store._client.get_or_create_collection(name="paper_chunks")
            )
        except Exception as e:
            logger.warning(f"Could not reset ChromaDB collection: {e}. Recreating from scratch.")
            import shutil
            # Close client, delete files, recreate
            del self._vector_store
            chroma_path = self.data_dir / "chroma_db"
            if chroma_path.exists():
                shutil.rmtree(chroma_path)
            self._vector_store = VectorStore(data_dir=self.data_dir)

        # Clear knowledge graph
        graph_path = self.data_dir / "knowledge_graph.graphml"
        if graph_path.exists():
            graph_path.unlink()
        self._knowledge_graph = KnowledgeGraph()

        # Clear communities
        comm_path = self.data_dir / "communities.json"
        if comm_path.exists():
            comm_path.unlink()
        self._communities = {}
        self._community_summaries = {}

        self._ingested_papers = []
        logger.info("Knowledge base cleared. Ready for fresh ingestion.")

    def _post_ingestion(self) -> None:
        """Run after ingesting new papers: entity resolution, community detection, persist."""
        logger.info("Post-ingestion: resolving entities...")
        merge_map = self._knowledge_graph.build_alias_merge_map()
        if merge_map:
            logger.info(f"  Merging {len(merge_map)} duplicate entities")
            self._knowledge_graph.resolve_entities(merge_map)

        logger.info("Post-ingestion: detecting communities...")
        self._communities = detect_communities(self._knowledge_graph)

        if self._communities:
            logger.info("Post-ingestion: generating community summaries...")
            self._community_summaries = generate_community_summaries(
                self._knowledge_graph, self._communities, self._llm
            )

        # Persist everything
        self._knowledge_graph.save(self.data_dir / "knowledge_graph.graphml")
        if self._communities:
            save_communities(
                self._communities,
                self._community_summaries,
                self.data_dir / "communities.json",
            )

        logger.info(
            f"Knowledge graph: {self._knowledge_graph.num_nodes()} nodes, "
            f"{self._knowledge_graph.num_edges()} edges, "
            f"{len(set(self._communities.values()))} communities"
        )

    def _guess_section(self, chunk_text: str, sections: dict[str, str]) -> str:
        """Guess which section a chunk belongs to based on text overlap."""
        for section_name, section_text in sections.items():
            if section_text and chunk_text[:100] in section_text:
                return section_name
        return "unknown"

    # ------------------------------------------------------------------
    # Summarization
    # ------------------------------------------------------------------

    def summarize(
        self,
        query: Optional[str] = None,
        prompt: Optional[str] = None,
        use_cot: bool = True,
        n_chunks: int = 10,
        paper_filter: Optional[str] = None,
        mode: str = "hybrid",
    ) -> str:
        """Summarize ingested papers based on a query.

        Args:
            query: Query to guide retrieval and summarization.
                If None, retrieves the most representative chunks.
            prompt: Custom system prompt (overrides domain CoT).
            use_cot: Whether to use Chain-of-Thought prompting.
            n_chunks: Number of chunks to retrieve for context.
            paper_filter: If set, only retrieve chunks from this paper.
            mode: Retrieval mode — "local", "graph", "global", or "hybrid".

        Returns:
            Summary text.
        """
        if self._vector_store.count() == 0:
            raise RuntimeError("No papers ingested. Call ingest_papers() first.")

        # Build hybrid retriever
        retriever = HybridRetriever(
            vector_store=self._vector_store,
            knowledge_graph=self._knowledge_graph if self._knowledge_graph.num_nodes() > 0 else None,
            communities=self._communities,
            community_summaries=self._community_summaries,
        )

        # Retrieve relevant chunks
        search_query = query or "Summarize the key findings, methods, and implications."
        retrieved = retriever.retrieve(
            query=search_query,
            mode=mode,
            n_results=n_chunks,
            paper_filter=paper_filter,
        )

        if not retrieved:
            raise RuntimeError("No relevant chunks found.")

        context_text = "\n\n---\n\n".join(chunk["text"] for chunk in retrieved)

        # Build prompt — method-level prompt overrides instance-level, which overrides domain default
        effective_prompt = prompt or self._user_prompt
        prompt_template = build_summarization_prompt(
            user_prompt=effective_prompt,
            domain_config=self.domain_config,
            use_cot=use_cot,
        )

        # Invoke LLM
        chain = prompt_template | self._llm
        response = chain.invoke({"text": context_text})

        return response.content

    # ------------------------------------------------------------------
    # Batch summarization (per-paper)
    # ------------------------------------------------------------------

    def summarize_papers(
        self,
        pdf_directory: str | Path,
        output_directory: str | Path,
        prompt: Optional[str] = None,
        use_cot: bool = True,
        output_format: str = "docx",
        n_chunks: int = 10,
        skip_existing: bool = True,
        mode: str = "hybrid",
    ) -> list[dict]:
        """Ingest and summarize all PDFs in a directory, saving outputs.

        Args:
            pdf_directory: Directory containing PDF files.
            output_directory: Directory for output files.
            prompt: Custom prompt (overrides domain CoT).
            use_cot: Whether to use Chain-of-Thought prompting.
            output_format: Output format ("docx", "json", "csv").
            n_chunks: Number of chunks per paper for summarization.
            skip_existing: Skip papers that already have output files.
            mode: Retrieval mode — "local", "graph", "global", or "hybrid".

        Returns:
            List of dicts with paper filename and status.
        """
        pdf_dir = Path(pdf_directory)
        out_dir = Path(output_directory)
        out_dir.mkdir(parents=True, exist_ok=True)

        pdf_files = sorted(pdf_dir.glob("*.pdf"))
        if not pdf_files:
            logger.warning(f"No PDF files found in {pdf_dir}")
            return []

        results = []
        total = len(pdf_files)

        for idx, pdf_path in enumerate(pdf_files, 1):
            filename = pdf_path.name
            stem = pdf_path.stem
            output_path = out_dir / f"Summary-{stem}.{output_format}"

            # Skip if output already exists
            if skip_existing and output_path.exists():
                logger.info(f"[{idx}/{total}] Skipping (output exists): {filename}")
                results.append({"filename": filename, "status": "skipped"})
                continue

            logger.info(f"[{idx}/{total}] Processing: {filename}")

            try:
                # Ingest this paper
                self._ingest_single_paper(pdf_path)

                # Summarize using only this paper's chunks
                summary = self.summarize(
                    prompt=prompt,
                    use_cot=use_cot,
                    n_chunks=n_chunks,
                    paper_filter=filename,
                    mode=mode,
                )

                # Save
                title = f"Summary - {stem}"
                if output_format == "json":
                    save_as_json(summary, output_path)
                elif output_format == "csv":
                    save_as_csv(summary, output_path)
                else:
                    save_as_docx(summary, output_path, title=title)

                logger.info(f"  Saved: {output_path.name}")
                results.append({"filename": filename, "status": "success"})

            except Exception as e:
                logger.error(f"  Error: {e}")
                results.append({"filename": filename, "status": f"error: {e}"})

        return results

    # ------------------------------------------------------------------
    # Graph querying
    # ------------------------------------------------------------------

    def query_graph(self, entity: str) -> dict:
        """Query the knowledge graph for information about an entity.

        Args:
            entity: Entity name to look up.

        Returns:
            Dict with entity info, neighbors, and edges.
        """
        info = self._knowledge_graph.get_entity(entity)
        if info is None:
            return {"error": f"Entity '{entity}' not found in knowledge graph"}

        return {
            "entity": entity,
            "info": info,
            "neighbors": self._knowledge_graph.neighbors(entity),
            "edges": self._knowledge_graph.get_edges_for_entity(entity),
            "community": self._communities.get(entity),
        }

    def get_central_entities(self, top_k: int = 10) -> list[tuple[str, float]]:
        """Get the most central entities in the knowledge graph."""
        return self._knowledge_graph.get_central_entities(top_k=top_k)

    def get_communities(self) -> dict[int, list[str]]:
        """Get community assignments as community_id → list of entities."""
        result: dict[int, list[str]] = {}
        for entity, comm_id in self._communities.items():
            result.setdefault(comm_id, []).append(entity)
        return result

    # ------------------------------------------------------------------
    # Explore / Visualize
    # ------------------------------------------------------------------

    def explore(
        self,
        output_path: str | Path = "./explore_dashboard.html",
        open_browser: bool = True,
    ) -> Path:
        """Generate and open an interactive HTML dashboard to explore the knowledge base.

        Shows the knowledge graph, vector store contents, community structure,
        and summary statistics in a browser-based interface.

        Args:
            output_path: Where to save the HTML file.
            open_browser: If True, opens in the default browser.

        Returns:
            Path to the generated HTML file.
        """
        from res_sum.explore.dashboard import generate_dashboard

        return generate_dashboard(
            knowledge_graph=self._knowledge_graph,
            vector_store=self._vector_store,
            communities=self._communities,
            community_summaries=self._community_summaries,
            output_path=output_path,
            open_browser=open_browser,
        )

    # ------------------------------------------------------------------
    # Utility
    # ------------------------------------------------------------------

    @property
    def vector_store(self) -> VectorStore:
        """Access the underlying vector store."""
        return self._vector_store

    @property
    def knowledge_graph(self) -> KnowledgeGraph:
        """Access the underlying knowledge graph."""
        return self._knowledge_graph

    @property
    def llm(self):
        """Access the underlying LLM."""
        return self._llm

    def stats(self) -> dict:
        """Return statistics about the current state."""
        return {
            "domain": self.domain_config.name,
            "llm_provider": self.llm_provider,
            "model": self.model,
            "total_chunks": self._vector_store.count(),
            "papers_ingested": len(self._ingested_papers),
            "graph_nodes": self._knowledge_graph.num_nodes(),
            "graph_edges": self._knowledge_graph.num_edges(),
            "communities": len(set(self._communities.values())) if self._communities else 0,
            "data_dir": str(self.data_dir),
        }
