import sys
from pathlib import Path

# Ensure repo root is on path for imports
repo_root = Path(run_dir).resolve().parent.parent
if str(repo_root) not in sys.path:
    sys.path.insert(0, str(repo_root))

import json
import random
import pandas as pd
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots

from notebooks.utils import load_run_data

data = load_run_data(run_dir)
metadata = data["metadata"]
stats = data["stats"]
regions = data["regions"]
bandwidths = data["bandwidths"]
run_id = run_id or data["run_id"]

print(f"Analyzing run: {run_id}")
print(f"Run directory: {run_dir}")

# Build overview table
sim = metadata.get("simulation", {})
fuzzer_cfg = metadata.get("fuzzer", {})
node_counts = metadata.get("node_counts", {})

overview = {
    "Run ID": run_id,
    "Total Nodes": sim.get("total_nodes", "N/A"),
    "Duration (s)": fuzzer_cfg.get("duration_secs", "N/A"),
    "Runner": fuzzer_cfg.get("runner", "N/A"),
    "Seed": fuzzer_cfg.get("seed", "N/A"),
    "Total Subnets": sim.get("total_subnets", "N/A"),
}

for client, count in sorted(node_counts.items()):
    overview[f"Nodes ({client})"] = count

pd.DataFrame([overview]).T.rename(columns={0: "Value"})

bandwidth_stats = stats.get("bandwidth", {})
bandwidth_events = bandwidth_stats.get("events", [])
bandwidth_summary = bandwidth_stats.get("summary", {})

if bandwidth_summary.get("warning"):
    print(f"⚠️ {bandwidth_summary['warning']}")

if not bandwidth_events:
    print("No bandwidth data available.")
