"""Typer CLI for llm-refresh-wheel."""
from __future__ import annotations

import json
import sys
from pathlib import Path
from typing import Optional

import typer
from rich.console import Console

from llm_refresh.config import CONFIG_PATH, AppConfig, get_config, reset_config
from llm_refresh.models import BufferStrategy
from llm_refresh.utils import (
    console,
    format_metrics_table,
    format_status_panel,
    print_error,
    print_info,
    print_success,
    print_warning,
)

app = typer.Typer(
    name="refresh-wheel",
    help="Keep your local LLM fresh without forgetting what it already knows.",
    no_args_is_help=True,
)
config_app = typer.Typer(help="Manage llm-refresh-wheel configuration.")
app.add_typer(config_app, name="config")


def _load_jsonl(path: Path) -> list[dict]:
    examples = []
    with path.open("r", encoding="utf-8") as f:
        for lineno, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                examples.append(json.loads(line))
            except json.JSONDecodeError as e:
                print_warning(f"Skipping line {lineno}: {e}")
    return examples


def _get_wheel(
    model: Optional[str] = None,
    strategy: str = "reservoir",
    state_path: Optional[str] = None,
) -> "RefreshWheel":  # type: ignore[name-defined]  # noqa: F821
    from llm_refresh.core import RefreshWheel

    cfg = get_config()
    buf_strategy = BufferStrategy(strategy)
    rw = RefreshWheel(
        model_name=model or cfg.model.name,
        buffer_strategy=buf_strategy,
        state_path=state_path,
    )
    if state_path and Path(state_path).exists():
        rw.load(state_path)
    return rw


# Default state path
DEFAULT_STATE = Path("~/.local/share/llm_refresh/state").expanduser()


@app.command()
def init() -> None:
    """Write default config.toml to ~/.config/llm_refresh/config.toml."""
    cfg = AppConfig()
    cfg.save(CONFIG_PATH)
    print_success(f"Config written to {CONFIG_PATH}")


@app.command()
def add(
    file: Path = typer.Argument(..., help="JSONL file with {\"text\": \"...\"} records"),
    state: str = typer.Option(str(DEFAULT_STATE), "--state", help="State directory"),
) -> None:
    """Add data to the replay buffer."""
    if not file.exists():
        print_error(f"File not found: {file}")
        raise typer.Exit(1)
    examples = _load_jsonl(file)
    if not examples:
        print_warning("No examples loaded.")
        raise typer.Exit(1)

    rw = _get_wheel(state_path=state)
    n = rw.add_data(examples)
    # Flush new data into buffer immediately for persistence
    rw.buffer.add(examples)
    rw.save(state)
    print_success(f"Added {n} examples to buffer. Buffer size: {len(rw.buffer)}")


@app.command()
def anchor(
    file: Path = typer.Argument(..., help="JSONL file for anchor evaluation set"),
    state: str = typer.Option(str(DEFAULT_STATE), "--state", help="State directory"),
) -> None:
    """Set the anchor evaluation dataset."""
    if not file.exists():
        print_error(f"File not found: {file}")
        raise typer.Exit(1)
    examples = _load_jsonl(file)
    if not examples:
        print_warning("No examples loaded.")
        raise typer.Exit(1)

    state_path = Path(state)
    state_path.mkdir(parents=True, exist_ok=True)
    # Save anchor separately
    anchor_path = state_path / "anchor.jsonl"
    with anchor_path.open("w") as f:
        for ex in examples:
            f.write(json.dumps(ex) + "\n")
    print_success(f"Anchor set: {len(examples)} examples → {anchor_path}")


@app.command()
def refresh(
    model: Optional[str] = typer.Option(None, "--model", "-m", help="HuggingFace model name"),
    epochs: int = typer.Option(1, "--epochs", "-e", help="Training epochs"),
    dry_run: bool = typer.Option(False, "--dry-run", help="Skip training, test full flow"),
    state: str = typer.Option(str(DEFAULT_STATE), "--state", help="State directory"),
    strategy: str = typer.Option("reservoir", "--strategy", help="Buffer strategy"),
) -> None:
    """Run one continual learning refresh cycle."""
    state_path = Path(state)
    rw = _get_wheel(model=model, strategy=strategy, state_path=state)

    # Load anchor
    anchor_path = state_path / "anchor.jsonl"
    if anchor_path.exists():
        rw.set_anchor(_load_jsonl(anchor_path))
    elif not dry_run:
        print_error("No anchor set found. Run `refresh-wheel anchor <file.jsonl>` first.")
        raise typer.Exit(1)
    else:
        # Dry-run with dummy anchor
        rw.set_anchor([{"text": "dry run anchor example"}])

    result = rw.refresh(epochs=epochs, dry_run=dry_run)
    rw.save(state)

    if dry_run:
        print_info("Dry-run complete. No model was loaded or trained.")
    else:
        console.print(f"\nBWT={result.bwt:.4f}  KRS={result.krs:.1f}  rolled_back={result.rolled_back}")


