"""Tracing infrastructure for observing data flow through the system.

Three output modes, all fed from the same event list:
  C) Live text log   — Tracer._print() streams events as they happen
  A) Terminal Gantt   — Tracer.gantt() renders a Unicode timeline per round
  B) HTML report      — report.generate_html(tracer, path) writes a self-contained SVG timeline

Usage:
    tracer = Tracer()
    sched = Scheduler(job, transport, tracer=tracer)
    await sched.run()
    tracer.gantt()
    tracer.summary()
"""

from __future__ import annotations

import os
import sys
import time
from dataclasses import dataclass, field


# ------------------------------------------------------------------
# Events
# ------------------------------------------------------------------

@dataclass
class Event:
    time: float  # seconds since tracer start
    kind: str  # push | pull | exec | round_start | round_end | aggregate
    rank: int | None  # None for scheduler-level events
    data: dict = field(default_factory=dict)


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"


# ------------------------------------------------------------------
# Tracer
# ------------------------------------------------------------------

class Tracer:
    def __init__(self, file=None):
        self._t0 = time.monotonic()
        self.events: list[Event] = []
        self._file = file or sys.stderr

    def emit(self, kind: str, rank: int | None = None, **data) -> Event:
        ev = Event(
            time=time.monotonic() - self._t0,
            kind=kind,
            rank=rank,
            data=data,
        )
        self.events.append(ev)
        self._print(ev)
        return ev

    # ------------------------------------------------------------------
    # C) Live text log
    # ------------------------------------------------------------------

    def _print(self, ev: Event):
        t = f"[{ev.time:7.2f}s]"
        node = f"node={ev.rank}" if ev.rank is not None else "      "
        err = ev.data.get("error")

        match ev.kind:
            case "push":
                sz = _fmt_bytes(ev.data.get("size_bytes", 0))
                dur = ev.data.get("duration_s", 0)
                path = os.path.basename(ev.data.get("path", ""))
                if err:
                    line = f"{t} push  {node}  {path:<20s} !! {err}"
                else:
                    line = f"{t} push  {node}  {path:<20s} {sz:>8s}  {dur:.3f}s"
            case "pull":
                sz = _fmt_bytes(ev.data.get("size_bytes", 0))
                dur = ev.data.get("duration_s", 0)
                path = os.path.basename(ev.data.get("path", ""))
                if err:
                    line = f"{t} pull  {node}  {path:<20s} !! {err}"
                else:
                    line = f"{t} pull  {node}  {path:<20s} {sz:>8s}  {dur:.3f}s"
            case "exec":
                dur = ev.data.get("duration_s", 0)
                exit_code = ev.data.get("exit_code", "?")
                if err:
                    line = f"{t} exec  {node}  !! {err}"
                else:
                    line = f"{t} exec  {node}  exit={exit_code:<3}  {dur:.3f}s"
                    output = ev.data.get("output_tail", "")
                    if output:
                        line += f"  | {output}"
            case "round_start":
                r = ev.data.get("round_num", "?")
                n = ev.data.get("active_nodes", "?")
                line = f"{t} {'─'*4} round {r} start ({n} workers) {'─'*20}"
            case "round_end":
                r = ev.data.get("round_num", "?")
                surv = ev.data.get("survivors", "?")
                total = ev.data.get("total_nodes", "?")
                dur = ev.data.get("duration_s", 0)
                line = f"{t} {'─'*4} round {r} end ({surv}/{total} survived, {dur:.2f}s) {'─'*12}"
            case "aggregate":
                n = ev.data.get("num_workers", "?")
                wnorm = ev.data.get("weight_norm", 0)
                dnorm = ev.data.get("delta_norm", 0)
                dur = ev.data.get("duration_s", 0)
                line = (
                    f"{t} agg   {n} workers"
                    f"  |W|={wnorm:.2f}"
                    f"  Δ={dnorm:.4f}"
                    f"  {dur:.3f}s"
                )
            case _:
                line = f"{t} {ev.kind:6s} {node}  {ev.data}"

        print(line, file=self._file, flush=True)

    # ------------------------------------------------------------------
    # helpers
    # ------------------------------------------------------------------

    def _collect_rounds(self) -> dict[int, dict]:
        """Group events into per-round buckets with timing metadata."""
        rounds: dict[int, dict] = {}
        current_round: int | None = None

        for ev in self.events:
            if ev.kind == "round_start":
                current_round = ev.data["round_num"]
                rounds[current_round] = {
                    "start": ev.time,
                    "end": ev.time,
                    "duration": 0.0,
                    "num_workers": ev.data.get("active_nodes", 0),
                    "survivors": 0,
                    "total": 0,
                    "events": [],
                    "weight_norm": 0.0,
                    "delta_norm": 0.0,
                }
            elif ev.kind == "round_end":
                r = ev.data["round_num"]
                if r in rounds:
                    rounds[r]["end"] = ev.time
                    rounds[r]["duration"] = ev.data.get("duration_s", 0)
                    rounds[r]["survivors"] = ev.data.get("survivors", 0)
                    rounds[r]["total"] = ev.data.get("total_nodes", 0)
            elif ev.kind == "aggregate":
                if current_round is not None and current_round in rounds:
                    rounds[current_round]["weight_norm"] = ev.data.get("weight_norm", 0)
                    rounds[current_round]["delta_norm"] = ev.data.get("delta_norm", 0)
                    rounds[current_round]["events"].append(ev)
            else:
                if current_round is not None and current_round in rounds:
                    rounds[current_round]["events"].append(ev)

        return rounds

    def _max_rank(self) -> int:
        ranks = [ev.rank for ev in self.events if ev.rank is not None]
        return max(ranks) if ranks else 0

    # ------------------------------------------------------------------
    # A) Terminal Gantt chart
    # ------------------------------------------------------------------

    def gantt(self, width: int = 60) -> str:
        """Render a Unicode Gantt chart showing data flow per round."""
        rounds = self._collect_rounds()
        if not rounds:
            return "(no rounds recorded)"

        max_rank = self._max_rank()

        lines: list[str] = []
        lines.append("")
        lines.append("  ▓ push   █ train   ▒ pull   ● agg   ✗ dead   ░ wait")
        lines.append("")

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

            lines.append(f"  round {rnum}  [{surv}/{total}]  {dur:.2f}s")

            def to_col(t: float) -> int:
                frac = (t - t0) / dur
                return max(0, min(width - 1, int(frac * width)))

            def fill(row: list[str], t_start: float, t_end: float, char: str):
                s = to_col(t_start)
                e = to_col(t_end)
                if s == e:
                    e = min(s + 1, width - 1)
                for i in range(s, e + 1):
                    if 0 <= i < width:
                        row[i] = char

            # --- scheduler row ---
            sched = list("░" * width)
            for ev in rd["events"]:
                d = ev.data.get("duration_s", 0)
                if ev.kind == "push":
                    fill(sched, ev.time - d, ev.time, "▓")
                elif ev.kind == "pull":
                    fill(sched, ev.time - d, ev.time, "▒")
                elif ev.kind == "aggregate":
                    fill(sched, ev.time - d, ev.time, "●")
            lines.append(f"    scheduler  {''.join(sched)}")

            # --- worker rows ---
            for rank in range(max_rank + 1):
                row = list("─" * width)
                rank_evs = [e for e in rd["events"] if e.rank == rank]
                dead_col = None

                for ev in rank_evs:
                    d = ev.data.get("duration_s", 0)
                    if ev.data.get("error"):
                        dead_col = to_col(ev.time)
                        fill(row, ev.time - d, ev.time, "✗")
                    elif ev.kind == "exec" and ev.data.get("exit_code", 0) != 0:
                        dead_col = to_col(ev.time)
                        fill(row, ev.time - d, ev.time, "✗")
                    elif ev.kind == "exec":
                        fill(row, ev.time - d, ev.time, "█")
                    elif ev.kind == "push":
                        fill(row, ev.time - d, ev.time, "▓")
                    elif ev.kind == "pull":
                        fill(row, ev.time - d, ev.time, "▒")

                # fill dead from failure point onward
                if dead_col is not None:
                    for i in range(dead_col, width):
                        if row[i] == "─":
                            row[i] = "✗"

                # worker was expected but produced nothing
                if not rank_evs and rank < rd.get("num_workers", 0):
                    row = list("✗" * width)

                lines.append(f"    worker {rank:<3}  {''.join(row)}")

            lines.append("")

        text = "\n".join(lines)
        print(text, file=self._file, flush=True)
        return text

    # ------------------------------------------------------------------
    # Summary table
    # ------------------------------------------------------------------

    def summary(self) -> str:
        """Render a round-by-round summary table."""
        rounds = self._collect_rounds()
        if not rounds:
            return "(no rounds recorded)"

        lines: list[str] = []
        lines.append("")

        header = (
            f"{'Round':<6} {'Workers':<9} {'Params↑':<10} {'Weights↓':<11}"
            f" {'Train(max)':<11} {'Agg':<8} {'Total':<8} {'Δ norm'}"
        )
        lines.append(header)
        lines.append("─" * len(header))

        total_push = 0
        total_pull = 0
        total_time = 0.0

        for rnum in sorted(rounds.keys()):
            rd = rounds[rnum]
            surv = rd["survivors"]
            total = rd["total"]

            push_bytes = sum(
                e.data.get("size_bytes", 0)
                for e in rd["events"]
                if e.kind == "push" and not e.data.get("error")
            )
            pull_bytes = sum(
                e.data.get("size_bytes", 0)
                for e in rd["events"]
                if e.kind == "pull" and not e.data.get("error")
            )
            exec_durs = [
                e.data.get("duration_s", 0)
                for e in rd["events"]
                if e.kind == "exec" and not e.data.get("error")
            ]
            agg_dur = sum(
                e.data.get("duration_s", 0)
                for e in rd["events"]
                if e.kind == "aggregate"
            )
            train_max = max(exec_durs) if exec_durs else 0
            dur = rd["duration"]
            delta = rd["delta_norm"]

            total_push += push_bytes
            total_pull += pull_bytes
            total_time += dur

            lines.append(
                f"{rnum:<6} {surv}/{total:<6} {_fmt_bytes(push_bytes):<10}"
                f" {_fmt_bytes(pull_bytes):<11} {train_max:<11.2f}"
                f" {agg_dur:<8.3f} {dur:<8.2f} {delta:.4f}"
            )

        lines.append("─" * len(header))
        lines.append(
            f"{'total':<6} {'':9} {_fmt_bytes(total_push):<10}"
            f" {_fmt_bytes(total_pull):<11} {'':11}"
            f" {'':8} {total_time:<8.2f}"
        )

        text = "\n".join(lines)
        print(text, file=self._file, flush=True)
        return text


