import json
from pathlib import Path

import pandas as pd
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from IPython.display import display

# Set default renderer for static HTML output
import plotly.io as pio
pio.renderers.default = "notebook"

# Resolve devnet_id
DATA_DIR = Path("../data")

if devnet_id is None:
    devnets_path = DATA_DIR / "devnets.json"
    if devnets_path.exists():
        with open(devnets_path) as f:
            devnets = json.load(f).get("devnets", [])
        if devnets:
            devnet_id = devnets[-1]["id"]
            print(f"Using latest devnet: {devnet_id}")
    else:
        raise ValueError("No devnets.json found. Run 'just detect-devnets' first.")

DEVNET_DIR = DATA_DIR / devnet_id
print(f"Loading data from: {DEVNET_DIR}")

# Load devnet metadata
with open(DATA_DIR / "devnets.json") as f:
    devnets_data = json.load(f)
    devnet_info = next((d for d in devnets_data["devnets"] if d["id"] == devnet_id), None)

if devnet_info:
    print(f"Devnet: {devnet_info['id']}")
    print(f"Duration: {devnet_info['duration_hours']:.1f} hours")
    print(f"Time: {devnet_info['start_time']} to {devnet_info['end_time']}")
    print(f"Slots: {devnet_info['start_slot']} \u2192 {devnet_info['end_slot']}")
    print(f"Clients: {', '.join(devnet_info['clients'])}")

# Unified client list from devnet metadata (includes all containers via cAdvisor)
all_clients = sorted(devnet_info["clients"])
n_cols = min(len(all_clients), 2)
n_rows = -(-len(all_clients) // n_cols)

# Load head slot data
head_df = pd.read_parquet(DEVNET_DIR / "head_slot.parquet")
head_df = head_df.groupby(["client", "metric", "timestamp"], as_index=False)["value"].max()
print(f"Head slot: {len(head_df)} records, clients: {sorted(head_df['client'].unique())}")

# Load finality data
finality_df = pd.read_parquet(DEVNET_DIR / "finality_metrics.parquet")
finality_df = finality_df.groupby(["client", "metric", "timestamp"], as_index=False)["value"].max()
print(f"Finality: {len(finality_df)} records, clients: {sorted(finality_df['client'].unique())}")
print(f"Finality metrics: {sorted(finality_df['metric'].unique())}")

# Load fork choice reorgs
reorgs_df = pd.read_parquet(DEVNET_DIR / "fork_choice_reorgs.parquet")
reorgs_df = reorgs_df.groupby(["client", "timestamp"], as_index=False)["value"].max()
print(f"Reorgs: {len(reorgs_df)} records, clients: {sorted(reorgs_df['client'].unique())}")

print(f"\nAll clients ({len(all_clients)}): {all_clients}")

# Assign a distinct color per client for overlaid charts
import plotly.express as px
client_colors = {client: px.colors.qualitative.Plotly[i % len(px.colors.qualitative.Plotly)] for i, client in enumerate(all_clients)}

fig = make_subplots(
    rows=3, cols=1,
    subplot_titles=["Head Slot Progression", "Current-to-Head Distance", "Finalized Slot Progression"],
    vertical_spacing=0.08,
    shared_xaxes=True,
)

# 1. Head slot progression - all clients overlaid
head_slot_df = head_df[head_df["metric"] == "lean_head_slot"]
for client in all_clients:
    cdf = head_slot_df[head_slot_df["client"] == client].sort_values("timestamp")
    if not cdf.empty:
        fig.add_trace(
            go.Scatter(
                x=cdf["timestamp"], y=cdf["value"],
                name=client, legendgroup=client,
                line=dict(color=client_colors[client]),
            ),
            row=1, col=1,
        )

# Add current_slot as a dashed reference line (any client, they're all the same)
current_slot_df = head_df[head_df["metric"] == "lean_current_slot"]
ref_client = all_clients[0]
ref = current_slot_df[current_slot_df["client"] == ref_client].sort_values("timestamp")
if not ref.empty:
    fig.add_trace(
        go.Scatter(
            x=ref["timestamp"], y=ref["value"],
            name="current_slot", legendgroup="current_slot",
            line=dict(color="#ccc", dash="dot", width=1),
        ),
        row=1, col=1,
    )
fig.update_yaxes(title_text="Slot", row=1, col=1)

# 2. Current-to-head distance - all clients overlaid
current = head_df[head_df["metric"] == "lean_current_slot"][["client", "timestamp", "value"]].rename(columns={"value": "current_slot"})
head_only = head_df[head_df["metric"] == "lean_head_slot"][["client", "timestamp", "value"]].rename(columns={"value": "head_slot"})
overview_lag = current.merge(head_only, on=["client", "timestamp"], how="inner")
overview_lag["lag"] = overview_lag["current_slot"] - overview_lag["head_slot"]

for client in all_clients:
    cdf = overview_lag[overview_lag["client"] == client].sort_values("timestamp")
    if not cdf.empty:
        fig.add_trace(
            go.Scatter(
                x=cdf["timestamp"], y=cdf["lag"],
                name=client, legendgroup=client,
                showlegend=False,
                line=dict(color=client_colors[client]),
            ),
            row=2, col=1,
        )
fig.update_yaxes(title_text="Slots behind", row=2, col=1)

# 3. Finalized slot progression - all clients overlaid
fin_slot_df = finality_df[finality_df["metric"] == "lean_latest_finalized_slot"]
for client in all_clients:
    cdf = fin_slot_df[fin_slot_df["client"] == client].sort_values("timestamp")
    if not cdf.empty:
        fig.add_trace(
            go.Scatter(
                x=cdf["timestamp"], y=cdf["value"],
                name=client, legendgroup=client,
                showlegend=False,
                line=dict(color=client_colors[client]),
            ),
            row=3, col=1,
        )
fig.update_yaxes(title_text="Slot", row=3, col=1)

fig.update_layout(
    height=800,
    legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="left", x=0),
)
fig.show()

