"""Graph builder - creates RustGraph directly from money flows.

No NetworkX dependency. RustGraph (petgraph) is the primary graph structure.
Edge attributes are stored in a separate Python dict when needed.
"""

from typing import List, Dict, Any, Optional, Tuple
from collections import defaultdict

from loguru import logger

# Import Rust backend
try:
    from chainswarm_analyzers_baseline._rust import RustGraph
except ImportError:
    raise ImportError(
        "Rust backend required for graph building. "
        "Build with: maturin develop --release"
    )


class GraphData:
    """Container for graph + edge attributes.

    RustGraph stores structure and weights.
    Edge attributes stored separately for access when needed.
    """

    def __init__(
        self,
        graph: RustGraph,
        edge_attrs: Optional[Dict[Tuple[str, str], Dict[str, Any]]] = None,
        node_attrs: Optional[Dict[str, Dict[str, Any]]] = None,
    ):
        self.graph = graph
        self.edge_attrs = edge_attrs or {}
        self.node_attrs = node_attrs or {}

    def get_edge_attr(self, source: str, target: str, attr: str, default: Any = None) -> Any:
        """Get edge attribute by name."""
        key = (source, target)
        if key in self.edge_attrs:
            return self.edge_attrs[key].get(attr, default)
        return default

    def get_node_attr(self, node: str, attr: str, default: Any = None) -> Any:
        """Get node attribute by name."""
        if node in self.node_attrs:
            return self.node_attrs[node].get(attr, default)
        return default

    def set_node_attr(self, node: str, attr: str, value: Any) -> None:
        """Set node attribute."""
        if node not in self.node_attrs:
            self.node_attrs[node] = {}
        self.node_attrs[node][attr] = value


def build_graph(money_flows: List[Dict[str, Any]], store_attrs: bool = True) -> GraphData:
    """Build RustGraph from money flows.

    Args:
        money_flows: List of flow dicts with from_address, to_address, amount_usd_sum, etc.
        store_attrs: Whether to store edge attributes (for later access)

    Returns:
        GraphData containing RustGraph and optional edge/node attributes
    """
    sources = []
    targets = []
    weights = []
    edge_attrs = {} if store_attrs else None

    for flow in money_flows:
        from_addr = flow['from_address']
        to_addr = flow['to_address']
        amount_usd_sum = float(flow['amount_usd_sum'])

        sources.append(from_addr)
        targets.append(to_addr)
        weights.append(amount_usd_sum)

        if store_attrs:
            edge_attrs[(from_addr, to_addr)] = {
                'amount_usd_sum': amount_usd_sum,
                'tx_count': int(flow.get('tx_count', 1)),
                'first_seen_timestamp': int(flow.get('first_seen_timestamp', 0)),
                'last_seen_timestamp': int(flow.get('last_seen_timestamp', 0)),
                'active_days': int(flow.get('active_days', 1)),
                'avg_tx_size_usd': float(flow.get('avg_tx_size_usd', amount_usd_sum)),
                'unique_assets': int(flow.get('unique_assets', 1)),
                'dominant_asset': flow.get('dominant_asset', ''),
                'hourly_pattern': flow.get('hourly_pattern', []),
                'weekly_pattern': flow.get('weekly_pattern', []),
                'reciprocity_ratio': float(flow.get('reciprocity_ratio', 0.0)),
                'is_bidirectional': bool(flow.get('is_bidirectional', False)),
            }

    # Build RustGraph
    graph = RustGraph()
    graph.add_edges(sources, targets, weights)

    node_count, edge_count = graph.stats()
    logger.info(f"Graph built: {node_count} nodes, {edge_count} edges")

    return GraphData(graph=graph, edge_attrs=edge_attrs)


def build_money_flow_graph(money_flows: List[Dict[str, Any]]) -> GraphData:
    """Build graph from money flows (compatibility wrapper).

    Returns GraphData instead of nx.DiGraph.
    """
    return build_graph(money_flows, store_attrs=True)