# ------------------------------------------------------------------
# TracingConnection wrapper
# ------------------------------------------------------------------

class TracingConnection:
    """Wraps any Connection to auto-trace push/pull/exec with timing and sizes.

    Emits events on both success and failure so the Gantt chart can show
    where workers died.
    """

    def __init__(self, inner, rank: int, tracer: Tracer):
        self._inner = inner
        self._rank = rank
        self._tracer = tracer

    # proxy unknown attributes to inner (e.g. MockConnection.kill, .root_dir)
    def __getattr__(self, name):
        return getattr(self._inner, name)

    async def exec(self, cmd: str, timeout: float = 30.0) -> tuple[int, str]:
        t0 = time.monotonic()
        try:
            ret, output = await self._inner.exec(cmd, timeout=timeout)
        except Exception as exc:
            dur = time.monotonic() - t0
            self._tracer.emit(
                "exec", rank=self._rank,
                duration_s=dur, exit_code=-1, output_tail="",
                error=str(exc),
            )
            raise

        dur = time.monotonic() - t0
        tail = ""
        for line in reversed(output.strip().splitlines()):
            if line.strip():
                tail = line.strip()[:80]
                break

        self._tracer.emit(
            "exec", rank=self._rank,
            duration_s=dur, exit_code=ret, output_tail=tail,
        )
        return ret, output

    async def push(self, local_path: str, remote_path: str) -> None:
        size = os.path.getsize(local_path) if os.path.exists(local_path) else 0
        t0 = time.monotonic()
        try:
            await self._inner.push(local_path, remote_path)
        except Exception as exc:
            dur = time.monotonic() - t0
            self._tracer.emit(
                "push", rank=self._rank,
                path=remote_path, size_bytes=0, duration_s=dur,
                error=str(exc),
            )
            raise

        dur = time.monotonic() - t0
        self._tracer.emit(
            "push", rank=self._rank,
            path=remote_path, size_bytes=size, duration_s=dur,
        )

    async def pull(self, remote_path: str, local_path: str) -> None:
        t0 = time.monotonic()
        try:
            await self._inner.pull(remote_path, local_path)
        except Exception as exc:
            dur = time.monotonic() - t0
            self._tracer.emit(
                "pull", rank=self._rank,
                path=remote_path, size_bytes=0, duration_s=dur,
                error=str(exc),
            )
            raise

        dur = time.monotonic() - t0
        size = os.path.getsize(local_path) if os.path.exists(local_path) else 0
        self._tracer.emit(
            "pull", rank=self._rank,
            path=remote_path, size_bytes=size, duration_s=dur,
        )

    async def close(self) -> None:
        await self._inner.close()