fig = make_subplots(
    rows=n_rows, cols=n_cols,
    subplot_titles=all_clients,
    vertical_spacing=0.12 / max(n_rows - 1, 1) * 2,
    horizontal_spacing=0.08,
)

colors = {"lean_head_slot": "#636EFA", "lean_current_slot": "#EF553B"}
labels = {"lean_head_slot": "head_slot", "lean_current_slot": "current_slot"}
legend_added = set()

for i, client in enumerate(all_clients):
    row = i // n_cols + 1
    col = i % n_cols + 1
    cdf = head_df[head_df["client"] == client]
    for metric in ["lean_current_slot", "lean_head_slot"]:
        mdf = cdf[cdf["metric"] == metric].sort_values("timestamp")
        if mdf.empty:
            continue
        label = labels[metric]
        show_legend = metric not in legend_added
        legend_added.add(metric)
        fig.add_trace(
            go.Scatter(
                x=mdf["timestamp"], y=mdf["value"],
                name=label, legendgroup=metric,
                showlegend=show_legend,
                line=dict(color=colors[metric]),
            ),
            row=row, col=col,
        )
    fig.update_yaxes(title_text="Slot", row=row, col=col)

fig.update_layout(
    title="Head Slot vs Current Slot by Client",
    height=270 * n_rows,
)
fig.show()

jf_metrics = ["lean_latest_justified_slot", "lean_latest_finalized_slot"]
jf_df = finality_df[finality_df["metric"].isin(jf_metrics)].copy()

head_only_all = head_df[head_df["metric"] == "lean_head_slot"].copy()
head_only_all["metric"] = "lean_head_slot"

combined = pd.concat([jf_df, head_only_all], ignore_index=True)

fig = make_subplots(
    rows=n_rows, cols=n_cols,
    subplot_titles=all_clients,
    vertical_spacing=0.12 / max(n_rows - 1, 1) * 2,
    horizontal_spacing=0.08,
)

colors = {
    "lean_head_slot": "#636EFA",
    "lean_latest_justified_slot": "#00CC96",
    "lean_latest_finalized_slot": "#EF553B",
}
labels = {
    "lean_head_slot": "head",
    "lean_latest_justified_slot": "justified",
    "lean_latest_finalized_slot": "finalized",
}
legend_added = set()

for i, client in enumerate(all_clients):
    row = i // n_cols + 1
    col = i % n_cols + 1
    cdf = combined[combined["client"] == client]
    for metric in ["lean_head_slot", "lean_latest_justified_slot", "lean_latest_finalized_slot"]:
        mdf = cdf[cdf["metric"] == metric].sort_values("timestamp")
        if mdf.empty:
            continue
        show_legend = metric not in legend_added
        legend_added.add(metric)
        fig.add_trace(
            go.Scatter(
                x=mdf["timestamp"], y=mdf["value"],
                name=labels[metric], legendgroup=metric,
                showlegend=show_legend,
                line=dict(color=colors[metric]),
            ),
            row=row, col=col,
        )
    fig.update_yaxes(title_text="Slot", row=row, col=col)