@app.command()
def status(
    state: str = typer.Option(str(DEFAULT_STATE), "--state", help="State directory"),
) -> None:
    """Show buffer stats, KRS, and last refresh time."""
    rw = _get_wheel(state_path=state)
    panel = format_status_panel(rw.status())
    console.print(panel)


@app.command()
def metrics(
    state: str = typer.Option(str(DEFAULT_STATE), "--state", help="State directory"),
) -> None:
    """Show full forgetting history as a table."""
    rw = _get_wheel(state_path=state)
    history = [r.to_dict() for r in rw.tracker.history]
    if not history:
        print_info("No refresh cycles recorded yet.")
        return
    table = format_metrics_table(history)
    console.print(table)


@app.command()
def eval(
    state: str = typer.Option(str(DEFAULT_STATE), "--state", help="State directory"),
    model: Optional[str] = typer.Option(None, "--model", "-m", help="Model name override"),
) -> None:
    """Compute perplexity on the anchor set."""
    state_path = Path(state)
    anchor_path = state_path / "anchor.jsonl"
    if not anchor_path.exists():
        print_error("No anchor set. Run `refresh-wheel anchor <file.jsonl>` first.")
        raise typer.Exit(1)

    anchor_data = _load_jsonl(anchor_path)
    rw = _get_wheel(model=model, state_path=state)
    rw.set_anchor(anchor_data)
    rw.load_model()
    result = rw.eval()
    print_success(f"Perplexity: {result.perplexity:.4f}  Loss: {result.loss:.4f}  N={result.dataset_size}")


@app.command()
def schedule(
    every: float = typer.Option(24.0, "--every", help="Interval in hours"),
    state: str = typer.Option(str(DEFAULT_STATE), "--state", help="State directory"),
    model: Optional[str] = typer.Option(None, "--model", "-m"),
) -> None:
    """Start a daemon that runs refresh every N hours."""
    try:
        import schedule as sched_lib
        import time
    except ImportError:
        print_error("Install the 'schedule' library: pip install schedule")
        raise typer.Exit(1)

    print_info(f"Scheduling refresh every {every} hours. Press Ctrl+C to stop.")

    def _run() -> None:
        rw = _get_wheel(model=model, state_path=state)
        state_path_obj = Path(state)
        anchor_path = state_path_obj / "anchor.jsonl"
        if anchor_path.exists():
            rw.set_anchor(_load_jsonl(anchor_path))
        try:
            rw.refresh()
        except Exception as e:
            print_error(f"Refresh failed: {e}")

    sched_lib.every(every).hours.do(_run)
    _run()  # Run immediately on start

    try:
        while True:
            sched_lib.run_pending()
            import time as t
            t.sleep(60)
    except KeyboardInterrupt:
        print_info("Scheduler stopped.")


# -------------------------------------------------------------------------
# Config sub-commands
# -------------------------------------------------------------------------

@config_app.command("show")
def config_show() -> None:
    """Pretty-print current configuration."""
    from rich.syntax import Syntax
    import toml

    cfg = get_config()
    toml_str = toml.dumps(cfg.model_dump())
    console.print(Syntax(toml_str, "toml", theme="monokai"))


@config_app.command("set")
def config_set(
    key: str = typer.Argument(..., help="Dot-notation key, e.g. model.rank"),
    value: str = typer.Argument(..., help="New value"),
) -> None:
    """Set a config value and persist to config.toml."""
    cfg = get_config()
    try:
        cfg.set_value(key, value)
    except KeyError as e:
        print_error(str(e))
        raise typer.Exit(1)
    cfg.save(CONFIG_PATH)
    print_success(f"Set {key} = {value}")
    reset_config()


def main() -> None:
    app()


if __name__ == "__main__":
    main()
