"""Feature extraction for heterogeneous financial risk graphs."""

from __future__ import annotations

from collections import Counter, defaultdict, deque
from dataclasses import dataclass, field
from typing import Any

from .models import RiskGraphEdge, RiskGraphNode


@dataclass(slots=True)
class RiskGraphFeatureSet:
    node_type_counts: dict[str, int] = field(default_factory=dict)
    edge_type_counts: dict[str, int] = field(default_factory=dict)
    meta_path_counts: dict[str, int] = field(default_factory=dict)
    degree_centrality: dict[str, float] = field(default_factory=dict)
    clusters: list[dict[str, Any]] = field(default_factory=list)
    neighbor_risk_density: dict[str, float] = field(default_factory=dict)

    def as_dict(self) -> dict[str, Any]:
        return {
            "node_type_counts": dict(self.node_type_counts),
            "edge_type_counts": dict(self.edge_type_counts),
            "meta_path_counts": dict(self.meta_path_counts),
            "degree_centrality": dict(self.degree_centrality),
            "clusters": list(self.clusters),
            "neighbor_risk_density": dict(self.neighbor_risk_density),
        }


class HeterogeneousRiskGraphFeatureBuilder:
    def build(
        self,
        nodes: list[RiskGraphNode],
        edges: list[RiskGraphEdge],
        *,
        risk_node_keys: set[str] | None = None,
    ) -> RiskGraphFeatureSet:
        node_by_key = {node.key: node for node in nodes}
        adjacency = _build_adjacency(edges)
        risk_keys = set(risk_node_keys or set())
        return RiskGraphFeatureSet(
            node_type_counts=dict(Counter(node.node_type for node in nodes)),
            edge_type_counts=dict(Counter(edge.edge_type for edge in edges)),
            meta_path_counts=_meta_path_counts(node_by_key, adjacency),
            degree_centrality=_degree_centrality(node_by_key, adjacency),
            clusters=_clusters(node_by_key, adjacency),
            neighbor_risk_density=_neighbor_risk_density(node_by_key, adjacency, risk_keys),
        )


def _build_adjacency(edges: list[RiskGraphEdge]) -> dict[str, list[tuple[str, str]]]:
    adjacency: dict[str, list[tuple[str, str]]] = defaultdict(list)
    for edge in edges:
        adjacency[edge.source_key].append((edge.target_key, edge.edge_type))
        adjacency[edge.target_key].append((edge.source_key, edge.edge_type))
    return adjacency


def _meta_path_counts(
    node_by_key: dict[str, RiskGraphNode],
    adjacency: dict[str, list[tuple[str, str]]],
) -> dict[str, int]:
    counts: Counter[str] = Counter()
    for source_key, first_hops in adjacency.items():
        source = node_by_key.get(source_key)
        if source is None:
            continue
        for middle_key, first_edge_type in first_hops:
            middle = node_by_key.get(middle_key)
            if middle is None:
                continue
            for target_key, second_edge_type in adjacency.get(middle_key, []):
                if target_key == source_key:
                    continue
                target = node_by_key.get(target_key)
                if target is None:
                    continue
                key = (
                    f"{source.node_type}->{first_edge_type}->{middle.node_type}"
                    f"->{second_edge_type}->{target.node_type}"
                )
                counts[key] += 1
    return dict(counts)


def _degree_centrality(
    node_by_key: dict[str, RiskGraphNode],
    adjacency: dict[str, list[tuple[str, str]]],
) -> dict[str, float]:
    denominator = max(1, len(node_by_key) - 1)
    return {
        node_key: round(len(adjacency.get(node_key, [])) / denominator, 4)
        for node_key in node_by_key
    }


def _clusters(
    node_by_key: dict[str, RiskGraphNode],
    adjacency: dict[str, list[tuple[str, str]]],
) -> list[dict[str, Any]]:
    visited: set[str] = set()
    clusters: list[dict[str, Any]] = []
    for start_key in node_by_key:
        if start_key in visited:
            continue
        queue: deque[str] = deque([start_key])
        visited.add(start_key)
        members: list[str] = []
        type_counts: Counter[str] = Counter()
        while queue:
            node_key = queue.popleft()
            members.append(node_key)
            type_counts[node_by_key[node_key].node_type] += 1
            for next_key, _ in adjacency.get(node_key, []):
                if next_key in visited or next_key not in node_by_key:
                    continue
                visited.add(next_key)
                queue.append(next_key)
        clusters.append(
            {
                "size": len(members),
                "node_keys": sorted(members),
                "node_type_counts": dict(type_counts),
            }
        )
    return sorted(clusters, key=lambda item: item["size"], reverse=True)


def _neighbor_risk_density(
    node_by_key: dict[str, RiskGraphNode],
    adjacency: dict[str, list[tuple[str, str]]],
    risk_keys: set[str],
) -> dict[str, float]:
    density: dict[str, float] = {}
    for node_key in node_by_key:
        neighbors = [target for target, _ in adjacency.get(node_key, [])]
        if not neighbors:
            density[node_key] = 0.0
            continue
        risk_neighbor_count = sum(1 for target in neighbors if target in risk_keys)
        density[node_key] = round(risk_neighbor_count / len(neighbors), 4)
    return density