fig.update_layout(
    title="Head, Justified & Finalized Slot by Client",
    height=270 * n_rows,
)
fig.show()

current_df = head_df[head_df["metric"] == "lean_current_slot"][["client", "timestamp", "value"]].rename(columns={"value": "current_slot"})
head_only = head_df[head_df["metric"] == "lean_head_slot"][["client", "timestamp", "value"]].rename(columns={"value": "head_slot"})
lag_df = current_df.merge(head_only, on=["client", "timestamp"], how="inner")
lag_df["lag"] = lag_df["current_slot"] - lag_df["head_slot"]

fig = make_subplots(
    rows=n_rows, cols=n_cols,
    subplot_titles=all_clients,
    vertical_spacing=0.12 / max(n_rows - 1, 1) * 2,
    horizontal_spacing=0.08,
)

for i, client in enumerate(all_clients):
    row = i // n_cols + 1
    col = i % n_cols + 1
    cdf = lag_df[lag_df["client"] == client].sort_values("timestamp")
    if not cdf.empty:
        fig.add_trace(
            go.Scatter(
                x=cdf["timestamp"], y=cdf["lag"],
                name=client, showlegend=False,
                line=dict(color="#636EFA"),
            ),
            row=row, col=col,
        )
    else:
        fig.add_trace(
            go.Scatter(x=[None], y=[None], showlegend=False, hoverinfo='skip'),
            row=row, col=col,
        )
        _n = (row - 1) * n_cols + col
        _s = "" if _n == 1 else str(_n)
        fig.add_annotation(
            text="No data available",
            xref=f"x{_s} domain", yref=f"y{_s} domain",
            x=0.5, y=0.5,
            showarrow=False,
            font=dict(size=12, color="#999"),
        )
    fig.update_yaxes(title_text="Slots behind", row=row, col=col)

fig.update_layout(
    title="Current-to-Head Distance (current_slot - head_slot)",
    height=270 * n_rows,
)
fig.show()

head_ts = head_df[head_df["metric"] == "lean_head_slot"][["client", "timestamp", "value"]].rename(columns={"value": "head_slot"})
just_ts = finality_df[finality_df["metric"] == "lean_latest_justified_slot"][["client", "timestamp", "value"]].rename(columns={"value": "justified_slot"})
justification_lag = head_ts.merge(just_ts, on=["client", "timestamp"], how="inner")
justification_lag["lag"] = justification_lag["head_slot"] - justification_lag["justified_slot"]

fig = make_subplots(
    rows=n_rows, cols=n_cols,
    subplot_titles=all_clients,
    vertical_spacing=0.12 / max(n_rows - 1, 1) * 2,
    horizontal_spacing=0.08,
)

for i, client in enumerate(all_clients):
    row = i // n_cols + 1
    col = i % n_cols + 1
    cdf = justification_lag[justification_lag["client"] == client].sort_values("timestamp")
    if not cdf.empty:
        fig.add_trace(
            go.Scatter(
                x=cdf["timestamp"], y=cdf["lag"],
                name=client, showlegend=False,
                line=dict(color="#636EFA"),
            ),
            row=row, col=col,
        )
    else:
        fig.add_trace(
            go.Scatter(x=[None], y=[None], showlegend=False, hoverinfo='skip'),
            row=row, col=col,
        )
        _n = (row - 1) * n_cols + col
        _s = "" if _n == 1 else str(_n)
        fig.add_annotation(
            text="No data available",
            xref=f"x{_s} domain", yref=f"y{_s} domain",
            x=0.5, y=0.5,
            showarrow=False,
            font=dict(size=12, color="#999"),
        )
    fig.update_yaxes(title_text="Slots", row=row, col=col)

fig.update_layout(
    title="Head-to-Justified Distance (head_slot - justified_slot)",
    height=270 * n_rows,
)
fig.show()

fin_ts = finality_df[finality_df["metric"] == "lean_latest_finalized_slot"][["client", "timestamp", "value"]].rename(columns={"value": "finalized_slot"})
finality_lag = just_ts.merge(fin_ts, on=["client", "timestamp"], how="inner")
finality_lag["lag"] = finality_lag["justified_slot"] - finality_lag["finalized_slot"]

fig = make_subplots(
    rows=n_rows, cols=n_cols,
    subplot_titles=all_clients,
    vertical_spacing=0.12 / max(n_rows - 1, 1) * 2,
    horizontal_spacing=0.08,
)

