"""OpOp TUI -- the optimizer optimizer is watching.

Launch with:
    opop monitor --log path/to/train.log --control path/to/control.json
"""

import json
import math
import os
import time
from typing import Dict, List, Optional

from rich.panel import Panel
from rich.table import Table
from rich.text import Text

from textual.app import App, ComposeResult
from textual.binding import Binding
from textual.containers import Horizontal, Vertical, VerticalScroll
from textual.widgets import (
    Button, DataTable, Footer, Header, Label,
    RichLog, Sparkline, Static, TabbedContent, TabPane, TextArea,
)

from opop._parsers import (
    parse_line, ProgressEvent, EvalEvent, GenEvent,
    CheckpointEvent, ControlEvent, DiagEvent, InfoLine,
)
from opop._gpu import poll_gpus


# ── Rendering helpers ────────────────────────────────────────────────────────

def _format_eta(seconds: Optional[float]) -> str:
    if seconds is None or seconds <= 0:
        return "???"
    s = int(seconds)
    d, s = divmod(s, 86400)
    h, s = divmod(s, 3600)
    m, _ = divmod(s, 60)
    if d > 0:
        return f"{d}d {h}h"
    if h > 0:
        return f"{h}h {m}m"
    return f"{m}m"


def _scale_bar(value: float, lo: float, hi: float, width: int = 10) -> str:
    """Render a value as a bar between lo and hi."""
    if hi <= lo:
        frac = 0.5
    else:
        # Log scale for scale values
        if lo > 0 and hi > 0:
            frac = (math.log(max(value, lo)) - math.log(lo)) / (math.log(hi) - math.log(lo))
        else:
            frac = (value - lo) / (hi - lo)
    frac = max(0.0, min(1.0, frac))
    filled = int(frac * width)
    return "[cyan]" + "=" * filled + "[/]" + " " * (width - filled)


def _mood(baseline: float) -> str:
    """Brain mood based on reward baseline."""
    if baseline > 0.02:
        return "[bold green]COOKING[/]"
    if baseline > 0.005:
        return "[green]LEARNING[/]"
    if baseline > -0.005:
        return "[yellow]VIBING[/]"
    if baseline > -0.02:
        return "[red]ADJUSTING[/]"
    return "[bold red]COPING[/]"


def render_brain_panel(brain: dict) -> Panel:
    """Render brain decisions as a Rich Panel."""
    baseline = brain.get('baseline_reward', brain.get('baseline', 0.0))

    table = Table(show_header=True, header_style="bold cyan",
                  box=None, padding=(0, 1), expand=True)
    table.add_column("Group", style="bold", width=8)
    table.add_column("Scale", justify="right", width=7)
    table.add_column("Clip", justify="right", width=5)
    table.add_column("", width=14)

    # Try HenOptBrain naming (proj_scale), then OptBrain naming (g0_scale)
    groups = [
        ('proj', 'proj', 0),
        ('mlp', 'mlp', 1),
        ('adv_W1', 'advisor_W1', 2),
        ('adv_W2', 'advisor_W2', 3),
        ('bias', 'bias', 4),
    ]

    for short, full, gi in groups:
        scale = brain.get(f'{full}_scale', brain.get(f'g{gi}_scale'))
        if scale is None:
            continue
        clip = brain.get(f'{full}_urgency', brain.get(f'g{gi}_clip', 1.0))

        # Color scale
        if scale > 1.5:
            sc = f"[green]{scale:.2f}x[/]"
        elif scale < 0.5:
            sc = f"[red]{scale:.2f}x[/]"
        else:
            sc = f"[yellow]{scale:.2f}x[/]"

        bar = _scale_bar(scale, 0.01, 10.0)
        table.add_row(short, sc, f"{clip:.1f}", f"[{bar}]")

    step = brain.get('step', '?')
    title = f"Brain  |  {_mood(baseline)}  |  step {step}"
    subtitle = f"reward baseline: {baseline:+.4f}"

    return Panel(table, title=title, subtitle=subtitle,
                 border_style="bright_magenta", height=12)


