"""Real-time training monitor.

Watches training logs and alerts on plateaus, NaN loss, gradient instability,
and other training problems.

Supports multiple log formats:
- LeRobot text logs (step:1000 loss:0.8)
- wandb local files (wandb/latest-run/)
- Tensorboard event files (via tbparse, optional)
- JSON Lines ({"step": N, "loss": X})
- CSV files (step,loss,grad_norm)
"""

from __future__ import annotations

import csv
import json
import logging
import math
import time
from dataclasses import dataclass, field
from pathlib import Path

logger = logging.getLogger(__name__)


@dataclass
class MonitorState:
    """Running state of the training monitor."""

    last_step: int = 0
    last_loss: float = float("inf")
    min_loss: float = float("inf")
    plateau_start_step: int = 0
    plateau_start_loss: float = 0.0
    alerts_shown: set = field(default_factory=set)
    loss_history: list[float] = field(default_factory=list)
    step_history: list[int] = field(default_factory=list)


def monitor_training(
    training_dir: str,
    dataset: str | None = None,
    alert_on_plateau: bool = True,
    alert_on_nan: bool = True,
    poll_interval: float = 10.0,
) -> None:
    """Watch a training run and alert on problems.

    Supports log formats: LeRobot text logs, wandb, tensorboard (tbparse),
    JSON Lines, and CSV.

    Args:
        training_dir: Path to training output directory.
        dataset: Optional dataset ID for context.
        alert_on_plateau: Alert when loss plateaus.
        alert_on_nan: Alert on NaN/Inf loss.
        poll_interval: Seconds between log checks.
    """
    from orbit.utils.display import console

    training_path = Path(training_dir)
    if not training_path.exists():
        console.print(f"[red]Training directory not found: {training_path}[/red]")
        return

    console.print(f"\n[bold]Monitoring:[/bold] {training_path}")
    if dataset:
        console.print(f"[dim]Dataset: {dataset}[/dim]")
    console.print(f"[dim]Poll interval: {poll_interval}s — Press Ctrl+C to stop[/dim]\n")

    state = MonitorState()

    try:
        while True:
            # Parse all log formats
            all_entries = parse_all_logs(training_path)

            # Process new entries
            new_entries = [e for e in all_entries if e["step"] > state.last_step]

            for entry in new_entries:
                step = entry["step"]
                loss = entry["loss"]
                grad = entry.get("grad_norm")

                state.step_history.append(step)
                state.loss_history.append(loss)
                state.last_step = step

                # NaN check
                if alert_on_nan and (math.isnan(loss) or math.isinf(loss)):
                    alert_key = f"nan_{step}"
                    if alert_key not in state.alerts_shown:
                        state.alerts_shown.add(alert_key)
                        console.print(
                            f"  Step {step:>8,} | Loss: [red]NaN/Inf[/red] | "
                            f"[red bold]LOSS EXPLODED[/red bold]"
                        )
                        console.print(
                            "\n  [red]Training has failed.[/red] Common causes:"
                        )
                        console.print("    - Learning rate too high")
                        console.print("    - Dead joints causing gradient explosion")
                        if dataset:
                            console.print(f"    - Run: orbit debug {training_dir} --ai -d {dataset}")
                        else:
                            console.print(f"    - Run: orbit debug {training_dir} --ai")
                    continue

                # Normal progress display
                status_icon = _loss_trend_icon(loss, state)
                state.last_loss = loss
                state.min_loss = min(state.min_loss, loss)

                line = f"  Step {step:>8,} | Loss: {loss:.4f}"
                if grad is not None:
                    line += f" | Grad: {grad:.1f}"
                line += f" | {status_icon}"

                console.print(line)

                # Gradient spike check
                if grad is not None and grad > 100:
                    alert_key = f"grad_{step}"
                    if alert_key not in state.alerts_shown:
                        state.alerts_shown.add(alert_key)
                        console.print(
                            f"\n  [yellow]Gradient spike to {grad:.1f}[/yellow] — "
                            "consider gradient clipping (max_grad_norm=1.0)"
                        )

                # Plateau detection
                if alert_on_plateau and _detect_plateau(state, step):
                    alert_key = f"plateau_{step // 10000}"
                    if alert_key not in state.alerts_shown:
                        state.alerts_shown.add(alert_key)
                        plateau_len = step - state.plateau_start_step
                        console.print(
                            f"\n  [yellow]PLATEAU[/yellow] — loss hasn't improved in "
                            f"{plateau_len:,} steps"
                        )
                        console.print(
                            "  Common causes:"
                        )
                        console.print("    - Learning rate too low (try 3e-4)")
                        console.print("    - Data inconsistency")
                        if dataset:
                            console.print(
                                f"    - Run: orbit debug {training_dir} --ai -d {dataset}"
                            )
                        console.print()

            time.sleep(poll_interval)

    except KeyboardInterrupt:
        console.print("\n[dim]Monitor stopped.[/dim]")

    # Final summary
    if state.loss_history:
        console.print(f"\n[bold]Summary:[/bold]")
        console.print(
            f"  Steps: {state.step_history[0]:,} → {state.step_history[-1]:,}"
        )
        console.print(
            f"  Loss: {state.loss_history[0]:.4f} → {state.loss_history[-1]:.4f} "
            f"(min: {state.min_loss:.4f})"
        )
        console.print(f"\n  Next: orbit debug {training_dir}" + (f" -d {dataset}" if dataset else ""))
    console.print()


