"""Eval progress display built on pytui's Elm-style App.

Tails events.jsonl (written by evaluate() via logging), derives display state,
renders progress in alternate screen via a subprocess.

Usage:
    with progress_display(output_dir=output_dir):
        await evaluate(dataset, config)

Or detached mode (two terminals):
    # Terminal 1: Run eval (writes events.jsonl via logging)
    python run_eval.py --config config.py

    # Terminal 2: Watch progress
    python -m wafer_core.rollouts.progress_app /path/to/events.jsonl
"""

from __future__ import annotations

import json
import os
import signal
import subprocess
import sys
import time
from collections.abc import Generator
from contextlib import contextmanager
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

# ── ANSI colors ──────────────────────────────────────────────────────────────

DIM = "\x1b[90m"
RESET = "\x1b[0m"
GREEN = "\x1b[32m"
RED = "\x1b[31m"
YELLOW = "\x1b[33m"
CYAN = "\x1b[36m"
BOLD = "\x1b[1m"

# Event types the progress display understands.
# The "message" field in the JSONL record is the event type discriminator.
# Records with unrecognised message values are silently skipped.
KNOWN_EVENT_TYPES = frozenset({
    "eval_start",
    "eval_end",
    "sample_start",
    "sample_end",
    "turn",
    "modal_progress",
    "gepa_iteration",
    "gepa_accepted",
    "gepa_rejected",
    "sample_retry",
})


# ── State ────────────────────────────────────────────────────────────────────


@dataclass
class SampleState:
    """State of a single sample."""

    id: str
    name: str = ""
    turn: int = 0
    phase: str = ""
    score: float | None = None
    status: str = "started"  # started, complete, retry
    retry_attempt: int = 0
    start_time: float = field(default_factory=time.time)
    last_update: float = field(default_factory=time.time)


@dataclass
class RenderState:
    """Display state derived from events."""

    eval_name: str = ""
    total: int = 0
    completed: int = 0
    samples: dict[str, SampleState] = field(default_factory=dict)
    start_time: float = field(default_factory=time.time)
    gepa_iter: int | None = None
    gepa_total: int | None = None
    gepa_best: float | None = None
    scores: list[float] = field(default_factory=list)


# ── Pure functions ───────────────────────────────────────────────────────────


def _format_time(seconds: float) -> str:
    """Format seconds as H:MM:SS or M:SS."""
    h = int(seconds) // 3600
    m = int(seconds) % 3600 // 60
    s = int(seconds) % 60
    if h:
        return f"{h}:{m:02d}:{s:02d}"
    return f"{m}:{s:02d}"


def _format_bar(progress: float, width: int) -> str:
    """Render a progress bar with unicode blocks."""
    if width <= 0:
        return ""
    filled = progress * width
    full_blocks = int(filled)
    partial_idx = int(8 * filled) % 8
    partial = " ▏▎▍▌▋▊▉"[partial_idx] if partial_idx else ""
    bar = "█" * full_blocks + partial
    return bar.ljust(width)