def render_gpu_panel(gpus: list) -> Panel:
    """Render GPU status as a Rich Panel."""
    table = Table(show_header=True, box=None, header_style="bold blue",
                  padding=(0, 1), expand=True)
    table.add_column("GPU", width=8)
    table.add_column("Temp", justify="right", width=5)
    table.add_column("Util", justify="right", width=5)
    table.add_column("VRAM", width=10)
    table.add_column("W", justify="right", width=4)

    for gpu in gpus:
        # Color temp
        if gpu.temp_c > 80:
            tc = f"[red]{gpu.temp_c}C[/]"
        elif gpu.temp_c > 60:
            tc = f"[yellow]{gpu.temp_c}C[/]"
        else:
            tc = f"[green]{gpu.temp_c}C[/]"

        # Color util
        if gpu.util_pct > 90:
            uc = f"[green]{gpu.util_pct}%[/]"
        elif gpu.util_pct > 50:
            uc = f"[yellow]{gpu.util_pct}%[/]"
        else:
            uc = f"[dim]{gpu.util_pct}%[/]"

        mem = f"{gpu.mem_used_mb/1024:.1f}/{gpu.mem_total_mb/1024:.0f}G"
        pwr = f"{gpu.power_w:.0f}"

        table.add_row(f"{gpu.short_name}-{gpu.index}", tc, uc, mem, pwr)

    return Panel(table, title="GPUs", border_style="blue", height=len(gpus) + 4)


# ── CSS ──────────────────────────────────────────────────────────────────────

TCSS = """
Screen {
    layout: vertical;
}

#main-container {
    height: 1fr;
}

#left-panel {
    width: 70%;
    height: 100%;
    border-right: thick $accent;
}

#right-panel {
    width: 30%;
    height: 100%;
    padding: 0 1;
    overflow-y: auto;
}

/* ── Sparkline section ── */
#sparkline-row {
    height: 9;
    margin: 0 1;
}

#loss-spark-box, #acc-spark-box {
    width: 1fr;
    height: 100%;
    border: round $primary;
    padding: 0 1;
}

#loss-spark-box Label, #acc-spark-box Label {
    text-style: bold;
    color: $text;
    width: 100%;
    text-align: center;
}

Sparkline {
    height: 5;
}

#loss-sparkline > .sparkline--max-color {
    color: $error;
}
#loss-sparkline > .sparkline--min-color {
    color: $success;
}
#acc-sparkline > .sparkline--max-color {
    color: $success;
}
#acc-sparkline > .sparkline--min-color {
    color: $warning;
}

/* ── Metrics table ── */
#metrics-table {
    height: auto;
    max-height: 10;
    margin: 0 1;
}

/* ── Log stream ── */
#log-stream {
    height: 1fr;
    border: round $secondary;
    margin: 0 1;
}

/* ── Tabs ── */
TabbedContent {
    height: 1fr;
}

/* ── Run title ── */
#run-title {
    text-style: bold;
    background: $boost;
    color: $text;
    padding: 0 2;
    height: 1;
    text-align: center;
    width: 100%;
}

/* ── Right panel widgets ── */
#brain-panel {
    height: auto;
    max-height: 14;
    margin-bottom: 1;
}

#gpu-panel {
    height: auto;
    max-height: 10;
    margin-bottom: 1;
}

#control-panel {
    height: auto;
    max-height: 16;
    border: round $warning;
    padding: 0 1;
}

#control-label {
    text-style: bold;
    color: $warning;
    text-align: center;
    width: 100%;
    height: 1;
}

#control-editor {
    height: 8;
}

#btn-row {
    height: 3;
    align: center middle;
}

#btn-row Button {
    margin: 0 1;
}

/* ── Gen scroll ── */
#gen-scroll {
    height: 1fr;
    margin: 0 1;
}
"""