else:
    print(f"Bandwidth events: {bandwidth_summary.get('event_count', 0)}")
    print(f"Total bandwidth: {bandwidth_summary.get('total_bytes', 0) / 1024:.1f} KiB")
    print(
        "Duplicate-inclusive inbound gossip: "
        f"{bandwidth_summary.get('duplicate_inclusive_gossip_bytes', 0) / 1024:.1f} KiB"
    )

    df_bw = pd.DataFrame(bandwidth_events)
    df_bw = df_bw[df_bw["slot_offset_ms"].between(0, 4000)].copy()
    available_pairs = sorted(
        (str(host), int(slot))
        for host, slot in df_bw[["host", "slot"]].dropna().drop_duplicates().itertuples(index=False)
    )
    rng = random.Random(str(run_id))
    selected_pairs = sorted(rng.sample(available_pairs, min(3, len(available_pairs))))
    print(f"Selected host/slot pairs: {selected_pairs}")

    selected_pairs_set = set(selected_pairs)
    df_bw = df_bw[
        df_bw.apply(lambda row: (str(row["host"]), int(row["slot"])) in selected_pairs_set, axis=1)
    ].copy()

    if df_bw.empty:
        print("No bandwidth data available for selected host/slot pairs.")
    else:
        df_bw["plot_label"] = df_bw["host"] + " / slot " + df_bw["slot"].astype(str)
        df_bw["flow"] = df_bw["direction"] + " / " + df_bw["delivery"]
        df_bw["offset_bin_ms"] = (df_bw["slot_offset_ms"] // 100) * 100
        df_bw["offset_bin_ms"] = df_bw["offset_bin_ms"].clip(0, 4000)

        df_bw_timeline = (
            df_bw.groupby(["plot_label", "host", "slot", "offset_bin_ms", "message_kind", "flow"], as_index=False)["bytes"]
            .sum()
            .sort_values(["plot_label", "message_kind", "flow", "offset_bin_ms"])
        )
        df_bw_timeline["cumulative_kib"] = (
            df_bw_timeline.groupby(["plot_label", "message_kind", "flow"])["bytes"].cumsum() / 1024
        )

        fig = px.line(
            df_bw_timeline,
            x="offset_bin_ms",
            y="cumulative_kib",
            color="message_kind",
            line_dash="flow",
            facet_col="plot_label",
            facet_col_wrap=1,
            markers=True,
            title="Bandwidth During 3 Random Host Slots",
            labels={
                "offset_bin_ms": "Slot offset (ms)",
                "cumulative_kib": "Cumulative KiB",
                "message_kind": "Message kind",
                "flow": "Direction / delivery",
                "plot_label": "Host / slot",
            },
        )
        fig.update_xaxes(range=[0, 4000])
        fig.update_layout(height=900, legend=dict(orientation="h", y=-0.12))
        fig.show()

import yaml
from pathlib import Path

# ── Identify aggregator hosts from genesis ──────────────────────────
genesis_dir = Path(run_dir) / "genesis"
vc_path = genesis_dir / "validator-config.yaml"
aggregator_hosts: set[str] = set()
non_aggregator_hosts: set[str] = set()

if vc_path.exists():
    with open(vc_path) as f:
        vc = yaml.safe_load(f)
    for v in vc.get("validators", []):
        name: str = v.get("name", "")
        # validator-config uses underscores; bandwidth data uses hyphens
        host = name.replace("_", "-")
        if v.get("isAggregator"):
            aggregator_hosts.add(host)
        else:
            non_aggregator_hosts.add(host)

n_agg = len(aggregator_hosts)
n_non = len(non_aggregator_hosts)
print(f"Aggregator hosts ({n_agg}): {sorted(aggregator_hosts)}")
print(f"Non-aggregator hosts ({n_non}): {sorted(non_aggregator_hosts)[:5]}...")

# ── Compute per-topic totals from bandwidth slots ───────────────────
bw_slots = stats.get("bandwidth", {}).get("slots", [])

# Only gossipsub
gossip_slots = [s for s in bw_slots if s.get("protocol") == "gossip"]

# Exclude status_request/status_response — not topic-level gossip
TOPIC_KINDS = {"aggregation", "attestation", "block"}

# Aggregate: (message_kind, is_aggregator, direction) -> total bytes
from collections import defaultdict

totals: dict[tuple[str, bool, str], int] = defaultdict(int)
for s in gossip_slots:
    kind = s.get("message_kind", "")
    if kind not in TOPIC_KINDS:
        continue
    host = s.get("host", "")
    is_agg = host in aggregator_hosts
    direction = s.get("direction", "?")
    totals[(kind, is_agg, direction)] += s["bytes"]

# ── Build display table ─────────────────────────────────────────────
table_rows = []
for kind in sorted(TOPIC_KINDS):
    agg_in = totals.get((kind, True, "in"), 0)
    agg_out = totals.get((kind, True, "out"), 0)
    non_in = totals.get((kind, False, "in"), 0)
    non_out = totals.get((kind, False, "out"), 0)
    table_rows.append({
        "Topic": kind,
        "Agg In (GiB)": round(agg_in / 1e9, 3),
        "Agg Out (GiB)": round(agg_out / 1e9, 3),
        "Non-Agg In (GiB)": round(non_in / 1e9, 3),
        "Non-Agg Out (GiB)": round(non_out / 1e9, 3),
        "Avg In / agg node (MiB)": round(agg_in / n_agg / 1e6, 2) if n_agg else 0,
        "Avg In / non-agg node (MiB)": round(non_in / n_non / 1e6, 2) if n_non else 0,
    })

df_topics = pd.DataFrame(table_rows)
print("Gossipsub bandwidth breakdown by topic:")
print(df_topics.to_string(index=False))

# ── Average inbound bandwidth per node by topic ─────────────────────
# Total bytes divided by node count in each role — shows how much
# gossip data a typical aggregator vs non‑aggregator receives.
df_per_node = pd.DataFrame([
    {"Topic": kind, "Role": "Aggregator", "Avg inbound per node (MiB)": (
        totals.get((kind, True, "in"), 0) / n_agg / 1e6 if n_agg else 0
    )}
    for kind in sorted(TOPIC_KINDS)
] + [
    {"Topic": kind, "Role": "Non-aggregator", "Avg inbound per node (MiB)": (
        totals.get((kind, False, "in"), 0) / n_non / 1e6 if n_non else 0
    )}
    for kind in sorted(TOPIC_KINDS)
])

fig = px.bar(
    df_per_node,
    x="Topic",
    y="Avg inbound per node (MiB)",
    color="Role",
    barmode="group",
    title="Average Gossipsub Inbound Bandwidth per Node by Topic",
    labels={
        "Avg inbound per node (MiB)": "Avg inbound per node (MiB)",
        "Topic": "Message Kind",
    },
    text_auto=".1f",
)
fig.update_layout(height=450, legend=dict(orientation="h", y=1.12))
fig.show()

# Client distribution pie chart
if node_counts:
    fig = px.pie(
        names=list(node_counts.keys()),
        values=list(node_counts.values()),
        title=f"Client Distribution ({sum(node_counts.values())} nodes)",
        hole=0.4,
    )
    fig.update_traces(textposition="inside", textinfo="percent+label+value")
    fig.show()
else:
    print("No node count data available.")

# Region and bandwidth distribution
from collections import Counter

region_counts = Counter(regions.values()) if regions else Counter()
bw_counts = Counter(bandwidths.values()) if bandwidths else Counter()

fig = make_subplots(
    rows=1, cols=2,
    subplot_titles=("Region Distribution", "Bandwidth Distribution"),
)

if region_counts:
    fig.add_bar(
        x=list(region_counts.keys()),
        y=list(region_counts.values()),
        name="Regions",
        marker_color="steelblue",
        row=1, col=1,
    )

if bw_counts:
    fig.add_bar(
        x=list(bw_counts.keys()),
        y=list(bw_counts.values()),
        name="Bandwidth",
        marker_color="coral",
        row=1, col=2,
    )

fig.update_layout(showlegend=False, height=400)
fig.show()

blocks = stats.get("blocks", {})
block_slots = blocks.get("slots", [])
block_summary = blocks.get("summary", {})

if block_summary.get("warning"):
    print(f"\u26a0\ufe0f {block_summary['warning']}")

if not block_slots:
    print("No block data available.")
else:
    print(f"Blocks published: {block_summary.get('n_published', 0)}")
    print(f"Blocks received: {block_summary.get('n_received', 0)}")
    print(f"Slots with block size data: {block_summary.get('n_with_block_size', 0)}")
    print(f"Slots with block data: {len(block_slots)}")

# Block propagation latency scatter plot
if block_slots:
    rows = []
    for s in block_slots:
        pub_ms = s.get("published_ms")
        if pub_ms is None:
            continue
        for host, recv_ms in s.get("receive_timestamps_ms", {}).items():
            rows.append({
                "slot": s["slot"],
                "host": host,
                "client": host.rsplit("-", 1)[0] if "-" in host else host,
                "latency_ms": round(recv_ms - pub_ms, 1),
            })

    if rows:
        df = pd.DataFrame(rows)
        n_hosts = df["host"].nunique()

        fig = px.scatter(
            df, x="slot", y="latency_ms",
            color="client" if n_hosts > 20 else "host",
            title="Block Propagation Latency (publish \u2192 receive)",
            labels={"latency_ms": "Latency (ms)", "slot": "Slot", "client": "Client"},
        )
        fig.update_layout(
            height=500,
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=-0.35 if n_hosts <= 20 else -0.2,
                xanchor="center",
                x=0.5,
                font=dict(size=9),
            ),
            margin=dict(b=80),
        )
        fig.show()
    else:
        print("No publish\u2192receive latency data available.")