def _loss_trend_icon(current_loss: float, state: MonitorState) -> str:
    """Return a trend indicator for the loss."""
    if not state.loss_history:
        return "\u25bc dropping"

    prev = state.last_loss
    if math.isnan(current_loss) or math.isinf(current_loss):
        return "[red]EXPLODED[/red]"

    ratio = current_loss / max(prev, 1e-8)
    if ratio < 0.95:
        return "\u25bc dropping normally"
    elif ratio < 0.99:
        return "\u25bc approaching convergence"
    elif ratio < 1.01:
        return "\u25ac stable"
    elif ratio < 1.05:
        return "[yellow]\u25b2 slight increase[/yellow]"
    else:
        return "[red]\u25b2 INCREASING[/red]"


def _detect_plateau(state: MonitorState, current_step: int) -> bool:
    """Detect if training has plateaued."""
    if len(state.loss_history) < 10:
        return False

    # Check if loss hasn't improved significantly in last 10 entries
    recent = state.loss_history[-10:]
    oldest_recent = recent[0]
    newest_recent = recent[-1]

    ratio = newest_recent / max(oldest_recent, 1e-8)
    if ratio > 0.95 and current_step > 10_000:
        if state.plateau_start_step == 0:
            state.plateau_start_step = state.step_history[-10]
            state.plateau_start_loss = oldest_recent
        return True

    state.plateau_start_step = 0
    return False


# ---------------------------------------------------------------------------
# Multi-format log parsers
# ---------------------------------------------------------------------------


def parse_all_logs(training_path: Path) -> list[dict]:
    """Parse training logs from all supported formats.

    Supported formats (checked in order):
    1. Tensorboard event files (via tbparse, optional)
    2. JSON Lines (.jsonl files with {"step": N, "loss": X})
    3. CSV files (step,loss,grad_norm header)
    4. wandb local files
    5. LeRobot/plain text logs
    """
    entries: list[dict] = []

    # 1. Tensorboard events
    tb_entries = _parse_tensorboard(training_path)
    if tb_entries:
        entries.extend(tb_entries)

    # 2. JSON Lines
    entries.extend(_parse_jsonl(training_path))

    # 3. CSV
    entries.extend(_parse_csv_logs(training_path))

    # 4. wandb + text logs (from debugger)
    try:
        from orbit.debugger import _parse_lerobot_logs, _parse_wandb_logs

        entries.extend(_parse_lerobot_logs(training_path))
        entries.extend(_parse_wandb_logs(training_path))
    except Exception as e:
        import logging as _log
        _log.getLogger(__name__).debug("Failed to parse wandb/lerobot logs: %s", e)

    # Deduplicate by step (keep first occurrence)
    seen_steps: set[int] = set()
    unique = []
    for e in sorted(entries, key=lambda x: x.get("step", 0)):
        step = e.get("step", 0)
        if step not in seen_steps:
            seen_steps.add(step)
            unique.append(e)

    return unique


