#!/usr/bin/env python3
"""Model selector — choose preferred provider and model at startup.

Run directly:   python model_selector.py
Or via:         python run.py --model-selector
"""

import os
import sys
from pathlib import Path

from dotenv import set_key
from rich import box
from rich.console import Console
from rich.panel import Panel
from rich.rule import Rule
from rich.table import Table

from config import DIRECT_PROVIDER_PREFIXES, _resolve
from run_utils import _load_openswarm_dotenv, _openswarm_state_root

console = Console()
ENV_PATH = _openswarm_state_root() / ".env"

# ── Available models per provider ───────────────────────────────────────────

_PROVIDERS = [
    {
        "id": "deepseek",
        "name": "Deepseek",
        "env_key": "DEEPSEEK_API_KEY",
        "default_model": "deepseek/deepseek-v4-flash",
        "models": [
            ("deepseek/deepseek-v4-flash",   "Deepseek V4 Flash",   "Fast, cost-effective reasoning"),
            ("deepseek/deepseek-v4-pro",      "Deepseek V4 Pro",      "Higher quality reasoning"),
            ("deepseek/deepseek-reasoner",    "Deepseek Reasoner",     "Latest o1-style reasoner"),
        ],
    },
    {
        "id": "minimax",
        "name": "Minimax",
        "env_key": "MINIMAX_API_KEY",
        "default_model": "minimax/MiniMax-M3",
        "models": [
            ("minimax/MiniMax-M3",              "MiniMax-M3",              "Fast, cost-effective reasoning"),
            ("minimax/MiniMax-M2.7",             "MiniMax-M2.7",             "Balanced speed & quality"),
            ("minimax/MiniMax-M2.7-highspeed",   "MiniMax-M2.7-highspeed",   "Fast variant of M2.7"),
            ("minimax/MiniMax-M2.5",             "MiniMax-M2.5",             "Good value mid-tier"),
            ("minimax/MiniMax-M2.5-highspeed",   "MiniMax-M2.5-highspeed",   "Fast variant of M2.5"),
            ("minimax/MiniMax-M2.1",             "MiniMax-M2.1",             "Efficient standard model"),
            ("minimax/MiniMax-M2.1-highspeed",   "MiniMax-M2.1-highspeed",   "Fast variant of M2.1"),
            ("minimax/MiniMax-M2",               "MiniMax-M2",               "Legacy baseline"),
        ],
    },
]


# ── UI helpers (mirrors onboard.py style) ──────────────────────────────────────

try:
    import questionary
    from questionary import Choice, Style as QStyle
    import questionary.prompts.common as _qc_common

    _qc_common.INDICATOR_SELECTED = "✓"
    _HAS_QUESTIONARY = True
except ImportError:
    class Choice:
        def __init__(self, title: object, value: object | None = None) -> None:
            self.title = title
            self.value = title if value is None else value

    _HAS_QUESTIONARY = False


_QSTYLE = None
if _HAS_QUESTIONARY:
    _QSTYLE = QStyle([
        ("qmark",       "fg:#4fc3f7 bold"),
        ("question",    "bold"),
        ("answer",      "fg:#4fc3f7 bold"),
        ("pointer",     "fg:#4fc3f7 bold noreverse"),
        ("highlighted", "noreverse"),
        ("selected",    "fg:#4fc3f7 bold noreverse"),
        ("separator",   "fg:#555555 noreverse"),
        ("instruction", "fg:#555555 italic noreverse"),
        ("text",        "noreverse"),
    ])


def _ask_select(message: str, choices: list) -> object:
    if _HAS_QUESTIONARY:
        return questionary.select(message, choices=choices, style=_QSTYLE).ask()

    titles = [c.title if isinstance(c, Choice) else c for c in choices]
    values = [c.value if isinstance(c, Choice) else c for c in choices]
    console.print(f"\n[bold]{message}[/bold]")
    for i, title in enumerate(titles, 1):
        console.print(f"  [cyan]{i}.[/cyan] {title}")
    while True:
        raw = input("Enter number: ").strip()
        if raw.isdigit() and 1 <= int(raw) <= len(titles):
            return values[int(raw) - 1]
        console.print("[red]Invalid choice, try again.[/red]")


