"""Generate a self-contained HTML report with an SVG timeline from tracer events."""

from __future__ import annotations

import html
import os
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from .trace import Tracer

# ------------------------------------------------------------------
# colours
# ------------------------------------------------------------------

COLORS = {
    "push": "#3b82f6",
    "exec": "#22c55e",
    "pull": "#f97316",
    "aggregate": "#8b5cf6",
    "dead": "#ef4444",
}

BG = "#0f172a"
CARD = "#1e293b"
TEXT = "#e2e8f0"
MUTED = "#94a3b8"
GRID = "#334155"


def _fmt_bytes(n: int) -> str:
    if n < 1024:
        return f"{n}B"
    if n < 1024 * 1024:
        return f"{n / 1024:.1f}KB"
    return f"{n / (1024 * 1024):.1f}MB"


# ------------------------------------------------------------------
# SVG builder
# ------------------------------------------------------------------

def _build_svg(tracer: Tracer) -> str:
    rounds = tracer._collect_rounds()
    max_rank = tracer._max_rank()
    if not rounds:
        return '<svg xmlns="http://www.w3.org/2000/svg"></svg>'

    # layout constants
    label_w = 100
    right_pad = 20
    row_h = 28
    row_gap = 4
    round_gap = 24
    round_header_h = 22
    chart_w = 700
    total_w = label_w + chart_w + right_pad

    # compute total height
    rows_per_round = 1 + max_rank + 1  # scheduler + workers
    n_rounds = len(rounds)
    total_h = (
        n_rounds * (round_header_h + rows_per_round * (row_h + row_gap) + round_gap)
        + 40  # bottom time axis
    )

    parts: list[str] = []
    parts.append(
        f'<svg xmlns="http://www.w3.org/2000/svg" '
        f'viewBox="0 0 {total_w} {total_h}" '
        f'width="{total_w}" height="{total_h}">'
    )
    parts.append(
        f'<rect width="{total_w}" height="{total_h}" fill="{CARD}" rx="8"/>'
    )

    y_cursor = 12

    for rnum in sorted(rounds.keys()):
        rd = rounds[rnum]
        t0 = rd["start"]
        dur = rd["duration"] or 0.001
        surv = rd["survivors"]
        total_nodes = rd["total"]

        def x_pos(t: float) -> float:
            frac = max(0.0, min(1.0, (t - t0) / dur))
            return label_w + frac * chart_w

        def bar_w(d: float) -> float:
            return max(3, d / dur * chart_w)  # min 3px so it's visible

        # round header
        parts.append(
            f'<text x="12" y="{y_cursor + 14}" '
            f'font-size="13" font-family="monospace" fill="{TEXT}" font-weight="bold">'
            f"round {rnum}  [{surv}/{total_nodes}]  {dur:.2f}s</text>"
        )
        y_cursor += round_header_h

        # time gridlines
        n_ticks = 5
        for i in range(n_ticks + 1):
            frac = i / n_ticks
            gx = label_w + frac * chart_w
            parts.append(
                f'<line x1="{gx}" y1="{y_cursor}" '
                f'x2="{gx}" y2="{y_cursor + rows_per_round * (row_h + row_gap)}" '
                f'stroke="{GRID}" stroke-width="0.5" stroke-dasharray="4 4"/>'
            )
            t_label = frac * dur
            parts.append(
                f'<text x="{gx}" y="{y_cursor + rows_per_round * (row_h + row_gap) + 12}" '
                f'font-size="10" font-family="monospace" fill="{MUTED}" text-anchor="middle">'
                f"{t_label:.1f}s</text>"
            )

        # --- scheduler row ---
        row_y = y_cursor
        parts.append(
            f'<text x="12" y="{row_y + 18}" '
            f'font-size="11" font-family="monospace" fill="{MUTED}">scheduler</text>'
        )
        # background bar
        parts.append(
            f'<rect x="{label_w}" y="{row_y + 4}" width="{chart_w}" height="{row_h - 8}" '
            f'fill="{BG}" rx="3"/>'
        )
        for ev in rd["events"]:
            d = ev.data.get("duration_s", 0)
            if ev.kind in ("push", "pull", "aggregate"):
                color = COLORS.get(ev.kind, COLORS["push"])
                bx = x_pos(ev.time - d)
                bw = bar_w(d)
                tooltip = _tooltip(ev)
                parts.append(
                    f'<rect x="{bx}" y="{row_y + 6}" width="{bw}" height="{row_h - 12}" '
                    f'fill="{color}" rx="2" opacity="0.9">'
                    f"<title>{html.escape(tooltip)}</title></rect>"
                )
        y_cursor += row_h + row_gap

        # --- worker rows ---
        for rank in range(max_rank + 1):
            row_y = y_cursor
            parts.append(
                f'<text x="12" y="{row_y + 18}" '
                f'font-size="11" font-family="monospace" fill="{MUTED}">worker {rank}</text>'
            )
            parts.append(
                f'<rect x="{label_w}" y="{row_y + 4}" width="{chart_w}" height="{row_h - 8}" '
                f'fill="{BG}" rx="3"/>'
            )

            rank_evs = [e for e in rd["events"] if e.rank == rank]
            dead_x: float | None = None

            for ev in rank_evs:
                d = ev.data.get("duration_s", 0)
                is_err = bool(ev.data.get("error"))
                is_bad_exit = ev.kind == "exec" and ev.data.get("exit_code", 0) != 0

                if is_err or is_bad_exit:
                    color = COLORS["dead"]
                    dead_x = x_pos(ev.time)
                else:
                    color = COLORS.get(ev.kind, COLORS["push"])

                bx = x_pos(ev.time - d)
                bw = bar_w(d)
                tooltip = _tooltip(ev)
                parts.append(
                    f'<rect x="{bx}" y="{row_y + 6}" width="{bw}" height="{row_h - 12}" '
                    f'fill="{color}" rx="2" opacity="0.9">'
                    f"<title>{html.escape(tooltip)}</title></rect>"
                )

            # dead zone: hatched red from failure to end
            if dead_x is not None:
                dw = label_w + chart_w - dead_x
                if dw > 0:
                    parts.append(
                        f'<rect x="{dead_x}" y="{row_y + 6}" '
                        f'width="{dw}" height="{row_h - 12}" '
                        f'fill="{COLORS["dead"]}" rx="2" opacity="0.25"/>'
                    )

            # no events at all — full dead bar
            if not rank_evs and rank < rd.get("num_workers", 0):
                parts.append(
                    f'<rect x="{label_w}" y="{row_y + 6}" '
                    f'width="{chart_w}" height="{row_h - 12}" '
                    f'fill="{COLORS["dead"]}" rx="2" opacity="0.3">'
                    f"<title>worker {rank}: no response</title></rect>"
                )

            y_cursor += row_h + row_gap

        y_cursor += round_gap

    parts.append("</svg>")
    return "\n".join(parts)


