from __future__ import annotations

import asyncio
import os
from pathlib import Path
from typing import Annotated

import httpx
import torch
import typer
import uvicorn
from dotenv import load_dotenv
from rich.console import Console
from rich.table import Table

from .api import create_app
from .config import (
    load_es_config,
    load_rl_config,
    load_sft_config,
    load_worker_pool,
)
from .evaluate import evaluate_checkpoint
from .generate import generate_reward_dataset
from .io import load_rewards, load_tasks
from .model import RouterModel
from .orchestrator import FuguLiteOrchestrator
from .train_es import train_es
from .train_rl import train_rl
from .train_sft import train_sft
from .training_common import resolve_device

app = typer.Typer(no_args_is_help=True, help="Train and serve a compact Fugu-inspired router.")
console = Console()


def _load_environment() -> None:
    load_dotenv(override=False)


@app.command()
def doctor() -> None:
    """Check the local GPU and secret configuration without exposing the key."""
    _load_environment()
    table = Table(title="Fugu-Lite environment")
    table.add_column("Check")
    table.add_column("Value")
    table.add_row("Python/PyTorch", torch.__version__)
    table.add_row("CUDA available", str(torch.cuda.is_available()))
    if torch.cuda.is_available():
        properties = torch.cuda.get_device_properties(0)
        table.add_row("GPU", properties.name)
        table.add_row("VRAM", f"{properties.total_memory / 2**30:.1f} GiB")
        table.add_row("BF16", str(torch.cuda.is_bf16_supported()))
    table.add_row("OPENROUTER_API_KEY", "configured" if os.getenv("OPENROUTER_API_KEY") else "missing")
    console.print(table)


@app.command("list-models")
def list_models(
    contains: Annotated[str, typer.Option(help="Case-insensitive model ID filter.")] = "",
    limit: Annotated[int, typer.Option(min=1, max=500)] = 50,
) -> None:
    """List model slugs currently returned by OpenRouter."""
    _load_environment()
    key = os.getenv("OPENROUTER_API_KEY")
    headers = {"Authorization": f"Bearer {key}"} if key else {}
    response = httpx.get("https://openrouter.ai/api/v1/models", headers=headers, timeout=30)
    response.raise_for_status()
    models = response.json().get("data", [])
    needle = contains.casefold()
    ids = [item["id"] for item in models if needle in item.get("id", "").casefold()]
    for model_id in ids[:limit]:
        console.print(model_id)
    console.print(f"Shown {min(limit, len(ids))} of {len(ids)} matches")


@app.command()
def generate(
    workers: Annotated[Path, typer.Option(exists=True, dir_okay=False)],
    tasks: Annotated[Path, typer.Option(exists=True, dir_okay=False)],
    output: Annotated[Path, typer.Option()],
    limit: Annotated[int | None, typer.Option(min=1)] = None,
    repetitions: Annotated[
        int,
        typer.Option(
            min=1,
            help="Independent calls per worker/task. Rewards are averaged."
        ),
    ] = 1,
    resume: Annotated[
        bool,
        typer.Option(help="Skip task IDs already saved in output.")
    ] = False,
) -> None:
    """Call every worker for every task and create the offline reward matrix."""

    _load_environment()
    config = load_worker_pool(workers)

    asyncio.run(
        generate_reward_dataset(
            load_tasks(tasks),
            config,
            output,
            limit=limit,
            resume=resume,
            repetitions=repetitions,
        )
    )

@app.command("train-sft")
def train_sft_command(
    config: Annotated[Path, typer.Option(exists=True, dir_okay=False)],
    data: Annotated[Path, typer.Option(exists=True, dir_okay=False)],
    output: Annotated[Path, typer.Option()],
    checkpoint: Annotated[Path | None, typer.Option(exists=True, file_okay=False)] = None,
) -> None:
    """Warm-start the routing head with soft utility targets."""
    report = train_sft(
        load_rewards(data), load_sft_config(config), output, str(checkpoint) if checkpoint else None
    )
    console.print_json(data=report)


