"""Community detection on the knowledge graph using Leiden algorithm."""

from __future__ import annotations

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

import networkx as nx
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import SystemMessage, HumanMessage

from res_sum.graph.knowledge_graph import KnowledgeGraph

logger = logging.getLogger("res_sum")


def detect_communities(
    kg: KnowledgeGraph,
    resolution: float = 1.0,
) -> dict[str, int]:
    """Detect communities in the knowledge graph.

    Tries graspologic's Leiden algorithm first. Falls back to NetworkX's
    Louvain if graspologic is not installed.

    Args:
        kg: The knowledge graph.
        resolution: Resolution parameter (higher = more communities).

    Returns:
        Dict mapping entity name → community ID.
    """
    graph = kg.graph
    if graph.number_of_nodes() == 0:
        return {}

    # Convert to undirected for community detection
    undirected = graph.to_undirected()
    # Remove self-loops and isolated nodes for better results
    undirected.remove_edges_from(nx.selfloop_edges(undirected))

    # Try graspologic Leiden first
    try:
        from graspologic.partition import leiden

        partition = leiden(undirected, resolution=resolution)
        node_to_community = {}
        for node_id, community_id in partition.items():
            # graspologic uses integer node indices, need to map back
            if isinstance(node_id, int):
                nodes = list(undirected.nodes())
                if node_id < len(nodes):
                    node_to_community[nodes[node_id]] = int(community_id)
            else:
                node_to_community[str(node_id)] = int(community_id)
        logger.info(
            f"Leiden community detection: {len(set(node_to_community.values()))} "
            f"communities from {graph.number_of_nodes()} nodes"
        )
        return node_to_community

    except ImportError:
        logger.info("graspologic not installed, falling back to NetworkX Louvain")

    # Fallback: NetworkX Louvain
    try:
        communities = nx.community.louvain_communities(
            undirected, resolution=resolution, seed=42
        )
        node_to_community = {}
        for comm_id, members in enumerate(communities):
            for member in members:
                node_to_community[member] = comm_id
        logger.info(
            f"Louvain community detection: {len(communities)} communities "
            f"from {graph.number_of_nodes()} nodes"
        )
        return node_to_community

    except Exception as e:
        logger.warning(f"Community detection failed: {e}")
        # Assign all nodes to community 0
        return {node: 0 for node in graph.nodes()}


def generate_community_summaries(
    kg: KnowledgeGraph,
    communities: dict[str, int],
    llm: BaseChatModel,
) -> dict[int, str]:
    """Generate a text summary for each community using the LLM.

    Args:
        kg: The knowledge graph.
        communities: Dict mapping entity name → community ID.
        llm: The LLM to use for summarization.

    Returns:
        Dict mapping community ID → summary text.
    """
    # Group entities by community
    community_entities: dict[int, list[str]] = {}
    for entity, comm_id in communities.items():
        community_entities.setdefault(comm_id, []).append(entity)

    summaries: dict[int, str] = {}
    total = len(community_entities)

    for comm_id, entities in community_entities.items():
        logger.debug(f"  Summarizing community {comm_id + 1}/{total} ({len(entities)} entities)")

        # Build context: entity info + relationships within community
        context_parts = []
        for entity in entities[:30]:  # Cap to avoid huge prompts
            node_data = kg.get_entity(entity)
            if node_data:
                etype = node_data.get("entity_type", "")
                context_parts.append(f"- {entity} ({etype})")

            # Add relationships within this community
            for edge in kg.get_edges_for_entity(entity):
                target = edge.get("target", "")
                source = edge.get("source", "")
                other = target if source == entity else source
                if other in entities:
                    rel_type = edge.get("relation_type", "")
                    evidence = edge.get("evidence", "")
                    if evidence:
                        context_parts.append(f"  → {entity} {rel_type} {other}: {evidence}")

        if not context_parts:
            summaries[comm_id] = f"Community {comm_id}: {', '.join(entities[:5])}"
            continue

        context = "\n".join(context_parts)

        messages = [
            SystemMessage(content=(
                "You are a scientific knowledge analyst. Summarize the following "
                "group of related entities and their relationships in 2-3 sentences. "
                "Focus on what connects these entities and what the key theme or "
                "finding is."
            )),
            HumanMessage(content=f"Entities and relationships:\n\n{context}"),
        ]

        try:
            response = llm.invoke(messages)
            summaries[comm_id] = response.content.strip()
        except Exception as e:
            logger.warning(f"Failed to summarize community {comm_id}: {e}")
            summaries[comm_id] = f"Community {comm_id}: {', '.join(entities[:5])}"

    return summaries


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

def save_communities(
    communities: dict[str, int],
    summaries: dict[int, str],
    path: str | Path,
) -> None:
    """Save community data to JSON."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    data = {
        "assignments": communities,
        "summaries": {str(k): v for k, v in summaries.items()},
    }
    with open(path, "w") as f:
        json.dump(data, f, indent=2)
    logger.info(f"Communities saved to {path}")


def load_communities(path: str | Path) -> tuple[dict[str, int], dict[int, str]]:
    """Load community data from JSON."""
    path = Path(path)
    if not path.exists():
        return {}, {}
    with open(path) as f:
        data = json.load(f)
    communities = data.get("assignments", {})
    summaries = {int(k): v for k, v in data.get("summaries", {}).items()}
    return communities, summaries