def _ask_confirm(message: str, default: bool = True) -> bool:
    if _HAS_QUESTIONARY:
        return questionary.confirm(message, default=default, style=_QSTYLE).ask()
    prompt = f"{message} [{'Y/n' if default else 'y/N'}]: "
    raw = input(prompt).strip().lower()
    return default if not raw else raw in ("y", "yes")


# ── availability check ─────────────────────────────────────────────────────────

def _provider_available(env_key: str) -> bool:
    _load_openswarm_dotenv()
    return bool(os.getenv(env_key))


# ── main selector ──────────────────────────────────────────────────────────────

def run_model_selector(*, force: bool = False) -> str | None:
    """Show the model selector UI and return the chosen model, or None if cancelled."""
    _load_openswarm_dotenv()

    existing_model = os.getenv("DEFAULT_MODEL", "")

    console.print()
    console.print(Panel.fit(
        "[bold cyan]SwarmAgency[/bold cyan]  [dim]—  model selector[/dim]\n"
        "[dim]Choose your preferred AI provider and model.[/dim]",
        border_style="cyan",
        padding=(1, 4),
    ))

    # ── provider ───────────────────────────────────────────────────────────────
    available_providers = [
        p for p in _PROVIDERS if _provider_available(p["env_key"])
    ]

    if not available_providers:
        console.print(Panel(
            "[red]No API keys found.[/red]\n\n"
            "Run [bold]python onboard.py[/bold] first to configure your API keys.\n"
            "Then run [bold]python model_selector.py[/bold] to select a model.",
            border_style="red",
            padding=(1, 3),
        ))
        return None

    if len(available_providers) == 1 and not force:
        chosen_provider = available_providers[0]
        console.print(f"\n[dim]Only one provider configured:[/dim] [cyan]{chosen_provider['name']}[/cyan]")
    else:
        console.print()
        console.print(Rule("[bold]Step 1  ·  Provider[/bold]", style="cyan"))

        provider_choices = [
            Choice(title=p["name"], value=p) for p in available_providers
        ]
        chosen_provider = _ask_select("Select provider:", provider_choices)

    # ── model ─────────────────────────────────────────────────────────────────
    console.print()
    console.print(Rule(f"[bold]Step 2  ·  Model  ({chosen_provider['name']})[/bold]", style="cyan"))

    model_choices = [
        Choice(title=model_name, value=model_id)
        for model_id, model_name, _ in chosen_provider["models"]
    ]

    # Pre-select current model if it matches the provider
    current_model = None
    if existing_model.startswith(chosen_provider["id"] + "/"):
        current_model = existing_model

    if current_model:
        console.print(f"  [dim]Current:[/dim] [cyan]{existing_model}[/cyan]")
        if _ask_confirm("  Keep current model?", default=True):
            final_model = current_model
        else:
            final_model = _ask_select("Select model:", model_choices)
    else:
        final_model = _ask_select("Select model:", model_choices)

    # ── write to .env ──────────────────────────────────────────────────────────
    _write_env({"DEFAULT_MODEL": final_model})

    # ── summary ───────────────────────────────────────────────────────────────
    console.print()
    console.print(Rule("[bold green]Model saved[/bold green]", style="green"))
    table = Table(box=box.SIMPLE, show_header=False, padding=(0, 2))
    table.add_column(style="dim", no_wrap=True)
    table.add_column()
    table.add_row("Provider", f"[cyan]{chosen_provider['name']}[/cyan]")
    table.add_row("Model",    f"[cyan]{final_model}[/cyan]")
    table.add_row(".env",     f"[cyan]{ENV_PATH}[/cyan]")
    console.print(table)
    console.print()

    return final_model


def _write_env(updates: dict) -> None:
    ENV_PATH.parent.mkdir(parents=True, exist_ok=True)
    if not ENV_PATH.exists():
        ENV_PATH.write_text("", encoding="utf-8")
    for key, value in updates.items():
        if value:
            set_key(str(ENV_PATH), key, value)


# ── entry point ───────────────────────────────────────────────────────────────

if __name__ == "__main__":
    try:
        result = run_model_selector(force=True)
        if result:
            console.print("[dim]Restart with:[/dim] python run.py")
        else:
            console.print("\n[dim]Model selection cancelled.[/dim]")
            sys.exit(0)
    except KeyboardInterrupt:
        console.print("\n\n[dim]Model selection cancelled.[/dim]\n")
        sys.exit(0)