@app.command("train-rl")
def train_rl_command(
    config: Annotated[Path, typer.Option(exists=True, dir_okay=False)],
    data: Annotated[Path, typer.Option(exists=True, dir_okay=False)],
    output: Annotated[Path, typer.Option()],
    checkpoint: Annotated[Path | None, typer.Option(exists=True, file_okay=False)] = None,
) -> None:
    """Optimize worker selection with offline contextual-bandit RL."""
    report = train_rl(
        load_rewards(data), load_rl_config(config), output, str(checkpoint) if checkpoint else None
    )
    console.print_json(data=report)


@app.command("train-es")
def train_es_command(
    config: Annotated[Path, typer.Option(exists=True, dir_okay=False)],
    data: Annotated[Path, typer.Option(exists=True, dir_okay=False)],
    output: Annotated[Path, typer.Option()],
    checkpoint: Annotated[Path | None, typer.Option(exists=True, file_okay=False)] = None,
) -> None:
    """Run Sakana-style separable CMA-ES on the cached router features."""
    report = train_es(
        load_rewards(data), load_es_config(config), output, str(checkpoint) if checkpoint else None
    )
    console.print_json(data=report)


@app.command()
def evaluate(
    checkpoint: Annotated[Path, typer.Option(exists=True, file_okay=False)],
    data: Annotated[Path, typer.Option(exists=True, dir_okay=False)],
    split: Annotated[str, typer.Option()] = "test",
    output: Annotated[Path | None, typer.Option()] = None,
) -> None:
    """Compare the router with oracle, best-fixed-worker, and random baselines."""
    report = evaluate_checkpoint(load_rewards(data), checkpoint, split, output)
    console.print_json(data=report)


@app.command()
def route(
    checkpoint: Annotated[Path, typer.Option(exists=True, file_okay=False)],
    prompt: Annotated[str, typer.Option()],
    domain: Annotated[str, typer.Option()] = "general",
) -> None:
    """Run only the local routing head; this spends no API credits."""
    model = RouterModel.from_checkpoint(checkpoint, device=resolve_device())
    model.eval()
    from .schemas import route_text

    encoded = model.tokenizer(
        [route_text(prompt, domain)],
        return_tensors="pt",
        truncation=True,
        max_length=model.router_config.max_length,
    )
    with torch.no_grad():
        logits = model(
            encoded["input_ids"].to(model.device_ref),
            encoded["attention_mask"].to(model.device_ref),
        )[0]
    probabilities = torch.softmax(logits, dim=-1).cpu().tolist()
    console.print_json(
        data={
            "selected_worker": model.worker_ids[max(range(len(probabilities)), key=probabilities.__getitem__)],
            "probabilities": dict(zip(model.worker_ids, probabilities)),
        }
    )


@app.command()
def ask(
    checkpoint: Annotated[Path, typer.Option(exists=True, file_okay=False)],
    workers: Annotated[Path, typer.Option(exists=True, dir_okay=False)],
    prompt: Annotated[str, typer.Option()],
    domain: Annotated[str, typer.Option()] = "general",
) -> None:
    """Route locally, then call the selected OpenRouter/local worker."""
    _load_environment()

    async def run():
        orchestrator = FuguLiteOrchestrator(checkpoint, workers)
        try:
            return await orchestrator.answer(prompt, domain)
        finally:
            await orchestrator.close()

    console.print_json(data=asyncio.run(run()))


@app.command()
def serve(
    checkpoint: Annotated[Path, typer.Option(exists=True, file_okay=False)],
    workers: Annotated[Path, typer.Option(exists=True, dir_okay=False)],
    host: Annotated[str, typer.Option()] = "127.0.0.1",
    port: Annotated[int, typer.Option(min=1, max=65535)] = 8080,
) -> None:
    """Serve one OpenAI-compatible endpoint backed by the learned router."""
    _load_environment()
    uvicorn.run(create_app(str(checkpoint), str(workers)), host=host, port=port)


if __name__ == "__main__":
    app()