def derive_state(events: list[dict[str, Any]]) -> RenderState:
    """Derive current display state from event stream. Stateless."""
    state = RenderState()

    for event in events:
        event_type = event.get("message")
        if event_type is None or event_type not in KNOWN_EVENT_TYPES:
            continue

        sample_id = event.get("sample_id", "")

        if event_type == "eval_start":
            state.eval_name = event.get("eval_name", "eval")
            state.total = event.get("total", 0)
            if ts := event.get("timestamp"):
                from datetime import datetime

                state.start_time = datetime.fromisoformat(ts.replace("Z", "+00:00")).timestamp()
            else:
                state.start_time = time.time()

        elif event_type == "sample_start":
            state.samples[sample_id] = SampleState(
                id=sample_id,
                name=event.get("sample_name", sample_id),
                last_update=time.time(),
                phase="starting",
            )

        elif event_type == "turn":
            if sample_id in state.samples:
                if "turn" in event:
                    state.samples[sample_id].turn = event["turn"]
                state.samples[sample_id].last_update = time.time()
                status = event.get("status", "running")
                state.samples[sample_id].phase = status

        elif event_type == "modal_progress":
            if sample_id in state.samples:
                state.samples[sample_id].phase = event.get("phase", "")
                state.samples[sample_id].last_update = time.time()

        elif event_type == "sample_end":
            if sample_id in state.samples:
                state.samples[sample_id].status = "complete"
                score = event.get("score")
                if score is not None:
                    state.samples[sample_id].score = score
                    state.scores.append(score)
                state.completed += 1

        elif event_type == "sample_retry":
            if sample_id in state.samples:
                state.samples[sample_id].status = "retry"
                state.samples[sample_id].retry_attempt = event.get("attempt", 1)

        elif event_type == "gepa_iteration":
            state.gepa_iter = event.get("iter")
            state.gepa_total = event.get("total")
            state.gepa_best = event.get("best")

        elif event_type == "eval_end":
            pass  # Handled by update() -> Cmd.quit()

    return state


def _format_sample_row(sample: SampleState, width: int) -> str:
    """Format a single sample row."""
    name = sample.name[:25].ljust(25)
    turn_info = f"T:{sample.turn}"

    if sample.status == "complete":
        if sample.score is not None:
            icon = f"{GREEN}✓{RESET}" if sample.score > 0.5 else f"{RED}✗{RESET}"
            return f"  {name} {turn_info:>4}  {icon} score={sample.score:.2f}"
        else:
            return f"  {name} {turn_info:>4}  {GREEN}✓{RESET} done"
    elif sample.status == "retry":
        return f"  {name} {turn_info:>4}  {YELLOW}⟳{RESET} retry (attempt {sample.retry_attempt})"
    else:
        phase = sample.phase or "running"
        return f"  {name} {turn_info:>4}  {CYAN}{phase}...{RESET}"


# ── Render function ──────────────────────────────────────────────────────────


def render(state: RenderState, width: int, height: int) -> list[str]:
    """Render state to list of lines. Pure function."""
    lines: list[str] = []

    # Header
    if state.gepa_iter is not None:
        header = (
            f"GEPA iter {state.gepa_iter}/{state.gepa_total or '?'} │ best: {state.gepa_best:.0%}"
            if state.gepa_best
            else f"GEPA iter {state.gepa_iter}/{state.gepa_total or '?'}"
        )
    else:
        elapsed = time.time() - state.start_time
        if state.total > 0:
            progress = state.completed / state.total
            bar_width = min(30, width - 50)
            bar = _format_bar(progress, bar_width)
            eta = ""
            if state.completed > 0 and progress < 1.0:
                remaining = (elapsed / progress) - elapsed
                eta = f" <{_format_time(remaining)}"
            header = (
                f"{state.eval_name}: {BOLD}{state.completed}/{state.total}{RESET} "
                f"|{bar}| {100 * progress:3.0f}% "
                f"[{_format_time(elapsed)}{eta}]"
            )
        else:
            header = f"{state.eval_name}: {state.completed} samples [{_format_time(elapsed)}]"

    lines.append(header[:width])

    # Active samples, sorted by most recently updated
    active = [s for s in state.samples.values() if s.status != "complete" and s.phase]
    active.sort(key=lambda s: s.last_update, reverse=True)

    max_samples = max(1, height - 4)
    for sample in active[:max_samples]:
        lines.append(_format_sample_row(sample, width))

    total_in_flight = sum(1 for s in state.samples.values() if s.status != "complete")
    hidden = total_in_flight - len(active[:max_samples])
    if hidden > 0:
        lines.append(f"{DIM}  ... and {hidden} more in flight{RESET}")

    if state.scores:
        mean_score = sum(state.scores) / len(state.scores)
        lines.append(f"{DIM}score: {mean_score:.1%}{RESET}")

    return lines


# ── Elm architecture (subprocess entry point) ────────────────────────────────


