"""Graph partitioning engine for multi-agent swarm orchestration."""

from __future__ import annotations

import os
import sqlite3
from collections import Counter, defaultdict

import networkx as nx

from roam.db.connection import batched_in
from roam.graph.clusters import detect_clusters


def partition_for_agents(
    G: nx.DiGraph,
    conn: sqlite3.Connection,
    n_agents: int,
    target_files: list[str] | None = None,
) -> dict:
    """Partition the symbol graph into non-overlapping agent work zones.

    Algorithm:
    1. If *target_files* is given, extract the subgraph for those files.
    2. Run Louvain community detection.
    3. Merge or split clusters to match *n_agents*.
    4. For each partition compute write files, read-only files,
       shared interfaces, and contracts.

    Returns a dict with ``agents``, ``merge_order``, ``conflict_probability``,
    and ``shared_interfaces``.
    """
    if n_agents < 1:
        n_agents = 1

    # -- 1. Scope the graph ----------------------------------------
    if target_files:
        target_set = set(target_files)
        keep = {n for n in G.nodes if G.nodes[n].get("file_path", "") in target_set}
        if not keep:
            # Try prefix matching (directories)
            keep = {n for n in G.nodes if any(G.nodes[n].get("file_path", "").startswith(t) for t in target_set)}
        if keep:
            G = G.subgraph(keep).copy()

    if len(G) == 0:
        return _empty_result(n_agents)

    # -- 2. Detect clusters ----------------------------------------
    cluster_map = detect_clusters(G)
    if not cluster_map:
        # Assign every node to cluster 0
        cluster_map = {n: 0 for n in G.nodes}

    # Group nodes by cluster
    groups: dict[int, set[int]] = defaultdict(set)
    for node_id, cid in cluster_map.items():
        groups[cid].add(node_id)

    # -- 3. Adjust cluster count to match n_agents -----------------
    partitions = _adjust_cluster_count(G, groups, n_agents)

    # -- 4. Build agent descriptors --------------------------------
    agents = _build_agent_descriptors(G, conn, partitions)

    # -- 5. Shared interfaces --------------------------------------
    shared_interfaces = _find_shared_interfaces(G, conn, partitions)

    # -- 6. Conflict probability -----------------------------------
    conflict_prob = compute_conflict_probability(G, partitions)

    # -- 7. Merge order --------------------------------------------
    merge_order = compute_merge_order(G, partitions)

    # Count write conflicts (files appearing in multiple write lists)
    all_write_files: list[str] = []
    for a in agents:
        all_write_files.extend(a["write_files"])
    file_counts = Counter(all_write_files)
    write_conflicts = sum(1 for c in file_counts.values() if c > 1)

    return {
        "agents": agents,
        "merge_order": merge_order,
        "conflict_probability": round(conflict_prob, 4),
        "shared_interfaces": [si["symbol"] for si in shared_interfaces],
        "write_conflicts": write_conflicts,
    }


def compute_conflict_probability(G: nx.DiGraph, partitions: list[dict[str, set[int]]]) -> float:
    """Ratio of cross-partition edges to total edges."""
    if len(G.edges) == 0:
        return 0.0

    # Build node -> partition index
    node_part: dict[int, int] = {}
    for idx, p in enumerate(partitions):
        for n in p["nodes"]:
            node_part[n] = idx

    cross = 0
    total = 0
    for u, v in G.edges:
        if u in node_part and v in node_part:
            total += 1
            if node_part[u] != node_part[v]:
                cross += 1

    return cross / total if total > 0 else 0.0