# Block propagation latency percentiles per slot
if "df" in globals() and not df.empty:
    df_pct = (
        df.groupby("slot")["latency_ms"]
        .quantile([0.50, 0.95, 0.99])
        .unstack()
        .rename(columns={0.50: "p50", 0.95: "p95", 0.99: "p99"})
        .reset_index()
    )

    fig = go.Figure()
    for col in ["p50", "p95", "p99"]:
        fig.add_scatter(
            x=df_pct["slot"],
            y=df_pct[col],
            mode="lines+markers",
            name=col,
        )

    fig.update_layout(
        title="Block Propagation Latency Percentiles per Slot",
        xaxis_title="Slot",
        yaxis_title="Latency (ms)",
        height=450,
    )
    fig.show()
else:
    print("No publish\u2192receive latency data available.")

# Block size per slot
block_size_rows = [
    {"slot": s["slot"], "block_size_kib": s["block_size_bytes"] / 1024}
    for s in block_slots
    if "block_size_bytes" in s
]

if block_size_rows:
    df_block_sizes = pd.DataFrame(block_size_rows)
    fig = px.line(
        df_block_sizes,
        x="slot",
        y="block_size_kib",
        markers=True,
        title="Block Size per Slot",
        labels={"slot": "Slot", "block_size_kib": "Block size (KiB)"},
    )
    fig.update_layout(height=400)
    fig.show()
else:
    print("No block size data available.")

attestations = stats.get("attestations", {})
coverage = attestations.get("coverage", {})
cov_slots = coverage.get("slots", [])
cov_summary = coverage.get("summary", {})