@dataclass(frozen=True)
class Model:
    """Immutable model for the pytui App."""

    events: tuple[dict[str, Any], ...] = ()
    done: bool = False


@dataclass(frozen=True)
class NewEvent:
    line: str


def _parse_event(line: str) -> dict[str, Any] | None:
    try:
        return json.loads(line)
    except json.JSONDecodeError:
        return None


def update(model: Model, msg: object) -> tuple[Model, Any]:
    from ._pytui import Cmd, KeyPress

    match msg:
        case KeyPress(key="q" | "\x03"):
            return model, Cmd.quit()
        case NewEvent(line=line):
            event = _parse_event(line)
            if event is None:
                return model, Cmd.none()
            new_events = model.events + (event,)
            if event.get("message") == "eval_end":
                return Model(events=new_events, done=True), Cmd.quit()
            return Model(events=new_events, done=model.done), Cmd.none()
    return model, Cmd.none()


def view(model: Model, width: int, height: int) -> list[str]:
    state = derive_state(list(model.events))
    lines = render(state, width, height)

    # Pad to fill screen, add footer
    footer = f"{DIM}q: quit{RESET}"
    while len(lines) < height - 1:
        lines.append("")
    lines.append(footer)
    return lines[:height]


# ── Context manager (spawns subprocess) ──────────────────────────────────────


@contextmanager
def progress_display(
    output_dir: Path | str,
    disable: bool = False,
) -> Generator[Path, None, None]:
    """Context manager that spawns a pytui progress display in a subprocess.

    The subprocess owns the terminal (alternate screen, raw mode via /dev/tty).
    The parent process writes events.jsonl via logging as usual.

    Args:
        output_dir: Directory containing events.jsonl
        disable: If True, skip progress display (verbose mode)
    """
    if disable:
        yield Path(output_dir)
        return

    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    events_file = output_dir / "events.jsonl"

    proc = subprocess.Popen(
        [sys.executable, "-m", "wafer_core.rollouts.progress_app", str(events_file)],
        # Don't touch stdout/stderr — child inherits the terminal.
        # pytui opens /dev/tty directly for input, writes to stdout for rendering.
        #
        # start_new_session: put subprocess in its own session so it doesn't
        # receive SIGINT from Ctrl+C. The parent gets SIGINT and terminates
        # the subprocess in the finally block. Without this, pytui's raw mode
        # swallows Ctrl+C (terminal driver doesn't generate SIGINT in raw mode)
        # and the parent never sees it.
        stdin=subprocess.DEVNULL,
        start_new_session=True,
    )

    try:
        yield output_dir
    finally:
        # Give the subprocess a chance to clean up its terminal state.
        proc.send_signal(signal.SIGTERM)
        try:
            proc.wait(timeout=5.0)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait()

        # Defensive: restore terminal to sane state in case the subprocess
        # didn't clean up (e.g. got SIGKILL, or crashed before restoring).
        # stty sane resets raw mode, re-enables echo, etc.
        tty_fd = os.open("/dev/tty", os.O_RDWR)
        try:
            subprocess.run(["stty", "sane"], stdin=tty_fd)
        finally:
            os.close(tty_fd)


# ── CLI entry point (subprocess runs this) ────────────────────────────────────

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python -m wafer_core.rollouts.progress_app <events.jsonl>", file=sys.stderr)
        sys.exit(1)

    from ._pytui import App, Cmd, Sub

    events_file = sys.argv[1]

    # Model needs the path for subscriptions. We use a non-frozen wrapper
    # to thread the path through — subscriptions() needs it but it's static.
    _events_path = events_file

    # Patch subscriptions to close over the path
    def _subscriptions(model: Model) -> Sub:
        if model.done:
            return Sub.none()
        return Sub.file_tail(_events_path, lambda line: NewEvent(line=line))

    App(
        init=(Model(), Cmd.none()),
        update=update,
        view=view,
        subscriptions=_subscriptions,
        alternate_screen=False,
    ).run()