def compute_merge_order(G: nx.DiGraph, partitions: list[dict[str, set[int]]]) -> list[int]:
    """Topological sort of partitions by dependency direction (leaves first).

    Returns a list of 1-based agent IDs in merge order.
    """
    if not partitions:
        return []

    # Build node -> partition index
    node_part: dict[int, int] = {}
    for idx, p in enumerate(partitions):
        for n in p["nodes"]:
            node_part[n] = idx

    # Build partition dependency graph
    PG = nx.DiGraph()
    for idx in range(len(partitions)):
        PG.add_node(idx)

    for u, v in G.edges:
        pu = node_part.get(u)
        pv = node_part.get(v)
        if pu is not None and pv is not None and pu != pv:
            PG.add_edge(pu, pv)

    # Topological sort on condensed version (handles cycles)
    cond = nx.condensation(PG)
    scc_map = cond.graph["mapping"]

    scc_order = list(nx.topological_sort(cond))
    # Reverse: leaves first (SCCs with no outgoing edges come first)
    scc_order = list(reversed(scc_order))

    # Map SCC order back to partition IDs (1-based)
    order: list[int] = []
    for scc_id in scc_order:
        members = [pid for pid, sid in scc_map.items() if sid == scc_id]
        members.sort()
        for pid in members:
            order.append(pid + 1)  # 1-based

    return order


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _empty_result(n_agents: int) -> dict:
    """Return a valid but empty result."""
    agents = [
        {
            "id": i + 1,
            "write_files": [],
            "read_only_files": [],
            "symbols_owned": 0,
            "contracts": [],
            "cluster_label": f"empty-{i + 1}",
        }
        for i in range(n_agents)
    ]
    return {
        "agents": agents,
        "merge_order": list(range(1, n_agents + 1)),
        "conflict_probability": 0.0,
        "shared_interfaces": [],
        "write_conflicts": 0,
    }


def _adjust_cluster_count(
    G: nx.DiGraph,
    groups: dict[int, set[int]],
    n_agents: int,
) -> list[dict[str, set[int]]]:
    """Merge or split clusters until we have exactly *n_agents* partitions.

    Returns a list of dicts with key ``nodes`` (set of node IDs).
    """
    # Start with sorted cluster groups (biggest first)
    sorted_groups = sorted(groups.values(), key=len, reverse=True)

    # If we already have the right number, done
    if len(sorted_groups) == n_agents:
        return [{"nodes": s} for s in sorted_groups]

    # If more clusters than agents: merge smallest. A min-heap keyed on
    # size replaces the former re-sort-twice-per-iteration loop (flagged by
    # the algo detector); the tiebreaker counter keeps the sets themselves
    # out of tuple comparison.
    if len(sorted_groups) > n_agents:
        import heapq
        from itertools import count

        tiebreak = count()
        heap = [(len(g), next(tiebreak), g) for g in sorted_groups]
        heapq.heapify(heap)
        while len(heap) > n_agents:
            _, _, smallest = heapq.heappop(heap)
            _, _, second_smallest = heapq.heappop(heap)
            merged = smallest | second_smallest
            heapq.heappush(heap, (len(merged), next(tiebreak), merged))
        groups = sorted((g for _, _, g in heap), key=len, reverse=True)
        return [{"nodes": s} for s in groups]

    # If fewer clusters than agents: split largest
    while len(sorted_groups) < n_agents:
        sorted_groups.sort(key=len)
        largest = sorted_groups.pop()
        if len(largest) < 2:
            # Cannot split a singleton; add it back and create an empty one
            sorted_groups.append(largest)
            sorted_groups.append(set())
            continue
        # Split by betweenness centrality cut
        sub = G.subgraph(largest)
        try:
            # Speed guard (mirrors graph/pagerank.py): exact betweenness is
            # O(V*E). k-sample the pivots on large subgraphs — the cut-node pick
            # is a heuristic, so approximate centrality is fine, and the
            # try/except already falls back to bisection. seed keeps the sampled
            # pick reproducible across invocations.
            _m = sub.number_of_nodes()
            if _m <= 1000:
                bc = nx.betweenness_centrality(sub)
            else:
                _k = min(_m, max(200, int(_m**0.5 * 5)))
                bc = nx.betweenness_centrality(sub, k=_k, seed=1)
            # Remove the highest-betweenness node to split
            cut_node = max(bc, key=bc.get)
            remaining = largest - {cut_node}
            # Get connected components of the subgraph without cut_node
            sub_remaining = G.subgraph(remaining).to_undirected()
            components = list(nx.connected_components(sub_remaining))
            if len(components) >= 2:
                # Take the two largest components; assign cut_node to the first
                components.sort(key=len, reverse=True)
                first = components[0] | {cut_node}
                rest = set()
                for c in components[1:]:
                    rest |= c
                sorted_groups.append(first)
                sorted_groups.append(rest)
            else:
                # Could not split by removing cut node; split by bisection
                half = len(largest) // 2
                nodes_list = sorted(largest)
                sorted_groups.append(set(nodes_list[:half]))
                sorted_groups.append(set(nodes_list[half:]))
        except (nx.NetworkXError, ValueError):
            # Fallback: simple bisection for graph/centrality failures.
            half = len(largest) // 2
            nodes_list = sorted(largest)
            sorted_groups.append(set(nodes_list[:half]))
            sorted_groups.append(set(nodes_list[half:]))

    # Trim excess (shouldn't happen but safety)
    partitions = [{"nodes": s} for s in sorted_groups[:n_agents]]
    return partitions