def add_node_volume_attributes(graph_data: GraphData) -> None:
    """Add total volume attribute to nodes.

    Uses RustGraph.volume_map() for efficient calculation.
    """
    volume_map = graph_data.graph.volume_map()

    for node, (in_vol, out_vol) in volume_map.items():
        graph_data.set_node_attr(node, 'total_volume_usd', in_vol + out_vol)
        graph_data.set_node_attr(node, 'in_volume_usd', in_vol)
        graph_data.set_node_attr(node, 'out_volume_usd', out_vol)


def extract_addresses_from_flows(flows: List[Dict[str, Any]]) -> List[str]:
    """Extract unique addresses from flows."""
    addresses_set = set()
    for flow in flows:
        addresses_set.add(flow['from_address'])
        addresses_set.add(flow['to_address'])
    return sorted(list(addresses_set))


def build_flows_index_by_address(graph_data: GraphData) -> Dict[str, List[Dict[str, Any]]]:
    """Build index of flows by address.

    Uses RustGraph.edges() to get edge list.
    """
    flows_by_address: Dict[str, List[Dict[str, Any]]] = defaultdict(list)

    for src, tgt, weight in graph_data.graph.edges():
        attrs = graph_data.edge_attrs.get((src, tgt), {})
        flow = {
            'from_address': src,
            'to_address': tgt,
            'amount_usd_sum': weight,
            'tx_count': attrs.get('tx_count', 1),
        }
        flows_by_address[src].append(flow)
        flows_by_address[tgt].append(flow)

    return dict(flows_by_address)


# Utility functions for working with raw arrays
def flows_to_arrays(money_flows: List[Dict[str, Any]]) -> Tuple[List[str], List[str], List[float]]:
    """Convert money flows to raw arrays for Rust.

    Returns:
        (sources, targets, weights) tuple
    """
    sources = []
    targets = []
    weights = []

    for flow in money_flows:
        sources.append(flow['from_address'])
        targets.append(flow['to_address'])
        weights.append(float(flow['amount_usd_sum']))

    return sources, targets, weights


def build_graph_from_arrays(
    sources: List[str],
    targets: List[str],
    weights: List[float]
) -> RustGraph:
    """Build RustGraph directly from arrays.

    Use this when you don't need edge attributes.
    """
    graph = RustGraph()
    graph.add_edges(sources, targets, weights)
    return graph


def build_graph_data_from_arrays(
    sources: List[str],
    targets: List[str],
    weights: List[float],
    store_edge_set: bool = True
) -> GraphData:
    """Build GraphData directly from arrays (fastest path).

    This is the preferred method when loading from Rust parquet reader.
    Only stores edge set for O(1) lookup, not full edge attributes.

    Args:
        sources: List of source addresses
        targets: List of target addresses
        weights: List of edge weights (amount_usd_sum)
        store_edge_set: If True, stores (src, tgt) -> weight dict for edge lookup

    Returns:
        GraphData with RustGraph and minimal edge attrs
    """
    graph = RustGraph()
    graph.add_edges(sources, targets, weights)

    # Store minimal edge attrs - just the weight for O(1) edge existence check
    edge_attrs = None
    if store_edge_set:
        edge_attrs = {
            (src, tgt): {'amount_usd_sum': weight}
            for src, tgt, weight in zip(sources, targets, weights)
        }

    node_count, edge_count = graph.stats()
    logger.info(f"Graph built from arrays: {node_count} nodes, {edge_count} edges")

    return GraphData(graph=graph, edge_attrs=edge_attrs)


def extract_addresses_from_arrays(sources: List[str], targets: List[str]) -> List[str]:
    """Extract unique addresses from source/target arrays.

    Faster than extract_addresses_from_flows when you already have arrays.
    """
    addresses_set = set(sources)
    addresses_set.update(targets)
    return sorted(list(addresses_set))
