"""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, replace
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"
BG_SELECTED = "\x1b[48;5;236m"  # Dark gray background for selected row
BLUE = "\x1b[34m"

# 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",
})

# Event types in per-sample JSONL files
SAMPLE_EVENT_TYPES = frozenset({
    "text_delta",
    "thinking_delta",
    "assistant_message",
    "tool_execution",
    "llm_call",
    "modal_progress",
    "turn",
    "sample_start",
    "sample_end",
})


# ── 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, selected: bool = False) -> 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}"
            row = f"  {name} {turn_info:>4}  {icon} score={sample.score:.2f}"
        else:
            row = f"  {name} {turn_info:>4}  {GREEN}✓{RESET} done"
    elif sample.status == "retry":
        row = f"  {name} {turn_info:>4}  {YELLOW}⟳{RESET} retry (attempt {sample.retry_attempt})"
    else:
        phase = sample.phase or "running"
        row = f"  {name} {turn_info:>4}  {CYAN}{phase}...{RESET}"

    if selected:
        # Add selection highlight
        row = f"{BG_SELECTED}{row}{RESET}"

    return row


# ── Viewport formatting ──────────────────────────────────────────────────────


def format_sample_event(
    event: dict[str, Any], last_turn: int | None = None
) -> tuple[str | None, int | None]:
    """Format a per-sample JSONL event into a readable line.

    Returns (formatted_line, new_turn_number).
    formatted_line is None if the event should not be displayed.
    new_turn_number is the turn number if this event changes it, else None.
    """
    msg = event.get("message", "")

    if msg == "text_delta":
        # Streaming text — return raw text for accumulation
        return event.get("text", ""), None

    elif msg == "thinking_delta":
        # Extended thinking — show dimmed
        text = event.get("text", "")
        return (f"{DIM}{text}{RESET}" if text else None), None

    elif msg == "assistant_message":
        # Final assistant message — skip if we've been streaming
        # The caller handles dedup
        return None, None

    elif msg == "tool_execution":
        tool = event.get("tool_name", "tool")
        duration = event.get("duration_ms", 0)
        status = event.get("result_summary", "")
        is_error = event.get("is_error", False)
        color = RED if is_error else CYAN
        return f"{color}→ {tool}{RESET} ({duration / 1000:.1f}s) {status}", None

    elif msg == "llm_call":
        provider = event.get("provider", "")
        model = event.get("model", "")
        tokens_in = event.get("tokens_in", 0)
        tokens_out = event.get("tokens_out", 0)
        duration = event.get("duration_ms", 0)
        model_short = model.split("/")[-1] if model else provider
        return (
            f"{DIM}◆ {model_short} {tokens_in}→{tokens_out} tok ({duration / 1000:.1f}s){RESET}",
            None,
        )

    elif msg == "modal_progress":
        phase = event.get("phase", "")
        return (f"{YELLOW}⟳ {phase}...{RESET}" if phase else None), None

    elif msg == "turn":
        turn = event.get("turn")
        # Only show turn separator when turn number changes
        if turn is not None and turn != last_turn:
            return f"{DIM}── turn {turn} ──{RESET}", turn
        return None, turn if turn is not None else last_turn

    return None, None


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


def render_header(state: RenderState, width: int) -> str:
    """Render the progress header line."""
    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)}]"

    return header[:width]


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

    lines.append(render_header(state, 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

    # Viewport state
    focus: str = "list"  # "list" or "viewport"
    selected_idx: int = 0  # Index into active samples list
    viewport_lines: tuple[str, ...] = ()  # Formatted lines for viewport
    viewport_scroll: int = 0  # Scroll position in viewport
    viewport_auto_follow: bool = True  # Auto-scroll to bottom on new content
    tailing_sample_id: str | None = None  # Which sample file we're tailing

    # For streaming text accumulation
    streaming_text: str = ""

    # For deduping turn separators
    last_turn: int | None = None


@dataclass(frozen=True)
class NewEvent:
    """New line from events.jsonl."""

    line: str


@dataclass(frozen=True)
class SampleEvent:
    """New line from per-sample JSONL file."""

    line: str


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


def _get_active_samples(events: tuple[dict[str, Any], ...]) -> list[SampleState]:
    """Get list of active (non-complete) samples from events."""
    state = derive_state(list(events))
    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)
    return active


def _get_selected_sample_id(events: tuple[dict[str, Any], ...], selected_idx: int) -> str | None:
    """Get the sample_id of the currently selected sample."""
    active = _get_active_samples(events)
    if 0 <= selected_idx < len(active):
        return active[selected_idx].id
    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()

        # Focus switching
        case KeyPress(key="\t"):  # Tab
            new_focus = "viewport" if model.focus == "list" else "list"
            return replace(model, focus=new_focus), Cmd.none()

        # Navigation in list focus
        case KeyPress(key="j" | "\x1b[B") if model.focus == "list":  # Down
            active = _get_active_samples(model.events)
            new_idx = min(model.selected_idx + 1, len(active) - 1) if active else 0
            new_sample_id = _get_selected_sample_id(model.events, new_idx)
            # Reset viewport when selection changes
            if new_sample_id != model.tailing_sample_id:
                return replace(
                    model,
                    selected_idx=new_idx,
                    tailing_sample_id=new_sample_id,
                    viewport_lines=(),
                    viewport_scroll=0,
                    viewport_auto_follow=True,
                    streaming_text="",
                    last_turn=None,
                ), Cmd.none()
            return replace(model, selected_idx=new_idx), Cmd.none()

        case KeyPress(key="k" | "\x1b[A") if model.focus == "list":  # Up
            new_idx = max(model.selected_idx - 1, 0)
            new_sample_id = _get_selected_sample_id(model.events, new_idx)
            if new_sample_id != model.tailing_sample_id:
                return replace(
                    model,
                    selected_idx=new_idx,
                    tailing_sample_id=new_sample_id,
                    viewport_lines=(),
                    viewport_scroll=0,
                    viewport_auto_follow=True,
                    streaming_text="",
                    last_turn=None,
                ), Cmd.none()
            return replace(model, selected_idx=new_idx), Cmd.none()

        # Navigation in viewport focus
        case KeyPress(key="j" | "\x1b[B") if model.focus == "viewport":  # Down
            new_scroll = model.viewport_scroll + 1
            return replace(
                model, viewport_scroll=new_scroll, viewport_auto_follow=False
            ), Cmd.none()

        case KeyPress(key="k" | "\x1b[A") if model.focus == "viewport":  # Up
            new_scroll = max(0, model.viewport_scroll - 1)
            return replace(
                model, viewport_scroll=new_scroll, viewport_auto_follow=False
            ), Cmd.none()

        case KeyPress(key="G") if model.focus == "viewport":  # Jump to bottom
            new_scroll = max(0, len(model.viewport_lines) - 1)
            return replace(model, viewport_scroll=new_scroll, viewport_auto_follow=True), Cmd.none()

        case KeyPress(key="g"):  # gg = top (simplified, just g for now)
            if model.focus == "viewport":
                return replace(model, viewport_scroll=0, viewport_auto_follow=False), Cmd.none()

        # Events from events.jsonl
        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 replace(model, events=new_events, done=True), Cmd.quit()

            # Auto-select first sample if none selected
            new_model = replace(model, events=new_events)
            if model.tailing_sample_id is None:
                sample_id = _get_selected_sample_id(new_events, model.selected_idx)
                if sample_id:
                    new_model = replace(new_model, tailing_sample_id=sample_id)

            return new_model, Cmd.none()

        # Events from per-sample JSONL file
        case SampleEvent(line=line):
            event = _parse_event(line)
            if event is None:
                return model, Cmd.none()

            formatted, new_turn = format_sample_event(event, model.last_turn)
            msg_type = event.get("message", "")

            # Handle streaming text accumulation
            if msg_type == "text_delta":
                new_streaming = model.streaming_text + (formatted or "")
                return replace(model, streaming_text=new_streaming), Cmd.none()

            # When we get a non-delta event, flush accumulated streaming text
            new_lines = list(model.viewport_lines)
            if model.streaming_text:
                # Wrap streaming text to reasonable width and add
                for line in model.streaming_text.split("\n"):
                    if line:
                        new_lines.append(line)
                new_streaming = ""
            else:
                new_streaming = model.streaming_text

            # Add the formatted event line
            if formatted and msg_type != "text_delta":
                new_lines.append(formatted)

            # Auto-scroll if enabled
            new_scroll = model.viewport_scroll
            if model.viewport_auto_follow:
                new_scroll = max(0, len(new_lines) - 1)

            return replace(
                model,
                viewport_lines=tuple(new_lines),
                viewport_scroll=new_scroll,
                streaming_text=new_streaming,
                last_turn=new_turn if new_turn is not None else model.last_turn,
            ), Cmd.none()

    return model, Cmd.none()


def view(model: Model, width: int, height: int) -> list[str]:
    """Render the split view: sample list on top, viewport on bottom."""
    state = derive_state(list(model.events))
    lines: list[str] = []

    # Header
    lines.append(render_header(state, width))

    # Calculate layout
    # Sample list gets ~1/3 of space (min 4 lines), viewport gets rest
    list_height = max(4, min(height // 3, 10))
    viewport_height = height - list_height - 3  # -3 for header, separator, footer

    # Active samples with selection highlight
    active = _get_active_samples(model.events)

    # Render sample list
    for i, sample in enumerate(active[: list_height - 1]):
        selected = i == model.selected_idx
        row = _format_sample_row(sample, width, selected=selected)
        lines.append(row)

    # Pad list area
    while len(lines) < list_height:
        lines.append("")

    # Show hidden count
    if len(active) > list_height - 1:
        hidden = len(active) - (list_height - 1)
        lines[list_height - 1] = f"{DIM}  ... and {hidden} more{RESET}"

    # Separator with focus indicator and selected sample name
    selected_name = ""
    if model.tailing_sample_id:
        active = _get_active_samples(model.events)
        if 0 <= model.selected_idx < len(active):
            selected_name = f" [{active[model.selected_idx].name[:20]}]"

    if model.focus == "list":
        label = "samples"
    else:
        label = "viewport"
    # Build separator to fill width
    prefix = f"─── {label}{selected_name} "
    fill = "─" * max(0, width - len(prefix))
    lines.append(f"{DIM}{prefix}{fill}{RESET}")

    # Viewport
    if model.tailing_sample_id:
        # Include any streaming text as a "virtual" line at the end
        display_lines = list(model.viewport_lines)
        if model.streaming_text:
            # Show streaming text (last 200 chars to keep it manageable)
            stream_preview = model.streaming_text[-200:]
            if len(model.streaming_text) > 200:
                stream_preview = "..." + stream_preview
            display_lines.append(f"{CYAN}{stream_preview}{RESET}")

        if display_lines:
            # Clamp scroll
            max_scroll = max(0, len(display_lines) - viewport_height)
            scroll = min(model.viewport_scroll, max_scroll)

            # Render visible lines
            visible = display_lines[scroll : scroll + viewport_height]
            for line in visible:
                # Truncate long lines (but preserve ANSI - crude truncation for now)
                # TODO: proper ANSI-aware truncation
                if len(line) > width + 20:  # Allow some slack for ANSI codes
                    lines.append(line[: width + 10] + "...")
                else:
                    lines.append(line)
        else:
            lines.append(f"{DIM}  (waiting for events...){RESET}")

        # Pad viewport
        while len(lines) < list_height + 1 + viewport_height:
            lines.append("")
    else:
        # No sample selected
        lines.append(f"{DIM}  (select a sample to view its output){RESET}")
        while len(lines) < list_height + 1 + viewport_height:
            lines.append("")

    # Footer
    if model.focus == "list":
        footer = f"{DIM}j/k: select  Tab: focus viewport  q: quit{RESET}"
    else:
        follow_status = "follow" if model.viewport_auto_follow else "manual"
        footer = f"{DIM}j/k: scroll  G: bottom  Tab: focus list  [{follow_status}]  q: quit{RESET}"

    # Pad to fill screen
    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]
    output_dir = Path(events_file).parent
    samples_dir = output_dir / "samples"

    def _subscriptions(model: Model) -> Sub:
        if model.done:
            return Sub.none()

        subs = [Sub.file_tail(events_file, lambda line: NewEvent(line=line))]

        # Tail the selected sample's file if we have one
        if model.tailing_sample_id:
            sample_file = samples_dir / f"{model.tailing_sample_id}.jsonl"
            if sample_file.exists():
                subs.append(Sub.file_tail(str(sample_file), lambda line: SampleEvent(line=line)))

        return Sub.batch(*subs)

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