def _node_partition_index(partitions: list[dict[str, set[int]]]) -> dict[int, int]:
    """flatten partitions into ``{node_id: partition_index}``."""
    node_part: dict[int, int] = {}
    for idx, p in enumerate(partitions):
        for n in p["nodes"]:
            node_part[n] = idx
    return node_part


def _fetch_node_metadata(conn: sqlite3.Connection, node_ids: list[int]):
    """load (file, name, signature) for every node in one batch."""
    node_to_file: dict[int, str] = {}
    node_to_name: dict[int, str] = {}
    node_to_sig: dict[int, str] = {}
    if not node_ids:
        return node_to_file, node_to_name, node_to_sig
    rows = batched_in(
        conn,
        "SELECT s.id, s.name, s.signature, f.path FROM symbols s JOIN files f ON s.file_id = f.id WHERE s.id IN ({ph})",
        node_ids,
    )
    for r in rows:
        node_to_file[r["id"]] = r["path"].replace("\\", "/")
        node_to_name[r["id"]] = r["name"]
        node_to_sig[r["id"]] = r["signature"] or r["name"]
    return node_to_file, node_to_name, node_to_sig


def _file_majority_owners(node_part: dict[int, int], node_to_file: dict[int, str]) -> dict[int, set[str]]:
    """assign each file exclusively to its majority-vote partition."""
    file_partition_counts: dict[str, Counter] = defaultdict(Counter)
    for n, pidx in node_part.items():
        fp = node_to_file.get(n)
        if fp:
            file_partition_counts[fp][pidx] += 1
    partition_write_files: dict[int, set[str]] = defaultdict(set)
    for fp, counts in file_partition_counts.items():
        owner = counts.most_common(1)[0][0]
        partition_write_files[owner].add(fp)
    return partition_write_files


def _read_only_files_for(
    G, nodes: set[int], node_part: dict[int, int], node_to_file: dict[int, str], write_files: set[str]
) -> set[str]:
    """files with cross-partition edges but not owned by this agent."""
    read_only: set[str] = set()
    for n in nodes:
        for neighbour in (*G.predecessors(n), *G.successors(n)):
            if neighbour in nodes or neighbour not in node_part:
                continue
            fp = node_to_file.get(neighbour)
            if fp and fp not in write_files:
                read_only.add(fp)
    return read_only


def _boundary_contracts(
    G, nodes: set[int], node_part: dict[int, int], node_to_name: dict[int, str], node_to_sig: dict[int, str]
) -> list[str]:
    """list cross-partition symbols this agent must not change."""
    contracts: list[str] = []
    seen: set[str] = set()
    for n in nodes:
        for neighbour in (*G.successors(n), *G.predecessors(n)):
            if neighbour in nodes or neighbour not in node_part:
                continue
            name = node_to_name.get(neighbour)
            sig = node_to_sig.get(neighbour)
            if not (name and sig):
                continue
            contract = f"do NOT modify {sig} signature"
            if contract not in seen:
                seen.add(contract)
                contracts.append(contract)
    return contracts