# ── App ──────────────────────────────────────────────────────────────────────

class OpOpApp(App):
    """The OpOp Monitor -- a TUI for watching the optimizer optimize."""

    CSS = TCSS
    TITLE = "OpOp Monitor"
    SUB_TITLE = "the optimizer optimizer is watching"

    BINDINGS = [
        Binding("q", "quit", "Quit"),
        Binding("d", "switch_tab('tab-dash')", "Dashboard", show=True),
        Binding("l", "switch_tab('tab-logs')", "Logs", show=True),
        Binding("g", "switch_tab('tab-gen')", "Generations", show=True),
        Binding("c", "focus_control", "Control", show=True),
        Binding("r", "force_refresh", "Refresh", show=True),
    ]

    def __init__(self, log_paths: List[str], control_path: Optional[str],
                 control_dir: Optional[str], gpu_poll_interval: int = 5,
                 tail_interval: float = 0.5):
        super().__init__()
        self.log_paths = log_paths or []
        self.control_path = control_path
        self.control_dir = control_dir
        self.gpu_poll_interval = gpu_poll_interval
        self.tail_interval = tail_interval

        # Per-run tracking
        self._log_positions: Dict[str, int] = {}
        self._run_data: Dict[str, dict] = {}  # run_name -> latest metrics
        self._loss_data: Dict[str, List[float]] = {}
        self._acc_data: Dict[str, List[float]] = {}

        # Global
        self._brain_state: dict = {}
        self._gen_entries: List[GenEvent] = []
        self._eval_history: List[EvalEvent] = []
        self._primary_run: Optional[str] = None
        self._metrics_keys: Dict[str, str] = {}  # run_name -> row_key in table

        # Multiline gen accumulator: [gen] output spans multiple lines
        self._gen_accum: Optional[str] = None

    # ── Compose ──────────────────────────────────────────────────────────────

    def compose(self) -> ComposeResult:
        yield Header(show_clock=True)
        with Horizontal(id="main-container"):
            with Vertical(id="left-panel"):
                with TabbedContent(id="tabs"):
                    with TabPane("Dashboard", id="tab-dash"):
                        with Vertical():
                            yield Static(
                                "OpOp Monitor -- waiting for data...",
                                id="run-title")
                            with Horizontal(id="sparkline-row"):
                                with Vertical(id="loss-spark-box"):
                                    yield Label("Loss")
                                    yield Sparkline([], id="loss-sparkline")
                                with Vertical(id="acc-spark-box"):
                                    yield Label("Accuracy")
                                    yield Sparkline([], id="acc-sparkline")
                            yield DataTable(id="metrics-table")
                    with TabPane("Logs", id="tab-logs"):
                        yield RichLog(
                            id="log-stream", highlight=True,
                            markup=True, max_lines=2000)
                    with TabPane("Generations", id="tab-gen"):
                        with VerticalScroll(id="gen-scroll"):
                            yield RichLog(
                                id="gen-log", highlight=True,
                                markup=True, max_lines=500)
            with Vertical(id="right-panel"):
                yield Static(
                    Panel("Waiting for brain data...",
                          title="Brain", border_style="dim"),
                    id="brain-panel")
                yield Static(
                    Panel("Polling GPUs...",
                          title="GPUs", border_style="dim"),
                    id="gpu-panel")
                with Vertical(id="control-panel"):
                    yield Label("Live Control", id="control-label")
                    yield TextArea("",
                                   id="control-editor",
                                   show_line_numbers=True)
                    with Horizontal(id="btn-row"):
                        yield Button("Apply", id="btn-apply",
                                     variant="success")
                        yield Button("Reset", id="btn-reset",
                                     variant="warning")
        yield Footer()

    # ── Mount ────────────────────────────────────────────────────────────────

    def on_mount(self) -> None:
        # Metrics table columns
        mt = self.query_one("#metrics-table", DataTable)
        mt.add_columns("Run", "Step", "Loss", "CE", "Acc",
                        "Top1", "ETA")

        # Load control file
        if self.control_path and os.path.exists(self.control_path):
            try:
                with open(self.control_path) as f:
                    self.query_one("#control-editor", TextArea).text = f.read()
            except Exception:
                pass

        # Seed log positions near end of file
        for path in self.log_paths:
            if os.path.exists(path):
                size = os.path.getsize(path)
                # Start from last 100KB to catch recent history
                self._log_positions[path] = max(0, size - 102400)

        # Start timers
        self.set_interval(self.tail_interval, self._tail_logs)
        if self.gpu_poll_interval > 0:
            self.set_interval(self.gpu_poll_interval, self._poll_gpus)

    # ── Log tailing ──────────────────────────────────────────────────────────

    def _run_name(self, path: str) -> str:
        return os.path.basename(path).replace('.log', '')

    async def _tail_logs(self) -> None:
        for path in self.log_paths:
            if not os.path.exists(path):
                continue
            pos = self._log_positions.get(path, 0)
            try:
                with open(path, 'r', errors='replace') as f:
                    f.seek(pos)
                    new_lines = f.readlines()
                    self._log_positions[path] = f.tell()
            except Exception:
                continue

            run = self._run_name(path)
            if self._primary_run is None:
                self._primary_run = run

            for line in new_lines:
                stripped = line.rstrip('\n')

                # Multiline gen accumulation: [gen] output spans lines
                if self._gen_accum is not None:
                    # Continuation: lines that don't start with [
                    if stripped.strip() and not stripped.strip().startswith('['):
                        self._gen_accum += '\n' + stripped
                        continue
                    else:
                        # Flush accumulated gen
                        from opop._parsers import parse_line as _pl, GenEvent
                        gev = _pl(self._gen_accum)
                        if isinstance(gev, GenEvent):
                            self._gen_entries.append(gev)
                            self._update_gen_panel(gev)
                            log_widget = self.query_one("#log-stream", RichLog)
                            log_widget.write(f"[cyan]{run} [gen][/]")
                        self._gen_accum = None

                # Start gen accumulation
                if '    [gen]' in stripped or stripped.strip().startswith('[gen]'):
                    from opop._parsers import parse_line as _pl, GenEvent
                    gev = _pl(stripped)
                    if isinstance(gev, GenEvent):
                        # Single-line gen, process immediately
                        self._gen_entries.append(gev)
                        self._update_gen_panel(gev)
                        log_widget = self.query_one("#log-stream", RichLog)
                        log_widget.write(f"[cyan]{run} [gen][/]")
                    else:
                        # Multiline gen — start accumulating
                        self._gen_accum = stripped
                    continue

                self._process_line(stripped, run)

    def _process_line(self, line: str, run: str) -> None:
        event = parse_line(line)
        log_widget = self.query_one("#log-stream", RichLog)

        if isinstance(event, ProgressEvent):
            self._handle_progress(event, run)
            # Only show every 5th progress line in log to reduce noise
            if event.batch % 50 == 0:
                log_widget.write(f"[dim]{run}[/] {line.strip()}")

        elif isinstance(event, EvalEvent):
            self._eval_history.append(event)
            log_widget.write(f"[green bold]{run} {line.strip()}[/]")

        elif isinstance(event, GenEvent):
            self._gen_entries.append(event)
            self._update_gen_panel(event)
            log_widget.write(f"[cyan]{run} [gen][/]")

        elif isinstance(event, CheckpointEvent):
            log_widget.write(
                f"[yellow bold]{run} [ckpt] {event.size_mb:.0f}MB[/]")

        elif isinstance(event, ControlEvent):
            log_widget.write(
                f"[magenta bold]{run} [ctrl] "
                f"{event.param}: {event.old_value} -> {event.new_value}[/]")

        elif isinstance(event, DiagEvent):
            parts = " ".join(f"{k}={v:.4f}" for k, v in event.metrics.items())
            log_widget.write(f"[blue]{run} [diag] {parts}[/]")

        elif isinstance(event, InfoLine):
            if any(kw in event.text for kw in
                   ['Optimizer:', 'Resumed', 'Freed', 'TRAIN',
                    'BUILD', 'LOAD', 'HenonTiedLM', 'GenMLP']):
                log_widget.write(f"[bold]{run} {event.text}[/]")

    def _handle_progress(self, ev: ProgressEvent, run: str) -> None:
        # Track loss/acc per run
        if run not in self._loss_data:
            self._loss_data[run] = []
            self._acc_data[run] = []

        self._loss_data[run].append(ev.loss)
        self._acc_data[run].append(ev.acc)

        # Keep last 500
        if len(self._loss_data[run]) > 500:
            self._loss_data[run] = self._loss_data[run][-500:]
            self._acc_data[run] = self._acc_data[run][-500:]

        # Update sparklines for primary run
        if run == self._primary_run:
            try:
                self.query_one("#loss-sparkline", Sparkline).data = \
                    self._loss_data[run][-200:]
                self.query_one("#acc-sparkline", Sparkline).data = \
                    self._acc_data[run][-200:]
            except Exception:
                pass

        # Store latest metrics
        eta_str = _format_eta(ev.remaining_seconds)
        top1 = ""
        if self._eval_history:
            top1 = f"{self._eval_history[-1].top1:.1f}%"

        self._run_data[run] = {
            'step': f"{ev.batch}/{ev.total_batches}",
            'loss': f"{ev.loss:.4f}",
            'ce': f"{ev.ce_loss:.4f}" if ev.ce_loss else "-",
            'acc': f"{ev.acc:.3f}",
            'top1': top1,
            'eta': eta_str,
            'pct': ev.pct,
            'batch': ev.batch,
            'total': ev.total_batches,
        }

        # Update title for primary run
        if run == self._primary_run:
            self.query_one("#run-title", Static).update(
                f"{run} | step {ev.batch}/{ev.total_batches} "
                f"({ev.pct:.0f}%) | loss={ev.loss:.4f} | "
                f"acc={ev.acc:.3f} | ETA: {eta_str}")

        # Brain state
        if ev.brain_proj_scale is not None:
            self._brain_state = {
                'proj_scale': ev.brain_proj_scale,
                'proj_urgency': ev.brain_proj_urgency,
                'mlp_scale': ev.brain_mlp_scale,
                'mlp_urgency': ev.brain_mlp_urgency,
                'baseline_reward': ev.brain_reward,
                'step': ev.batch,
            }
            try:
                self.query_one("#brain-panel", Static).update(
                    render_brain_panel(self._brain_state))
            except Exception:
                pass

        # Sophia state for brain panel when no brain
        elif ev.sophia_clip_pct is not None and not self._brain_state:
            sophia_info = {
                'step': ev.batch,
                'sophia_clip': ev.sophia_clip_pct,
                'sophia_gnorm': ev.sophia_gnorm,
                'sophia_maxu': ev.sophia_maxu,
            }
            try:
                self.query_one("#brain-panel", Static).update(
                    _render_sophia_panel(sophia_info))
            except Exception:
                pass

        # Update metrics table
        self._update_metrics_table(run)

    def _update_metrics_table(self, run: str) -> None:
        mt = self.query_one("#metrics-table", DataTable)
        data = self._run_data.get(run, {})
        if not data:
            return

        row = (run, data['step'], data['loss'], data['ce'],
               data['acc'], data['top1'], data['eta'])

        if run in self._metrics_keys:
            try:
                key = self._metrics_keys[run]
                mt.update_cell(key, "Step", data['step'])
                mt.update_cell(key, "Loss", data['loss'])
                mt.update_cell(key, "CE", data['ce'])
                mt.update_cell(key, "Acc", data['acc'])
                mt.update_cell(key, "Top1", data['top1'])
                mt.update_cell(key, "ETA", data['eta'])
            except Exception:
                pass
        else:
            key = mt.add_row(*row)
            self._metrics_keys[run] = key

    def _update_gen_panel(self, ev: GenEvent) -> None:
        gen_log = self.query_one("#gen-log", RichLog)
        gen_log.write("")
        gen_log.write(f"[bold cyan]Prompt:[/] \"{ev.prompt}\"")
        gen_log.write(f"[bold green]Generated:[/] \"{ev.generated[:200]}\"")
        if ev.ttc_info:
            gen_log.write(f"[dim]{ev.ttc_info}[/]")
        gen_log.write("[dim]" + "-" * 60 + "[/]")

    async def _poll_gpus(self) -> None:
        gpus = poll_gpus()
        if gpus:
            try:
                self.query_one("#gpu-panel", Static).update(
                    render_gpu_panel(gpus))
            except Exception:
                pass

    # ── Control editing ──────────────────────────────────────────────────────

    def on_button_pressed(self, event: Button.Pressed) -> None:
        if event.button.id == "btn-apply":
            self._apply_control()
        elif event.button.id == "btn-reset":
            self._reset_control()

    def _apply_control(self) -> None:
        if not self.control_path:
            self.notify("No control file specified (use --control)",
                        severity="error")
            return
        editor = self.query_one("#control-editor", TextArea)
        text = editor.text
        try:
            parsed = json.loads(text)
            with open(self.control_path, 'w') as f:
                json.dump(parsed, f, indent=2)
                f.write('\n')
            self.notify(
                f"Written: {os.path.basename(self.control_path)}",
                severity="information")
        except json.JSONDecodeError as e:
            self.notify(f"Invalid JSON: {e}", severity="error")
        except Exception as e:
            self.notify(f"Write failed: {e}", severity="error")

    def _reset_control(self) -> None:
        if self.control_path and os.path.exists(self.control_path):
            try:
                with open(self.control_path) as f:
                    self.query_one("#control-editor", TextArea).text = f.read()
                self.notify("Reloaded from disk")
            except Exception as e:
                self.notify(f"Load failed: {e}", severity="error")

    # ── Keybindings ──────────────────────────────────────────────────────────

    def action_switch_tab(self, tab_id: str) -> None:
        self.query_one("#tabs", TabbedContent).active = tab_id

    def action_focus_control(self) -> None:
        self.query_one("#control-editor", TextArea).focus()

    def action_force_refresh(self) -> None:
        # Reset positions to re-read recent history
        for path in self.log_paths:
            if os.path.exists(path):
                size = os.path.getsize(path)
                self._log_positions[path] = max(0, size - 102400)
        self.notify("Refreshing logs...")


# ── Sophia fallback panel ────────────────────────────────────────────────────

def _render_sophia_panel(info: dict) -> Panel:
    """Render sophia stats when no brain is active."""
    table = Table(show_header=False, box=None, padding=(0, 1), expand=True)
    table.add_column("Metric", style="bold")
    table.add_column("Value", justify="right")

    clip = info.get('sophia_clip', 0)
    gnorm = info.get('sophia_gnorm', 0)
    maxu = info.get('sophia_maxu', 0)

    clip_color = "green" if clip < 0.5 else "yellow" if clip < 0.9 else "red"
    table.add_row("Clip %", f"[{clip_color}]{clip*100:.1f}%[/]")
    table.add_row("Grad Norm", f"{gnorm:.2f}")
    table.add_row("Max Update", f"{maxu:.1f}")

    return Panel(table, title=f"Sophia | step {info.get('step', '?')}",
                 border_style="yellow", height=8)