for i, client in enumerate(all_clients):
    row = i // n_cols + 1
    col = i % n_cols + 1
    cdf = finality_lag[finality_lag["client"] == client].sort_values("timestamp")
    if not cdf.empty:
        fig.add_trace(
            go.Scatter(
                x=cdf["timestamp"], y=cdf["lag"],
                name=client, showlegend=False,
                line=dict(color="#636EFA"),
            ),
            row=row, col=col,
        )
    else:
        fig.add_trace(
            go.Scatter(x=[None], y=[None], showlegend=False, hoverinfo='skip'),
            row=row, col=col,
        )
        _n = (row - 1) * n_cols + col
        _s = "" if _n == 1 else str(_n)
        fig.add_annotation(
            text="No data available",
            xref=f"x{_s} domain", yref=f"y{_s} domain",
            x=0.5, y=0.5,
            showarrow=False,
            font=dict(size=12, color="#999"),
        )
    fig.update_yaxes(title_text="Slots", row=row, col=col)

fig.update_layout(
    title="Justified-to-Finalized Distance (justified_slot - finalized_slot)",
    height=270 * n_rows,
)
fig.show()

fig = make_subplots(
    rows=n_rows, cols=n_cols,
    subplot_titles=all_clients,
    vertical_spacing=0.12 / max(n_rows - 1, 1) * 2,
    horizontal_spacing=0.08,
)

for i, client in enumerate(all_clients):
    row = i // n_cols + 1
    col = i % n_cols + 1
    cdf = reorgs_df[reorgs_df["client"] == client].sort_values("timestamp")
    if not cdf.empty:
        fig.add_trace(
            go.Scatter(
                x=cdf["timestamp"], y=cdf["value"],
                name=client, showlegend=False,
                line=dict(color="#636EFA"),
            ),
            row=row, col=col,
        )
    else:
        fig.add_trace(
            go.Scatter(x=[None], y=[None], showlegend=False, hoverinfo='skip'),
            row=row, col=col,
        )
        _n = (row - 1) * n_cols + col
        _s = "" if _n == 1 else str(_n)
        fig.add_annotation(
            text="No data available",
            xref=f"x{_s} domain", yref=f"y{_s} domain",
            x=0.5, y=0.5,
            showarrow=False,
            font=dict(size=12, color="#999"),
        )
    fig.update_yaxes(title_text="Cumulative reorgs", row=row, col=col)

fig.update_layout(
    title="Fork Choice Reorgs by Client",
    height=270 * n_rows,
)
fig.show()

summary_rows = []

for client in all_clients:
    row = {"Client": client}

    # Current-to-head distance
    client_lag = lag_df[lag_df["client"] == client]["lag"]
    if not client_lag.empty:
        row["Avg C-H Dist."] = f"{client_lag.mean():.1f}"
        row["Max C-H Dist."] = f"{client_lag.max():.0f}"

    # Head-to-justified distance
    client_just = justification_lag[justification_lag["client"] == client]["lag"]
    if not client_just.empty:
        row["Avg H-J Dist."] = f"{client_just.mean():.1f}"
        row["Max H-J Dist."] = f"{client_just.max():.0f}"

    # Justified-to-finalized distance
    client_fin = finality_lag[finality_lag["client"] == client]["lag"]
    if not client_fin.empty:
        row["Avg J-F Dist."] = f"{client_fin.mean():.1f}"
        row["Max J-F Dist."] = f"{client_fin.max():.0f}"

    # Reorgs
    client_reorgs = reorgs_df[reorgs_df["client"] == client]["value"]
    if not client_reorgs.empty:
        row["Reorgs"] = f"{client_reorgs.max():.0f}"

    # Final head slot
    client_head = head_df[(head_df["client"] == client) & (head_df["metric"] == "lean_head_slot")]
    if not client_head.empty:
        row["Final Head Slot"] = f"{client_head['value'].max():.0f}"

    # Final finalized slot
    client_finalized = finality_df[(finality_df["client"] == client) & (finality_df["metric"] == "lean_latest_finalized_slot")]
    if not client_finalized.empty:
        row["Final Finalized Slot"] = f"{client_finalized['value'].max():.0f}"

    summary_rows.append(row)

if summary_rows:
    summary_df = pd.DataFrame(summary_rows).set_index("Client").fillna("-")
    display(summary_df)

print(f"\nDevnet: {devnet_id}")
if devnet_info:
    print(f"Duration: {devnet_info['duration_hours']:.1f} hours")