"""NetworkX-based knowledge graph construction, persistence, and querying."""

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Optional

import networkx as nx

from res_sum.schemas import ExtractedEntity, ExtractedRelationship, ExtractionResult

logger = logging.getLogger("res_sum")


class KnowledgeGraph:
    """A knowledge graph built from extracted entities and relationships.

    Backed by NetworkX. Nodes are entities, edges are relationships.
    Persists to GraphML format.
    """

    def __init__(self) -> None:
        self._graph = nx.MultiDiGraph()

    @property
    def graph(self) -> nx.MultiDiGraph:
        """Access the underlying NetworkX graph."""
        return self._graph

    # ------------------------------------------------------------------
    # Building the graph
    # ------------------------------------------------------------------

    def add_extraction_result(
        self,
        result: ExtractionResult,
        paper_filename: str = "",
    ) -> None:
        """Add entities and relationships from an extraction result."""
        for entity in result.entities:
            self._add_entity(entity, paper_filename)
        for rel in result.relationships:
            self._add_relationship(rel, paper_filename)

    def _add_entity(self, entity: ExtractedEntity, paper_filename: str) -> None:
        """Add or update a knowledge node."""
        name = entity.name.strip()
        if not name:
            return

        # Use node_type if available, fall back to entity_type for backward compat
        node_type = getattr(entity, "node_type", "") or entity.entity_type or "FINDING"
        description = getattr(entity, "description", "") or ""
        evidence = getattr(entity, "evidence", "") or ""
        confidence = getattr(entity, "confidence", "") or ""

        if self._graph.has_node(name):
            # Update existing node — merge source chunks and papers
            node_data = self._graph.nodes[name]

            existing_chunks = set(node_data.get("source_chunk_ids", "").split("|"))
            existing_chunks.add(entity.source_chunk_id)
            existing_chunks.discard("")
            node_data["source_chunk_ids"] = "|".join(sorted(existing_chunks))

            existing_papers = set(node_data.get("papers", "").split("|"))
            existing_papers.add(paper_filename)
            existing_papers.discard("")
            node_data["papers"] = "|".join(sorted(existing_papers))

            # Update description if new one is longer/better
            if len(description) > len(node_data.get("description", "")):
                node_data["description"] = description
        else:
            self._graph.add_node(
                name,
                entity_type=node_type,
                node_type=node_type,
                description=description,
                evidence=evidence,
                confidence=confidence,
                aliases="|".join(entity.aliases),
                source_chunk_ids=entity.source_chunk_id,
                papers=paper_filename,
            )

    def _add_relationship(
        self, rel: ExtractedRelationship, paper_filename: str
    ) -> None:
        """Add a relationship edge."""
        source = rel.source.strip()
        target = rel.target.strip()
        if not source or not target:
            return

        # Ensure both nodes exist (create minimal nodes if needed)
        if not self._graph.has_node(source):
            self._graph.add_node(source, entity_type="FINDING", node_type="FINDING", description="", evidence="", confidence="", aliases="", source_chunk_ids="", papers=paper_filename)
        if not self._graph.has_node(target):
            self._graph.add_node(target, entity_type="FINDING", node_type="FINDING", description="", evidence="", confidence="", aliases="", source_chunk_ids="", papers=paper_filename)

        self._graph.add_edge(
            source,
            target,
            relation_type=rel.relation_type,
            evidence=rel.evidence,
            source_chunk_id=rel.source_chunk_id,
            paper=paper_filename,
        )

    # ------------------------------------------------------------------
    # Entity resolution
    # ------------------------------------------------------------------

    def resolve_entities(
        self,
        merge_map: dict[str, str],
    ) -> None:
        """Merge duplicate entity nodes based on a merge map.

        Args:
            merge_map: Dict mapping alias → canonical name.
                e.g., {"gray wolf": "Canis lupus", "C. lupus": "Canis lupus"}
        """
        for alias, canonical in merge_map.items():
            if alias == canonical:
                continue
            if not self._graph.has_node(alias):
                continue

            # Ensure canonical node exists
            if not self._graph.has_node(canonical):
                # Rename the alias node to canonical
                nx.relabel_nodes(self._graph, {alias: canonical}, copy=False)
                continue

            # Merge alias node into canonical node
            alias_data = self._graph.nodes[alias]
            canonical_data = self._graph.nodes[canonical]

            # Merge attributes
            for attr in ("aliases", "source_chunk_ids", "papers"):
                existing = set(canonical_data.get(attr, "").split("|"))
                incoming = set(alias_data.get(attr, "").split("|"))
                combined = existing | incoming
                combined.add(alias)
                combined.discard("")
                canonical_data[attr] = "|".join(sorted(combined))

            # Move edges from alias to canonical
            for _, target, data in list(self._graph.edges(alias, data=True)):
                if target != canonical:
                    self._graph.add_edge(canonical, target, **data)
            for source, _, data in list(self._graph.in_edges(alias, data=True)):
                if source != canonical:
                    self._graph.add_edge(source, canonical, **data)

            self._graph.remove_node(alias)

    def build_alias_merge_map(self) -> dict[str, str]:
        """Build a merge map from node aliases and abbreviated names.

        Detects duplicates via:
        1. Exact alias matches (case-insensitive)
        2. Abbreviated genus names ("R. marina" matches "Rhinella marina")
        3. Nodes that are substrings of other nodes with the same entity type

        Returns:
            Dict mapping alias/duplicate node names to canonical node names.
        """
        import re

        merge_map: dict[str, str] = {}
        # Build lookup: lowercase name → canonical node name
        name_lookup: dict[str, str] = {}

        all_nodes = list(self._graph.nodes())

        for node in all_nodes:
            lower = node.lower().strip()
            if lower not in name_lookup:
                name_lookup[lower] = node

            aliases_str = self._graph.nodes[node].get("aliases", "")
            for alias in aliases_str.split("|"):
                alias = alias.strip()
                if not alias:
                    continue
                alias_lower = alias.lower()
                if alias_lower in name_lookup and name_lookup[alias_lower] != node:
                    merge_map[alias] = name_lookup[alias_lower]
                else:
                    name_lookup[alias_lower] = node

        # Detect abbreviated genus names: "R. marina" → "Rhinella marina"
        # Pattern: single capital letter + period + space + epithet
        abbrev_pattern = re.compile(r"^([A-Z])\.\s+(\S+.*)$")

        species_nodes = [
            n for n in all_nodes
            if self._graph.nodes[n].get("entity_type") == "SPECIES"
        ]

        for node in all_nodes:
            match = abbrev_pattern.match(node)
            if not match:
                continue
            first_letter = match.group(1).lower()
            epithet = match.group(2).lower()

            # Find full binomial that matches
            for full_node in species_nodes:
                if full_node == node:
                    continue
                parts = full_node.split()
                if len(parts) >= 2:
                    if (parts[0][0].lower() == first_letter
                            and parts[1].lower() == epithet
                            and node not in merge_map):
                        # "R. marina" → "Rhinella marina"
                        merge_map[node] = full_node
                        break

        # Also check if any non-SPECIES node is an exact case-insensitive
        # duplicate of another node (e.g., "fire ant" vs "Fire ant")
        seen_lower: dict[str, str] = {}
        for node in all_nodes:
            lower = node.lower().strip()
            if lower in seen_lower and seen_lower[lower] != node:
                # Keep the longer or more capitalized version as canonical
                existing = seen_lower[lower]
                if len(node) > len(existing):
                    merge_map[existing] = node
                    seen_lower[lower] = node
                elif node not in merge_map:
                    merge_map[node] = existing
            else:
                seen_lower[lower] = node

        return merge_map

    # ------------------------------------------------------------------
    # Querying
    # ------------------------------------------------------------------

    def neighbors(self, entity: str) -> list[str]:
        """Get neighboring entities."""
        if not self._graph.has_node(entity):
            return []
        return list(set(
            list(self._graph.successors(entity)) +
            list(self._graph.predecessors(entity))
        ))

    def get_entity(self, name: str) -> Optional[dict]:
        """Get entity node data."""
        if self._graph.has_node(name):
            return dict(self._graph.nodes[name])
        return None

    def get_edges_for_entity(self, name: str) -> list[dict]:
        """Get all edges connected to an entity."""
        edges = []
        if not self._graph.has_node(name):
            return edges
        for u, v, data in self._graph.edges(name, data=True):
            edges.append({"source": u, "target": v, **data})
        for u, v, data in self._graph.in_edges(name, data=True):
            edges.append({"source": u, "target": v, **data})
        return edges

    def get_central_entities(self, top_k: int = 10) -> list[tuple[str, float]]:
        """Get the most central entities by degree centrality."""
        if self._graph.number_of_nodes() == 0:
            return []
        centrality = nx.degree_centrality(self._graph)
        sorted_nodes = sorted(centrality.items(), key=lambda x: x[1], reverse=True)
        return sorted_nodes[:top_k]

    def get_entities_by_type(self, entity_type: str) -> list[str]:
        """Get all entities of a specific type."""
        return [
            node for node, data in self._graph.nodes(data=True)
            if data.get("entity_type") == entity_type
        ]

    def get_source_chunk_ids_for_entity(self, name: str) -> list[str]:
        """Get ChromaDB chunk IDs associated with an entity."""
        if not self._graph.has_node(name):
            return []
        chunk_str = self._graph.nodes[name].get("source_chunk_ids", "")
        return [c for c in chunk_str.split("|") if c]

    # ------------------------------------------------------------------
    # Stats
    # ------------------------------------------------------------------

    def num_nodes(self) -> int:
        return self._graph.number_of_nodes()

    def num_edges(self) -> int:
        return self._graph.number_of_edges()

    def summary(self) -> dict:
        """Return a summary of the knowledge graph."""
        type_counts: dict[str, int] = {}
        for _, data in self._graph.nodes(data=True):
            etype = data.get("entity_type", "UNKNOWN")
            type_counts[etype] = type_counts.get(etype, 0) + 1

        return {
            "num_nodes": self.num_nodes(),
            "num_edges": self.num_edges(),
            "entity_types": type_counts,
            "connected_components": nx.number_weakly_connected_components(self._graph),
        }

    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------

    def save(self, path: str | Path) -> None:
        """Save the graph to GraphML format."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        nx.write_graphml(self._graph, str(path))
        logger.info(f"Knowledge graph saved to {path} ({self.num_nodes()} nodes, {self.num_edges()} edges)")

    def load(self, path: str | Path) -> None:
        """Load a graph from GraphML format."""
        path = Path(path)
        if not path.exists():
            logger.warning(f"Graph file not found: {path}")
            return
        self._graph = nx.read_graphml(str(path))
        # Convert to MultiDiGraph if needed
        if not isinstance(self._graph, nx.MultiDiGraph):
            self._graph = nx.MultiDiGraph(self._graph)
        logger.info(f"Knowledge graph loaded from {path} ({self.num_nodes()} nodes, {self.num_edges()} edges)")
