"""Log line parsers for OpOp TUI.

Regex patterns derived from train_resonance.py print() calls.
Each parser returns a typed dataclass or None.
"""

import re
from dataclasses import dataclass, field
from typing import Optional, Dict, Union


# ── Event types ──────────────────────────────────────────────────────────────

@dataclass
class ProgressEvent:
    batch: int
    total_batches: int
    pct: float
    loss: float
    ce_loss: Optional[float] = None
    adv_mse: Optional[float] = None
    acc: float = 0.0
    # Sophia
    sophia_clip_pct: Optional[float] = None
    sophia_gnorm: Optional[float] = None
    sophia_maxu: Optional[float] = None
    # Resonance
    resonance_gnorm: Optional[float] = None
    # Brain
    brain_proj_scale: Optional[float] = None
    brain_proj_urgency: Optional[float] = None
    brain_mlp_scale: Optional[float] = None
    brain_mlp_urgency: Optional[float] = None
    brain_reward: Optional[float] = None
    # ETA
    remaining_seconds: Optional[float] = None


@dataclass
class EvalEvent:
    top1: float
    top5: float
    n_tokens: int


@dataclass
class GenEvent:
    ttc_info: str
    prompt: str
    generated: str


@dataclass
class CheckpointEvent:
    path: str
    size_mb: float


@dataclass
class ControlEvent:
    param: str
    old_value: str
    new_value: str


@dataclass
class DiagEvent:
    metrics: Dict[str, float]


@dataclass
class InfoLine:
    """Any non-event line worth showing."""
    text: str


# ── Regex patterns ───────────────────────────────────────────────────────────

# [10/219726] 0% loss=6.6654 (ce=6.6654) adv_mse=0.0027 acc=0.157 | brain: p=0.34x/2.7u m=0.34x/2.5u r=0.007 ~1055469s remaining
RE_PROGRESS = re.compile(
    r'\[(\d+)/(\d+)\]\s+(\d+)%\s+'
    r'loss=([\d.]+)'
    r'(?:\s+\(ce=([\d.]+)\))?'
    r'(?:\s+adv_mse=([\d.]+))?'
    r'\s+acc=([\d.]+)'
    r'(?:\s+\|\s+sophia:\s+([\d.]+)%clip\s+gnorm=([\d.]+)\s+maxu=([\d.]+))?'
    r'(?:\s+\|\s+resonance:\s+gnorm=([\d.]+))?'
    r'(?:\s+\|\s+brain:\s+p=([\d.]+)x/([\d.]+)u\s+m=([\d.]+)x/([\d.]+)u\s+r=(-?[\d.]+))?'
    r'\s+~(-?[\d]+)s\s+remaining'
)

# [eval] top1=9.3% top5=20.9% (1024 tok)
RE_EVAL = re.compile(
    r'\[eval\]\s+top1=([\d.]+)%\s+top5=([\d.]+)%\s+\((\d+)\s+tok'
)

# [gen] [ttc pos=80 depth=0]  "The meaning..." -> "generated text"
RE_GEN = re.compile(
    r'\[gen\]\s*(\[.*?\])?\s*"(.*?)"\s*->\s*"(.*)"',
    re.DOTALL
)

# [ckpt] Saved /path/to/ckpt.npz (2988.3 MB)
RE_CKPT = re.compile(
    r'\[ckpt\]\s+Saved\s+(.*?)\s+\(([\d.]+)\s+MB\)'
)

# [ctrl] lr: 0.001 -> 0.0001
RE_CTRL = re.compile(
    r'\[ctrl\]\s+(\w+):\s+(\S+)\s+->\s+(\S+)'
)

# [diag] mlp_W0_norm=0.0088 advisor_W1_norm=0.0034
RE_DIAG = re.compile(r'\[diag\]\s+(.*)')

# key=value pairs inside diag lines
RE_KV = re.compile(r'(\w+)=([\d.e+-]+)')


# ── Parser ───────────────────────────────────────────────────────────────────

ParseResult = Optional[Union[
    ProgressEvent, EvalEvent, GenEvent,
    CheckpointEvent, ControlEvent, DiagEvent, InfoLine
]]


def parse_line(line: str) -> ParseResult:
    """Parse a single log line into a typed event, or None for blank lines."""
    stripped = line.strip()
    if not stripped:
        return None

    # Progress
    m = RE_PROGRESS.search(stripped)
    if m:
        g = m.groups()
        return ProgressEvent(
            batch=int(g[0]),
            total_batches=int(g[1]),
            pct=float(g[2]),
            loss=float(g[3]),
            ce_loss=float(g[4]) if g[4] else None,
            adv_mse=float(g[5]) if g[5] else None,
            acc=float(g[6]),
            sophia_clip_pct=float(g[7]) if g[7] else None,
            sophia_gnorm=float(g[8]) if g[8] else None,
            sophia_maxu=float(g[9]) if g[9] else None,
            resonance_gnorm=float(g[10]) if g[10] else None,
            brain_proj_scale=float(g[11]) if g[11] else None,
            brain_proj_urgency=float(g[12]) if g[12] else None,
            brain_mlp_scale=float(g[13]) if g[13] else None,
            brain_mlp_urgency=float(g[14]) if g[14] else None,
            brain_reward=float(g[15]) if g[15] else None,
            remaining_seconds=float(g[16]) if g[16] else None,
        )

    # Eval
    m = RE_EVAL.search(stripped)
    if m:
        return EvalEvent(
            top1=float(m.group(1)),
            top5=float(m.group(2)),
            n_tokens=int(m.group(3)),
        )

    # Gen
    m = RE_GEN.search(stripped)
    if m:
        return GenEvent(
            ttc_info=m.group(1) or "",
            prompt=m.group(2),
            generated=m.group(3),
        )

    # Checkpoint
    m = RE_CKPT.search(stripped)
    if m:
        return CheckpointEvent(
            path=m.group(1),
            size_mb=float(m.group(2)),
        )

    # Control change
    m = RE_CTRL.search(stripped)
    if m:
        return ControlEvent(
            param=m.group(1),
            old_value=m.group(2),
            new_value=m.group(3),
        )

    # Diagnostics
    m = RE_DIAG.search(stripped)
    if m:
        metrics = {}
        for kv in RE_KV.finditer(m.group(1)):
            try:
                metrics[kv.group(1)] = float(kv.group(2))
            except ValueError:
                pass
        return DiagEvent(metrics=metrics)

    # Section headers and other info
    if stripped.startswith('[') or stripped.startswith('='):
        return InfoLine(text=stripped)

    return None