def _parse_tensorboard(training_path: Path) -> list[dict]:
    """Parse tensorboard event files using tbparse (optional dependency)."""
    entries: list[dict] = []

    # Find event files
    event_files = list(training_path.rglob("events.out.tfevents.*"))
    if not event_files:
        return entries

    try:
        from tbparse import SummaryReader

        for ef in event_files:
            try:
                reader = SummaryReader(str(ef.parent))
                df = reader.scalars
                if df.empty:
                    continue

                # Look for loss and grad_norm columns
                loss_tags = [t for t in df["tag"].unique()
                             if "loss" in t.lower() and "val" not in t.lower()]
                grad_tags = [t for t in df["tag"].unique()
                             if "grad" in t.lower()]

                if loss_tags:
                    loss_df = df[df["tag"] == loss_tags[0]]
                    for _, row in loss_df.iterrows():
                        entry: dict = {
                            "step": int(row["step"]),
                            "loss": float(row["value"]),
                            "source": "tensorboard",
                        }
                        entries.append(entry)

                if grad_tags and loss_tags:
                    grad_df = df[df["tag"] == grad_tags[0]]
                    grad_by_step = {int(r["step"]): float(r["value"]) for _, r in grad_df.iterrows()}
                    for e in entries:
                        if e.get("source") == "tensorboard" and e["step"] in grad_by_step:
                            e["grad_norm"] = grad_by_step[e["step"]]

            except Exception as e:
                logger.debug("Failed to parse tensorboard file %s: %s", ef, e)

    except ImportError:
        logger.debug("tbparse not installed — skipping tensorboard parsing. "
                      "Install with: pip install tbparse")

    return entries


def _parse_jsonl(training_path: Path) -> list[dict]:
    """Parse JSON Lines log files."""
    entries: list[dict] = []

    for jsonl_file in training_path.rglob("*.jsonl"):
        try:
            for line in jsonl_file.read_text(errors="ignore").splitlines():
                line = line.strip()
                if not line:
                    continue
                try:
                    row = json.loads(line)
                except json.JSONDecodeError:
                    continue

                step = row.get("step", row.get("global_step", row.get("_step")))
                loss = row.get("loss", row.get("train/loss", row.get("train_loss")))

                if step is not None and loss is not None:
                    entry: dict = {
                        "step": int(step),
                        "loss": float(loss),
                        "source": "jsonl",
                    }
                    grad = row.get("grad_norm", row.get("train/grad_norm"))
                    if grad is not None:
                        entry["grad_norm"] = float(grad)
                    entries.append(entry)
        except OSError:
            continue

    return entries


def _parse_csv_logs(training_path: Path) -> list[dict]:
    """Parse CSV log files with step,loss headers."""
    entries: list[dict] = []

    for csv_file in training_path.rglob("*.csv"):
        try:
            with open(csv_file, newline="") as f:
                reader = csv.DictReader(f)
                if reader.fieldnames is None:
                    continue

                # Detect column names
                step_col = None
                loss_col = None
                grad_col = None

                for col in reader.fieldnames:
                    col_lower = col.lower().strip()
                    if col_lower in ("step", "global_step", "iteration"):
                        step_col = col
                    elif col_lower in ("loss", "train_loss", "train/loss"):
                        loss_col = col
                    elif "grad" in col_lower:
                        grad_col = col

                if step_col is None or loss_col is None:
                    continue

                for row in reader:
                    try:
                        entry: dict = {
                            "step": int(row[step_col]),
                            "loss": float(row[loss_col]),
                            "source": "csv",
                        }
                        if grad_col and row.get(grad_col):
                            entry["grad_norm"] = float(row[grad_col])
                        entries.append(entry)
                    except (ValueError, KeyError):
                        continue
        except OSError:
            continue

    return entries