def _tooltip(ev) -> str:
    kind = ev.kind
    d = ev.data
    dur = d.get("duration_s", 0)
    err = d.get("error")
    if err:
        return f"{kind} node={ev.rank}: FAILED — {err}"
    if kind == "push":
        return f"push {_fmt_bytes(d.get('size_bytes', 0))} → node {ev.rank} ({dur:.3f}s)"
    if kind == "pull":
        return f"pull {_fmt_bytes(d.get('size_bytes', 0))} ← node {ev.rank} ({dur:.3f}s)"
    if kind == "exec":
        tail = d.get("output_tail", "")
        return f"train node {ev.rank}: {dur:.2f}s exit={d.get('exit_code', '?')}\n{tail}"
    if kind == "aggregate":
        return (
            f"aggregate: {d.get('num_workers', '?')} workers, "
            f"|W|={d.get('weight_norm', 0):.2f}, Δ={d.get('delta_norm', 0):.4f}"
        )
    return f"{kind}: {d}"


# ------------------------------------------------------------------
# Full HTML page
# ------------------------------------------------------------------

def generate_html(tracer: Tracer, path: str) -> str:
    """Write a self-contained HTML report to `path`. Returns the path."""
    rounds = tracer._collect_rounds()

    # compute summary stats
    n_rounds = len(rounds)
    total_time = sum(r["duration"] for r in rounds.values())
    total_push = 0
    total_pull = 0
    for rd in rounds.values():
        total_push += sum(
            e.data.get("size_bytes", 0) for e in rd["events"]
            if e.kind == "push" and not e.data.get("error")
        )
        total_pull += sum(
            e.data.get("size_bytes", 0) for e in rd["events"]
            if e.kind == "pull" and not e.data.get("error")
        )
    max_workers = max((r["total"] for r in rounds.values()), default=0)
    final_wnorm = list(rounds.values())[-1]["weight_norm"] if rounds else 0

    svg = _build_svg(tracer)

    page = f"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<title>Training Run Report</title>
<style>
* {{ margin: 0; padding: 0; box-sizing: border-box; }}
body {{
  background: {BG}; color: {TEXT};
  font-family: 'SF Mono', 'Cascadia Code', 'Consolas', 'Menlo', monospace;
  padding: 32px; max-width: 960px; margin: 0 auto;
}}
h1 {{ font-size: 20px; margin-bottom: 24px; }}
.stats {{
  display: grid; grid-template-columns: repeat(auto-fit, minmax(140px, 1fr));
  gap: 12px; margin-bottom: 32px;
}}
.stat {{
  background: {CARD}; border-radius: 8px; padding: 16px;
  border: 1px solid {GRID};
}}
.stat-value {{ font-size: 22px; font-weight: bold; }}
.stat-label {{ font-size: 11px; color: {MUTED}; margin-top: 4px; }}
.legend {{
  display: flex; gap: 20px; margin: 16px 0; flex-wrap: wrap;
}}
.legend-item {{ display: flex; align-items: center; gap: 6px; font-size: 12px; }}
.legend-dot {{ width: 12px; height: 12px; border-radius: 3px; flex-shrink: 0; }}
.timeline {{ margin-top: 8px; }}
.timeline svg {{ width: 100%; height: auto; }}
.event-log {{
  margin-top: 32px; background: {CARD}; border-radius: 8px;
  padding: 16px; border: 1px solid {GRID};
  max-height: 400px; overflow-y: auto;
}}
.event-log pre {{
  font-size: 11px; line-height: 1.6; color: {MUTED};
  white-space: pre-wrap;
}}
.event-log .err {{ color: {COLORS["dead"]}; }}
</style>
</head>
<body>
<h1>Training Run Report</h1>

