from __future__ import annotations

import math
from typing import Any

import networkx as nx
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go


def build_timeseries_fig(events: list[dict[str, Any]]) -> go.Figure:
    df = pd.DataFrame(events)
    if df.empty:
        return go.Figure()
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=df["step"], y=df["cumulative_tokens"], mode="lines+markers", name="Cumulative Tokens"))
    fig.add_trace(go.Scatter(x=df["step"], y=df["hot_tokens"], mode="lines+markers", name="Hot Layer Tokens"))
    fig.update_layout(title="NEXUS Evolution Over Time", xaxis_title="Step", yaxis_title="Tokens")
    return fig


def build_radar_fig(latest: dict[str, Any]) -> go.Figure:
    categories = ["Tokens", "Tool Calls", "Warm Memory", "Cold Memory", "Velocity"]
    values = [
        min(100, latest.get("cumulative_tokens", 0) / 150),
        min(100, latest.get("tool_calls", 0) * 20),
        min(100, latest.get("warm_entries", 0) * 10),
        min(100, latest.get("cold_count", 0) * 10),
        min(100, 1000 / max(1.0, latest.get("duration_ms", 1.0))),
    ]
    fig = go.Figure(
        data=go.Scatterpolar(r=values, theta=categories, fill="toself", name="Runtime Profile")
    )
    fig.update_layout(title="Runtime Pillar Utilization", polar=dict(radialaxis=dict(visible=True, range=[0, 100])))
    return fig


def build_heatmap_fig(events: list[dict[str, Any]]) -> go.Figure:
    df = pd.DataFrame(events)
    if df.empty:
        return go.Figure()
    z = [
        df["cumulative_tokens"].tolist(),
        df["hot_tokens"].tolist(),
        df["warm_entries"].tolist(),
        df["tool_calls"].tolist(),
    ]
    fig = go.Figure(
        data=go.Heatmap(
            z=z,
            y=["cumulative_tokens", "hot_tokens", "warm_entries", "tool_calls"],
            x=df["step"].tolist(),
            colorscale="Viridis",
        )
    )
    fig.update_layout(title="State/Economy Heatmap", xaxis_title="Step")
    return fig


def build_timeline_fig(events: list[dict[str, Any]]) -> go.Figure:
    df = pd.DataFrame(events)
    if df.empty:
        return go.Figure()
    fig = px.bar(df, x="duration_ms", y="stage", color="agent_id", orientation="h", title="Node Duration Timeline")
    return fig


def build_graph_fig(events: list[dict[str, Any]]) -> go.Figure:
    g = nx.DiGraph()
    g.add_edges_from([("plan", "research"), ("research", "synthesize")])
    if events and events[-1].get("stage") == "done":
        g.add_node("done")
        g.add_edge("synthesize", "done")

    pos = nx.spring_layout(g, seed=7)
    edge_x = []
    edge_y = []
    for src, dst in g.edges():
        x0, y0 = pos[src]
        x1, y1 = pos[dst]
        edge_x += [x0, x1, None]
        edge_y += [y0, y1, None]

    node_x = []
    node_y = []
    node_text = []
    latest_stage = events[-1].get("stage") if events else ""
    for n in g.nodes():
        x, y = pos[n]
        node_x.append(x)
        node_y.append(y)
        node_text.append(f"{n}{' (active)' if n == latest_stage else ''}")

    fig = go.Figure()
    fig.add_trace(go.Scatter(x=edge_x, y=edge_y, mode="lines", line=dict(width=1.5, color="#7f8c8d"), hoverinfo="none"))
    fig.add_trace(
        go.Scatter(
            x=node_x,
            y=node_y,
            mode="markers+text",
            marker=dict(size=26, color="#1f77b4"),
            text=node_text,
            textposition="bottom center",
            hoverinfo="text",
        )
    )
    fig.update_layout(title="Execution Graph State", showlegend=False, xaxis=dict(visible=False), yaxis=dict(visible=False))
    return fig


def build_sunburst_fig(latest: dict[str, Any]) -> go.Figure:
    total = max(1, latest.get("cumulative_tokens", 1))
    tool_weight = max(1, latest.get("tool_calls", 0) * 20)
    hot_weight = max(1, latest.get("hot_tokens", 0))
    synth_weight = max(1, total - tool_weight)
    labels = ["Runtime", "LLM", "Tools", "Hot Layer", "Synthesis"]
    parents = ["", "Runtime", "Runtime", "Runtime", "LLM"]
    values = [total + tool_weight + hot_weight, total, tool_weight, hot_weight, synth_weight]
    fig = go.Figure(go.Sunburst(labels=labels, parents=parents, values=values, branchvalues="total"))
    fig.update_layout(title="Token and Workload Composition")
    return fig


def build_ledger_area_fig(events: list[dict[str, Any]]) -> go.Figure:
    df = pd.DataFrame(events)
    if df.empty:
        return go.Figure()
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=df["step"], y=df.get("planner_used", 0), stackgroup="one", name="planner"))
    fig.add_trace(go.Scatter(x=df["step"], y=df.get("researcher_used", 0), stackgroup="one", name="researcher"))
    fig.add_trace(go.Scatter(x=df["step"], y=df.get("synthesizer_used", 0), stackgroup="one", name="synthesizer"))
    fig.update_layout(title="TEE Ledger Consumption by Agent", xaxis_title="Step", yaxis_title="Tokens")
    return fig


def build_graph_progress_fig(events: list[dict[str, Any]]) -> go.Figure:
    df = pd.DataFrame(events)
    if df.empty:
        return go.Figure()
    fig = go.Figure()
    fig.add_trace(go.Bar(x=df["step"], y=df.get("completed_nodes", 0), name="Completed Nodes"))
    fig.add_trace(go.Scatter(x=df["step"], y=df.get("ready_nodes", 0), mode="lines+markers", name="Ready Nodes"))
    fig.add_trace(go.Scatter(x=df["step"], y=df.get("mutation_count", 0), mode="lines+markers", name="Mutations"))
    fig.update_layout(title="DEG Progress and Runtime Mutation", xaxis_title="Step", yaxis_title="Count")
    return fig