if cov_summary.get("warning"):
    print(f"\u26a0\ufe0f {cov_summary['warning']}")

if not cov_slots:
    print("No attestation coverage data available.")
else:
    print(f"Slots with coverage data: {cov_summary.get('slots_with_data', 0)}")
    for field in ["median_slot_p50_nodes_to_95_attestations_ms",
                  "median_slot_p90_nodes_to_95_attestations_ms",
                  "median_slot_p95_nodes_to_95_attestations_ms"]:
        if field in cov_summary:
            label = field.replace("median_slot_", "").replace("_nodes_to_95_attestations_ms", "")
            print(f"  Median {label}: {cov_summary[field]} ms")

# Attestation latency per slot
if cov_slots:
    df_cov = pd.DataFrame(cov_slots)

    fig = go.Figure()
    if "p95_nodes_to_95_attestations_ms" in df_cov.columns:
        fig.add_scatter(
            x=df_cov["slot"], y=df_cov["p95_nodes_to_95_attestations_ms"],
            mode="lines+markers", name="p95 latency",
            line=dict(color="coral", dash="dot"),
        )

    fig.update_yaxes(title_text="Latency (ms)")
    fig.update_xaxes(title_text="Slot")
    fig.update_layout(title="Attestation Latency per Slot", height=450)
    fig.show()

# Attestation propagation CDFs for 3 random validators at their assigned slots
validator_propagation = attestations.get("validator_propagation", [])

if not validator_propagation:
    print("No validator attestation propagation data available.")
else:
    rng = random.Random(str(run_id))

    # Build (validator_id, slot) pairs
    pairs = [(vp["validator_id"], vp["slot"]) for vp in validator_propagation]
    all_validators = sorted(set(v for v, s in pairs))

    n_select = min(3, len(all_validators))
    selected_validators = sorted(rng.sample(all_validators, n_select))
    print(f"Selected validators: {selected_validators}")

    for validator_id in selected_validators:
        validator_slots = sorted([s for v, s in pairs if v == validator_id])
        if not validator_slots:
            continue
        selected_slot = rng.choice(validator_slots)
        vp = next(
            (x for x in validator_propagation
             if x["validator_id"] == validator_id and x["slot"] == selected_slot),
            None,
        )
        if vp is None:
            continue

        propagation_times = vp["propagation_times_ms"]
        n = len(propagation_times)
        y = [(i + 1) / n for i in range(n)]

        fig = go.Figure()
        fig.add_scatter(
            x=propagation_times,
            y=y,
            mode="lines",
            line_shape="hv",
            name=f"Validator {validator_id}, Slot {selected_slot}",
        )
        fig.update_layout(
            title=f"Attestation Propagation CDF — Validator {validator_id}, Slot {selected_slot}",
            xaxis_title="Propagation time (ms)",
            yaxis_title="CDF",
            yaxis=dict(range=[0, 1]),
            height=450,
            showlegend=False,
        )
        fig.show()

# Aggregated attestation propagation CDFs for 3 random slots
aggregation_propagation = attestations.get("aggregation_propagation", [])

if not aggregation_propagation:
    print("No aggregated attestation propagation data available.")
else:
    rng = random.Random(str(run_id) + "-aggregation")

    # Group by slot
    by_slot: dict[int, list[dict]] = {}
    for ap in aggregation_propagation:
        slot = ap["slot"]
        by_slot.setdefault(slot, []).append(ap)

    available_slots = sorted(by_slot.keys())
    n_select = min(3, len(available_slots))
    selected_slots = sorted(rng.sample(available_slots, n_select))
    print(f"Selected slots: {selected_slots}")

    for slot in selected_slots:
        messages = by_slot[slot]
        selected_message = rng.choice(messages)
        message_id = selected_message["message_id"]
        propagation_times = selected_message["propagation_times_ms"]
        n = len(propagation_times)
        y = [(i + 1) / n for i in range(n)]

        fig = go.Figure()
        fig.add_scatter(
            x=propagation_times,
            y=y,
            mode="lines",
            line_shape="hv",
            name=f"Slot {slot}, {message_id[:16]}...",
        )
        fig.update_layout(
            title=f"Aggregated Attestation Propagation CDF — Slot {slot}",
            xaxis_title="Propagation time (ms)",
            yaxis_title="CDF",
            yaxis=dict(range=[0, 1]),
            height=450,
            showlegend=False,
        )
        fig.show()