<div class="stats">
  <div class="stat">
    <div class="stat-value">{n_rounds}</div>
    <div class="stat-label">Rounds</div>
  </div>
  <div class="stat">
    <div class="stat-value">{max_workers}</div>
    <div class="stat-label">Workers</div>
  </div>
  <div class="stat">
    <div class="stat-value">{total_time:.1f}s</div>
    <div class="stat-label">Total Time</div>
  </div>
  <div class="stat">
    <div class="stat-value">{_fmt_bytes(total_push + total_pull)}</div>
    <div class="stat-label">Data Transferred</div>
  </div>
  <div class="stat">
    <div class="stat-value">{_fmt_bytes(total_push)}</div>
    <div class="stat-label">Params Pushed</div>
  </div>
  <div class="stat">
    <div class="stat-value">{_fmt_bytes(total_pull)}</div>
    <div class="stat-label">Weights Pulled</div>
  </div>
  <div class="stat">
    <div class="stat-value">{final_wnorm:.1f}</div>
    <div class="stat-label">Final |W|</div>
  </div>
</div>

<div class="legend">
  <div class="legend-item"><div class="legend-dot" style="background:{COLORS['push']}"></div>Push params</div>
  <div class="legend-item"><div class="legend-dot" style="background:{COLORS['exec']}"></div>Train</div>
  <div class="legend-item"><div class="legend-dot" style="background:{COLORS['pull']}"></div>Pull weights</div>
  <div class="legend-item"><div class="legend-dot" style="background:{COLORS['aggregate']}"></div>Aggregate</div>
  <div class="legend-item"><div class="legend-dot" style="background:{COLORS['dead']}"></div>Dead / Error</div>
</div>

<div class="timeline">
{svg}
</div>

<div class="event-log">
<pre>{_render_event_log(tracer)}</pre>
</div>

</body>
</html>"""

    path = os.path.abspath(path)
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    with open(path, "w") as f:
        f.write(page)
    return path


def _render_event_log(tracer: Tracer) -> str:
    """Render the text event log as HTML-escaped pre-formatted text."""
    lines = []
    for ev in tracer.events:
        err = ev.data.get("error")
        t = f"[{ev.time:7.2f}s]"
        node = f"node={ev.rank}" if ev.rank is not None else "      "

        match ev.kind:
            case "push" | "pull":
                sz = _fmt_bytes(ev.data.get("size_bytes", 0))
                dur = ev.data.get("duration_s", 0)
                p = os.path.basename(ev.data.get("path", ""))
                if err:
                    line = f'{t} {ev.kind:5s} {node}  {p:<20s} <span class="err">!! {html.escape(err)}</span>'
                else:
                    line = f"{t} {ev.kind:5s} {node}  {p:<20s} {sz:>8s}  {dur:.3f}s"
            case "exec":
                dur = ev.data.get("duration_s", 0)
                ec = ev.data.get("exit_code", "?")
                if err:
                    line = f'{t} exec  {node}  <span class="err">!! {html.escape(err)}</span>'
                else:
                    tail = html.escape(ev.data.get("output_tail", ""))
                    line = f"{t} exec  {node}  exit={ec:<3}  {dur:.3f}s"
                    if tail:
                        line += f"  | {tail}"
            case "round_start":
                r = ev.data.get("round_num", "?")
                n = ev.data.get("active_nodes", "?")
                line = f"{t} {'─'*4} round {r} start ({n} workers) {'─'*16}"
            case "round_end":
                r = ev.data.get("round_num", "?")
                s = ev.data.get("survivors", "?")
                tot = ev.data.get("total_nodes", "?")
                dur = ev.data.get("duration_s", 0)
                line = f"{t} {'─'*4} round {r} end ({s}/{tot} survived, {dur:.2f}s) {'─'*8}"
            case "aggregate":
                n = ev.data.get("num_workers", "?")
                wn = ev.data.get("weight_norm", 0)
                dn = ev.data.get("delta_norm", 0)
                dur = ev.data.get("duration_s", 0)
                line = f"{t} agg   {n} workers  |W|={wn:.2f}  Δ={dn:.4f}  {dur:.3f}s"
            case _:
                line = f"{t} {ev.kind}"

        lines.append(line)
    return "\n".join(lines)