def _cluster_label_for(nodes: set[int], node_to_file: dict[int, str], idx: int) -> str:
    """directory-majority label, with sensible fallbacks."""
    dirs = [os.path.dirname(node_to_file.get(n, "")).replace("\\", "/") for n in nodes if n in node_to_file]
    counts = Counter(dirs) if dirs else Counter()
    if not counts:
        return f"partition-{idx + 1}"
    label = counts.most_common(1)[0][0]
    if not label:
        return f"root-{idx + 1}"
    return label.rstrip("/").rsplit("/", 1)[-1] or f"partition-{idx + 1}"


def _build_agent_descriptors(
    G: nx.DiGraph,
    conn: sqlite3.Connection,
    partitions: list[dict[str, set[int]]],
) -> list[dict]:
    """Build per-agent descriptor dicts.

    orchestrator only. Per-step logic lives in
    ``_node_partition_index``, ``_fetch_node_metadata``,
    ``_file_majority_owners``, ``_read_only_files_for``,
    ``_boundary_contracts``, ``_cluster_label_for``.
    Cognitive complexity dropped from 161 to ~10.

    File ownership is determined by majority vote: each file is
    assigned exclusively to the partition that owns the most of its
    symbols. This guarantees no write overlap between agents.
    """
    node_part = _node_partition_index(partitions)
    node_to_file, node_to_name, node_to_sig = _fetch_node_metadata(conn, list(node_part.keys()))
    partition_write_files = _file_majority_owners(node_part, node_to_file)

    agents: list[dict] = []
    for idx, p in enumerate(partitions):
        nodes = p["nodes"]
        if not nodes:
            agents.append(
                {
                    "id": idx + 1,
                    "write_files": [],
                    "read_only_files": [],
                    "symbols_owned": 0,
                    "contracts": [],
                    "cluster_label": f"partition-{idx + 1}",
                }
            )
            continue
        write_files = partition_write_files.get(idx, set())
        read_only = _read_only_files_for(G, nodes, node_part, node_to_file, write_files)
        contracts = _boundary_contracts(G, nodes, node_part, node_to_name, node_to_sig)
        agents.append(
            {
                "id": idx + 1,
                "write_files": sorted(write_files),
                "read_only_files": sorted(read_only),
                "symbols_owned": len(nodes),
                "contracts": contracts[:10],
                "cluster_label": _cluster_label_for(nodes, node_to_file, idx),
            }
        )
    return agents


def _find_shared_interfaces(
    G: nx.DiGraph,
    conn: sqlite3.Connection,
    partitions: list[dict[str, set[int]]],
) -> list[dict]:
    """Find symbols at partition boundaries (high betweenness in cross-partition edges)."""
    node_part: dict[int, int] = {}
    for idx, p in enumerate(partitions):
        for n in p["nodes"]:
            node_part[n] = idx

    # Count cross-partition edges per node
    boundary_counts: Counter = Counter()
    for u, v in G.edges:
        pu = node_part.get(u)
        pv = node_part.get(v)
        if pu is not None and pv is not None and pu != pv:
            boundary_counts[u] += 1
            boundary_counts[v] += 1

    if not boundary_counts:
        return []

    # Take top boundary symbols
    top_ids = [nid for nid, _ in boundary_counts.most_common(20)]

    # Fetch symbol info
    result = []
    if top_ids:
        rows = batched_in(
            conn,
            "SELECT s.id, s.name, f.path FROM symbols s JOIN files f ON s.file_id = f.id WHERE s.id IN ({ph})",
            top_ids,
        )
        seen = set()
        for r in rows:
            sym_label = f"{r['path'].replace(chr(92), '/')}::{r['name']}"
            if sym_label not in seen:
                seen.add(sym_label)
                result.append(
                    {
                        "symbol": sym_label,
                        "boundary_edges": boundary_counts.get(r["id"], 0),
                    }
                )

    result.sort(key=lambda x: -x["boundary_edges"])
    return result[:10]