chain = stats.get("chain_status", {})
chain_slots = chain.get("slots", [])
chain_summary = chain.get("summary", {})

if chain_summary.get("warning"):
    print(f"\u26a0\ufe0f {chain_summary['warning']}")

if not chain_slots:
    print("No chain status data available.")
else:
    print(f"Slots with chain data: {chain_summary.get('slots_with_data', 0)}")
    print(f"Hosts with chain data: {chain_summary.get('hosts_with_data', 0)}")

# Chain finality heatmaps
if chain_slots:
    def host_sort_key(host):
        prefix, sep, suffix = host.rpartition("-")
        if sep and suffix.isdigit():
            return (prefix, int(suffix))
        return (host, -1)

    all_hosts = sorted({h for s in chain_slots for h in s.get("hosts", {})}, key=host_sort_key)
    n_hosts = len(all_hosts)

    for metric, key, colorscale, title in [
        ("head", "head_slot", "Blues", "Chain Head Slot per Node"),
        ("finalized", "latest_finalized_slot", "Greens", "Finalized Slot per Node"),
    ]:
        rows = []
        for s in chain_slots:
            slot = s["slot"]
            for host in all_hosts:
                hs = s.get("hosts", {}).get(host, {})
                rows.append({"slot": slot, "host": host, metric: hs.get(key, 0)})

        df_m = pd.DataFrame(rows)

        pivot = df_m.pivot(index="host", columns="slot", values=metric).reindex(all_hosts)
        y_title = "Node"
        tick_size = 9 if n_hosts <= 16 else 6
        row_height = 25 if n_hosts <= 30 else 8

        fig = go.Figure(data=go.Heatmap(
            z=pivot.values.tolist(),
            x=[str(c) for c in pivot.columns],
            y=list(pivot.index),
            colorscale=colorscale,
        ))
        fig.update_layout(
            title=title,
            xaxis_title="Slot", yaxis_title=y_title,
            height=max(400, len(pivot.index) * row_height),
            margin=dict(l=max(80, max(len(str(y)) for y in pivot.index) * 7)),
            yaxis=dict(tickfont=dict(size=tick_size)),
        )
        fig.show()

topo_path = data.get("topology_gml_path")

if topo_path and topo_path.exists():
    try:
        import networkx as nx
        G = nx.read_gml(topo_path)
        print(f"Topology: {G.number_of_nodes()} nodes, {G.number_of_edges()} edges")

        pos = nx.spring_layout(G, seed=42)

        edge_x = []
        edge_y = []
        for u, v in G.edges():
            x0, y0 = pos[u]
            x1, y1 = pos[v]
            edge_x.extend([x0, x1, None])
            edge_y.extend([y0, y1, None])

        node_x = [pos[n][0] for n in G.nodes()]
        node_y = [pos[n][1] for n in G.nodes()]

        node_labels = [str(n) for n in G.nodes()]
        node_regions = []
        for n in G.nodes():
            label = str(n)
            if label in regions:
                node_regions.append(regions[label])
            else:
                node_regions.append("unknown")

        region_labels = sorted(set(node_regions))
        region_ids = {region: i for i, region in enumerate(region_labels)}
        node_colors = [region_ids[region] for region in node_regions]

        fig = go.Figure()
        fig.add_scatter(x=edge_x, y=edge_y, mode="lines",
                        line=dict(width=0.5, color="#888"), hoverinfo="none")
        fig.add_scatter(
            x=node_x, y=node_y, mode="markers",
            marker=dict(
                size=8,
                color=node_colors,
                colorscale="Viridis",
                showscale=True,
                colorbar=dict(tickmode="array", tickvals=list(region_ids.values()), ticktext=region_labels),
            ),
            text=[f"{l} ({r})" for l, r in zip(node_labels, node_regions)],
            hoverinfo="text",
        )
        fig.update_layout(
            title="Underlay Network Topology",
            showlegend=False, height=600,
            xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
            yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
        )
        fig.show()
    except ImportError:
        print("networkx not installed. Install with: pip install networkx")
else:
    print("No topology.gml file found.")

# Warnings
warnings = stats.get("warnings", [])
if warnings:
    print(f"\u26a0\ufe0f {len(warnings)} warnings:")
    for w in warnings:
        print(f"  - {w}")
else:
    print("\u2705 No warnings.")