"""Yarn CLI — command-line interface for the Yarn sovereign AI platform."""
from __future__ import annotations

import json
import os
import sys
from pathlib import Path

import click
import requests as _requests

CONFIG_DIR = Path.home() / ".config" / "yarn"
CONFIG_FILE = CONFIG_DIR / "config.json"

# Keycloak OIDC defaults for password grant (public client — no secret needed)
_KEYCLOAK_TOKEN_URL = "https://id.au.yarn.prosodylabs.com.au/realms/sovereign/protocol/openid-connect/token"
_KEYCLOAK_CLIENT_ID = "yarn-cli"
_ACCOUNT_PORTAL_URL = "https://api.yarn.prosodylabs.com.au"


def _load_config() -> dict:
    if CONFIG_FILE.exists():
        return json.loads(CONFIG_FILE.read_text())
    return {}


def _save_config(cfg: dict) -> None:
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    CONFIG_FILE.write_text(json.dumps(cfg, indent=2) + "\n")


def _get_client():
    from .client import Client

    cfg = _load_config()
    api_key = os.environ.get("YARN_API_KEY") or cfg.get("api_key")
    base_url = os.environ.get("YARN_BASE_URL") or cfg.get("base_url")
    kwargs: dict = {}
    if api_key:
        kwargs["api_key"] = api_key
    if base_url:
        kwargs["base_url"] = base_url
    return Client(**kwargs)


def _table(rows: list[list[str]], headers: list[str]) -> str:
    """Simple table formatter."""
    all_rows = [headers] + rows
    widths = [max(len(str(cell)) for cell in col) for col in zip(*all_rows)]
    fmt = "  ".join(f"{{:<{w}}}" for w in widths)
    lines = [fmt.format(*headers), fmt.format(*("-" * w for w in widths))]
    for row in rows:
        lines.append(fmt.format(*row))
    return "\n".join(lines)


def _json_out(data) -> None:
    click.echo(json.dumps(data, indent=2, default=str))


# ---------------------------------------------------------------------------
# Root group
# ---------------------------------------------------------------------------

@click.group()
@click.version_option(package_name="yarn-au")
def main():
    """Yarn — sovereign AI platform CLI."""
    # Name clash guard: warn if JS yarn package manager shadows us
    if os.environ.get("YARN_NAME_CHECK") != "0":
        import shutil
        import subprocess
        js_yarn = shutil.which("yarn")
        if js_yarn and js_yarn != shutil.which("yarn-au"):
            try:
                out = subprocess.run(
                    [js_yarn, "--version"], capture_output=True, text=True, timeout=2,
                )
                # JS yarn returns a semver like "1.22.19" or "4.1.0"
                if out.returncode == 0 and out.stdout.strip()[:1].isdigit():
                    click.echo(
                        "Warning: JS package manager 'yarn' detected on PATH. "
                        "Use 'yarn-au' to avoid conflict, or set YARN_NAME_CHECK=0 to suppress.",
                        err=True,
                    )
            except Exception:
                pass


# ---------------------------------------------------------------------------
# Auth: login / logout / status / whoami
# ---------------------------------------------------------------------------

@main.group()
def auth():
    """Authentication — login, logout, API key management."""
    pass


@auth.command("login")
@click.option("--email", "-e", default=None, help="Email address")
@click.option("--password", "-p", default=None, help="Password (prompted if not provided)")
@click.option("--api-key", default=None, help="Use an existing API key directly")
@click.option("--base-url", default=None, help="Override API base URL")
@click.option("--key-name", default="yarn-cli", help="Name for the generated API key")
def auth_login(email: str | None, password: str | None, api_key: str | None,
               base_url: str | None, key_name: str):
    """Log in to Yarn.

    \b
    With --email/--password: authenticates via Keycloak, creates an API key.
    With --api-key: saves an existing key directly.

    \b
    Examples:
      yarn auth login                           # interactive prompts
      yarn auth login -e user@example.com       # prompt for password only
      yarn auth login --api-key yarn_abc123     # direct key
    """
    if api_key:
        cfg = _load_config()
        cfg["api_key"] = api_key
        if base_url:
            cfg["base_url"] = base_url
        _save_config(cfg)
        click.echo(f"Saved to {CONFIG_FILE}")
        return

    # Full Keycloak flow
    if not email:
        email = click.prompt("Email")
    if not password:
        password = click.prompt("Password", hide_input=True)

    # Step 1: Get JWT from Keycloak
    try:
        resp = _requests.post(_KEYCLOAK_TOKEN_URL, data={
            "grant_type": "password",
            "client_id": _KEYCLOAK_CLIENT_ID,
            "username": email,
            "password": password,
        }, timeout=15)
    except _requests.RequestException as exc:
        click.echo(f"Connection error: {exc}", err=True)
        sys.exit(1)

    if resp.status_code != 200:
        detail = resp.json().get("error_description", resp.text) if resp.headers.get("content-type", "").startswith("application/json") else resp.text
        click.echo(f"Authentication failed: {detail}", err=True)
        sys.exit(1)

    jwt_token = resp.json()["access_token"]

    # Step 2: Create API key via account portal
    portal_url = base_url.rstrip("/") if base_url else _ACCOUNT_PORTAL_URL
    try:
        resp = _requests.post(
            f"{portal_url}/v1/users/me/api-keys",
            headers={"Authorization": f"Bearer {jwt_token}", "Content-Type": "application/json"},
            json={"name": key_name},
            timeout=15,
        )
    except _requests.RequestException as exc:
        click.echo(f"Connection error: {exc}", err=True)
        sys.exit(1)

    if resp.status_code not in (200, 201):
        click.echo(f"Failed to create API key: {resp.text}", err=True)
        sys.exit(1)

    new_key = resp.json()["api_key"]

    # Step 3: Save config
    cfg = _load_config()
    cfg["api_key"] = new_key
    if base_url:
        cfg["base_url"] = base_url
    _save_config(cfg)

    click.echo(f"Logged in as {email}")
    click.echo(f"API key saved to {CONFIG_FILE}")


@auth.command("logout")
def auth_logout():
    """Remove saved credentials."""
    if CONFIG_FILE.exists():
        cfg = _load_config()
        cfg.pop("api_key", None)
        if cfg:
            _save_config(cfg)
        else:
            CONFIG_FILE.unlink()
        click.echo("Logged out.")
    else:
        click.echo("Not logged in.")


@auth.command("status")
def auth_status():
    """Check authentication status."""
    cfg = _load_config()
    api_key = os.environ.get("YARN_API_KEY") or cfg.get("api_key")
    if not api_key:
        click.echo("Not logged in. Run: yarn auth login")
        sys.exit(1)
    source = "YARN_API_KEY env" if os.environ.get("YARN_API_KEY") else str(CONFIG_FILE)
    click.echo(f"Key:    {api_key[:10]}...{api_key[-4:]}")
    click.echo(f"Source: {source}")


# Backward compat: keep top-level login/whoami
@main.command()
@click.option("--api-key", prompt="API key", hide_input=True, help="Your yarn_... API key")
@click.option("--base-url", default=None, help="Override API base URL")
def login(api_key: str, base_url: str | None):
    """Configure API key (use 'yarn auth login' for full flow)."""
    cfg = _load_config()
    cfg["api_key"] = api_key
    if base_url:
        cfg["base_url"] = base_url
    _save_config(cfg)
    click.echo(f"Saved to {CONFIG_FILE}")


@main.command()
def whoami():
    """Show current user info."""
    cfg = _load_config()
    api_key = os.environ.get("YARN_API_KEY") or cfg.get("api_key")
    if not api_key:
        click.echo("Not logged in. Run: yarn auth login", err=True)
        sys.exit(1)
    try:
        resp = _requests.get(
            f"{_ACCOUNT_PORTAL_URL}/v1/auth/validate-api-key",
            headers={"X-API-Key": api_key},
            timeout=15,
        )
        resp.raise_for_status()
        data = resp.json()
        click.echo(f"User:  {data.get('email', 'unknown')}")
        click.echo(f"Name:  {data.get('name', 'unknown')}")
    except Exception as e:
        click.echo(f"Error: {e}", err=True)
        sys.exit(1)


# ---------------------------------------------------------------------------
# GPUs
# ---------------------------------------------------------------------------

@main.command()
@click.option("--json", "as_json", is_flag=True, help="Output raw JSON")
@click.option("--sovereign", is_flag=True, help="Only show Australian-sovereign GPUs")
@click.option("--provider", "provider_filter", default=None, help="Filter by provider (e.g. lambdalabs, runpod)")
@click.option("--min-vram", type=int, default=None, help="Minimum VRAM per GPU in GB")
def gpus(as_json: bool, sovereign: bool, provider_filter: str | None, min_vram: int | None):
    """List available GPUs across all providers with pricing.

    Shows platform (local), cloud, and BYO GPUs sorted by price.
    Use --sovereign to restrict to Australian data residency.
    """
    client = _get_client()
    data = client.gpus()

    if as_json:
        _json_out(data)
        return

    all_gpus = data.get("gpus", [])

    # Apply filters
    if sovereign:
        all_gpus = [g for g in all_gpus if g.get("location") == "sovereign"
                     or g.get("compliance_tier") in ("sovereign", "sovereign-au")]
    if provider_filter:
        pf = provider_filter.lower()
        all_gpus = [g for g in all_gpus if g.get("provider", "").lower() == pf
                     or (g.get("source") == "platform" and pf in ("platform", "local"))
                     or (g.get("source") == "byo" and pf == "byo")]
    if min_vram:
        all_gpus = [g for g in all_gpus if g.get("gpu_memory_gb", 0) >= min_vram]

    # Filter out CPU-only instances
    all_gpus = [g for g in all_gpus if g.get("gpu_count", 0) > 0
                and "cpu" not in g.get("gpu_type", "").lower()]

    # Group by source
    platform = [g for g in all_gpus if g.get("source") == "platform"]
    cloud = [g for g in all_gpus if g.get("source") == "cloud"]
    byo = [g for g in all_gpus if g.get("source") == "byo"]

    # Sort cloud by price (cheapest first)
    cloud.sort(key=lambda g: g.get("price_per_hour_usd") or float("inf"))

    any_shown = False

    # Platform GPUs
    if platform:
        any_shown = True
        click.echo("Platform GPUs")
        rows = []
        for g in platform:
            vram = ""
            if "vram" in g:
                v = g["vram"]
                vram = f"{v.get('available_gb', '?')}/{v.get('total_gb', '?')}GB"
            else:
                vram = f"{g.get('gpu_memory_gb', '?')}GB"
            models = ", ".join(
                dm["model"] for dm in g.get("deployed_models", []) if dm.get("ready")
            ) or "-"
            rows.append([g.get("gpu_type", "?"), vram, g.get("status", "?"),
                         str(g.get("active_workloads", 0)), models])
        click.echo(_table(rows, ["GPU", "VRAM", "Status", "Workloads", "Models Loaded"]))

    # Cloud GPUs
    if cloud:
        if any_shown:
            click.echo()
        any_shown = True
        click.echo("Cloud GPUs (sorted by price)")
        rows = []
        for g in cloud:
            gpu_count = g.get("gpu_count", 1)
            gpu_label = g.get("gpu_type", "?")
            if gpu_count > 1 and not gpu_label.startswith(f"{gpu_count}x"):
                gpu_label = f"{gpu_count}x {gpu_label}"
            vram = g.get("gpu_memory_gb", 0)
            vram_str = f"{vram}GB" if vram else "-"
            price = g.get("price_per_hour_usd")
            price_str = f"${price:.2f}/hr" if price is not None else "-"
            regions = ", ".join(g.get("regions", [])) or "-"
            rows.append([gpu_label, vram_str, g.get("provider", "?"), price_str,
                         g.get("status", "?"), regions])
        click.echo(_table(rows, ["GPU", "VRAM", "Provider", "Price (USD)", "Status", "Regions"]))

    # BYO GPUs
    if byo:
        if any_shown:
            click.echo()
        any_shown = True
        click.echo("Your GPUs")
        rows = []
        for g in byo:
            vram = f"{g.get('gpu_memory_gb', '?')}GB"
            rows.append([g.get("gpu_type", "?"), vram, g.get("org_name", "-"),
                         g.get("status", "?")])
        click.echo(_table(rows, ["GPU", "VRAM", "Organization", "Status"]))

    if not any_shown:
        if sovereign:
            click.echo("No sovereign GPUs available. Try without --sovereign to see all providers.")
        else:
            click.echo("No GPUs available.")
        return

    # Models
    models = data.get("models", [])
    if models:
        click.echo()
        model_rows = [
            [m.get("model_id", "?"), f"{m.get('min_gpu_memory_gb', '?')}GB",
             f"{m.get('parameter_count_billions', '?')}B" if m.get("parameter_count_billions") else "-"]
            for m in models
        ]
        click.echo(_table(model_rows, ["Model", "Min VRAM", "Params"]))

    # Pricing
    pricing = data.get("pricing", {})
    model_prices = pricing.get("models", {})
    if model_prices:
        click.echo()
        price_rows = [
            [mid, f"${p.get('input_per_1m', '?')}", f"${p.get('output_per_1m', '?')}"]
            for mid, p in model_prices.items()
        ]
        default = pricing.get("default", {})
        price_rows.append(["(default)", f"${default.get('input_per_1m', '?')}", f"${default.get('output_per_1m', '?')}"])
        click.echo(_table(price_rows, ["Model", "Input/1M tok (AUD)", "Output/1M tok (AUD)"]))


# ---------------------------------------------------------------------------
# Chat (inference)
# ---------------------------------------------------------------------------

@main.command()
@click.argument("message")
@click.option("--model", "-m", default=None, help="Model ID")
@click.option("--stream", "-s", is_flag=True, help="Stream response tokens")
@click.option("--temperature", "-t", type=float, default=None)
@click.option("--max-tokens", type=int, default=None)
@click.option("--system", "system_prompt", default=None, help="System prompt")
def chat(message: str, model: str | None, stream: bool, temperature: float | None,
         max_tokens: int | None, system_prompt: str | None):
    """Send a chat message. Use quotes for multi-word messages."""
    client = _get_client()

    if not model:
        # Pick first available model
        try:
            models = client.chat.models()
            if models:
                model = models[0].get("id") or models[0].get("model_id", "")
        except Exception:
            pass
    if not model:
        click.echo("No model specified and none discovered. Use --model.", err=True)
        sys.exit(1)

    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": message})

    kwargs: dict = {}
    if temperature is not None:
        kwargs["temperature"] = temperature
    if max_tokens is not None:
        kwargs["max_tokens"] = max_tokens

    result = client.chat.completions(model, messages, stream=stream, **kwargs)

    if stream:
        for chunk in result:
            delta = chunk.get("choices", [{}])[0].get("delta", {})
            text = delta.get("content", "")
            if text:
                click.echo(text, nl=False)
        click.echo()
    else:
        text = result.get("choices", [{}])[0].get("message", {}).get("content", "")
        click.echo(text)


# ---------------------------------------------------------------------------
# Embeddings
# ---------------------------------------------------------------------------

@main.command()
@click.argument("text")
@click.option("--model", "-m", default=None, help="Embedding model ID")
def embed(text: str, model: str | None):
    """Generate embeddings for text."""
    client = _get_client()
    if not model:
        click.echo("--model required for embeddings", err=True)
        sys.exit(1)
    result = client.embeddings.create(model, text)
    _json_out(result)


# ---------------------------------------------------------------------------
# Jobs
# ---------------------------------------------------------------------------

@main.group()
def job():
    """Training job management."""
    pass


@job.command("submit")
@click.argument("script_or_dir")
@click.option("--name", "-n", default=None, help="Job name")
@click.option("--gpu", default="rtx-4090", help="GPU type")
@click.option("--gpu-count", type=int, default=1)
@click.option("--pip", "pip_packages", multiple=True, help="pip packages to install")
@click.option("--env", "env_vars", multiple=True, help="ENV=VALUE pairs")
@click.option("--secret", "secret_names", multiple=True, help="Secret names to mount")
@click.option("--description", "-d", default="")
@click.option("--max-hours", type=float, default=0)
@click.option("--image", default="")
@click.option("--entrypoint", default="")
@click.option("--compute-target", default=None, help="Compute target (hub or spoke ID)")
@click.option("--provider", default=None, help="Preferred GPU provider (e.g. lambdalabs, runpod)")
@click.option("--region", default=None, help="Preferred region")
@click.option("--sovereignty", is_flag=True, help="Restrict to Australian-sovereign GPUs")
@click.option("--min-vram", type=int, default=None, help="Minimum VRAM per GPU in GB")
def job_submit(script_or_dir: str, name: str | None, gpu: str, gpu_count: int,
               pip_packages: tuple, env_vars: tuple, secret_names: tuple,
               description: str, max_hours: float, image: str, entrypoint: str,
               compute_target: str | None, provider: str | None, region: str | None,
               sovereignty: bool, min_vram: int | None):
    """Submit a training job (script file or directory)."""
    client = _get_client()
    path = Path(script_or_dir)

    kwargs: dict = {
        "gpu": gpu,
        "gpu_count": gpu_count,
        "description": description,
    }
    if name:
        kwargs["name"] = name
    else:
        kwargs["name"] = path.stem
    if pip_packages:
        kwargs["pip_packages"] = list(pip_packages)
    if env_vars:
        kwargs["env_vars"] = dict(kv.split("=", 1) for kv in env_vars)
    if secret_names:
        kwargs["secret_names"] = list(secret_names)
    if max_hours:
        kwargs["max_runtime_hours"] = max_hours
    if image:
        kwargs["image"] = image
    if entrypoint:
        kwargs["entrypoint"] = entrypoint
    if compute_target:
        kwargs["compute_target"] = compute_target
    if provider:
        kwargs["provider"] = provider
    if region:
        kwargs["region"] = region
    if sovereignty:
        kwargs["sovereignty"] = True
    if min_vram:
        kwargs["min_vram"] = min_vram

    if path.is_dir():
        kwargs["directory"] = str(path)
    elif path.is_file():
        kwargs["script"] = str(path)
    else:
        click.echo(f"Not found: {path}", err=True)
        sys.exit(1)

    result = client.jobs.submit(**kwargs)
    click.echo(f"Job submitted: {result.get('id', result.get('job_id', '?'))}")
    click.echo(f"Status: {result.get('status', '?')}")


@job.command("list")
@click.option("--json", "as_json", is_flag=True)
def job_list(as_json: bool):
    """List all jobs."""
    client = _get_client()
    jobs = client.jobs.list()
    if as_json:
        _json_out(jobs)
        return
    if not jobs:
        click.echo("No jobs.")
        return
    rows = [
        [j.get("id", j.get("job_id", "?"))[:12], j.get("name", "-"),
         j.get("status", "?"), j.get("gpu_type", "-"), j.get("created_at", "-")[:19]]
        for j in jobs
    ]
    click.echo(_table(rows, ["ID", "Name", "Status", "GPU", "Created"]))


@job.command("status")
@click.argument("job_id")
def job_status(job_id: str):
    """Get job status."""
    client = _get_client()
    result = client.jobs.get(job_id)
    click.echo(f"Job:    {result.get('id', result.get('job_id', '?'))}")
    click.echo(f"Name:   {result.get('name', '-')}")
    click.echo(f"Status: {result.get('status', '?')}")
    click.echo(f"GPU:    {result.get('gpu_type', '-')}")
    if result.get("created_at"):
        click.echo(f"Created: {result['created_at']}")


@job.command("logs")
@click.argument("job_id")
@click.option("--follow", "-f", is_flag=True, help="Stream logs until job completes")
@click.option("--tail", type=int, default=500)
def job_logs(job_id: str, follow: bool, tail: int):
    """Fetch job logs."""
    client = _get_client()
    if follow:
        for line in client.jobs.stream_logs(job_id, tail=tail):
            click.echo(line)
    else:
        result = client.jobs.logs(job_id, tail=tail)
        for line in result.get("logs", []):
            click.echo(line)


@job.command("cancel")
@click.argument("job_id")
def job_cancel(job_id: str):
    """Cancel a running job."""
    client = _get_client()
    client.jobs.cancel(job_id)
    click.echo(f"Cancelled: {job_id}")


@job.command("estimate")
@click.option("--gpu", default="rtx-4090")
@click.option("--gpu-count", type=int, default=1)
@click.option("--hours", type=float, default=4)
def job_estimate(gpu: str, gpu_count: int, hours: float):
    """Estimate job cost."""
    client = _get_client()
    result = client.jobs.estimate(gpu=gpu, gpu_count=gpu_count, max_runtime_hours=hours)
    click.echo(f"Estimated cost: ${result.get('estimated_cost_aud', '?')} AUD")
    click.echo(f"Rate: ${result.get('rate_per_gpu_hour_aud', '?')}/GPU-hr")


# ---------------------------------------------------------------------------
# Sessions
# ---------------------------------------------------------------------------

@main.group()
def session():
    """Interactive GPU session management."""
    pass


@session.command("start")
@click.option("--name", "-n", default="default")
@click.option("--gpu", default="rtx-4090")
@click.option("--gpu-count", type=int, default=1)
@click.option("--timeout", "idle_timeout", type=int, default=60, help="Idle timeout in minutes")
@click.option("--image", default="")
@click.option("--compute-target", default=None, help="Compute target (hub or spoke ID)")
@click.option("--provider", default=None, help="Preferred GPU provider (e.g. lambdalabs, runpod)")
@click.option("--region", default=None, help="Preferred region")
@click.option("--sovereignty", is_flag=True, help="Restrict to Australian-sovereign GPUs")
@click.option("--min-vram", type=int, default=None, help="Minimum VRAM per GPU in GB")
def session_start(name: str, gpu: str, gpu_count: int, idle_timeout: int,
                  image: str, compute_target: str | None, provider: str | None,
                  region: str | None, sovereignty: bool, min_vram: int | None):
    """Start an interactive GPU session."""
    client = _get_client()
    kwargs: dict = {"gpu": gpu, "gpu_count": gpu_count, "idle_timeout_minutes": idle_timeout}
    if image:
        kwargs["image"] = image
    if compute_target:
        kwargs["compute_target"] = compute_target
    if provider:
        kwargs["provider"] = provider
    if region:
        kwargs["region"] = region
    if sovereignty:
        kwargs["sovereignty"] = True
    if min_vram:
        kwargs["min_vram"] = min_vram
    sess = client.sessions.create(name, **kwargs)
    click.echo(f"Session: {sess.id}")
    click.echo(f"Status:  {sess.status}")
    if sess.dashboard_url:
        click.echo(f"Dashboard: {sess.dashboard_url}")


@session.command("list")
@click.option("--json", "as_json", is_flag=True)
def session_list(as_json: bool):
    """List sessions."""
    client = _get_client()
    sessions = client.sessions.list()
    if as_json:
        _json_out(sessions)
        return
    if not sessions:
        click.echo("No sessions.")
        return
    rows = [
        [s.get("id", "?")[:12], s.get("name", "-"), s.get("status", "?"),
         s.get("gpu_type", "-"), s.get("created_at", "-")[:19]]
        for s in sessions
    ]
    click.echo(_table(rows, ["ID", "Name", "Status", "GPU", "Created"]))


@session.command("status")
@click.argument("session_id")
def session_status(session_id: str):
    """Get session details."""
    client = _get_client()
    sess = client.sessions.get(session_id)
    click.echo(f"Session:   {sess.id}")
    click.echo(f"Name:      {sess.name}")
    click.echo(f"Status:    {sess.status}")
    click.echo(f"GPU:       {sess.gpu_type}")
    if sess.dashboard_url:
        click.echo(f"Dashboard: {sess.dashboard_url}")
    if sess.ray_address:
        click.echo(f"Ray:       {sess.ray_address}")


@session.command("stop")
@click.argument("session_id")
def session_stop(session_id: str):
    """Stop a session."""
    client = _get_client()
    client.sessions.delete(session_id)
    click.echo(f"Stopped: {session_id}")


@session.command("extend")
@click.argument("session_id")
@click.option("--minutes", type=int, default=60, help="Minutes to add")
def session_extend(session_id: str, minutes: int):
    """Extend session timeout."""
    client = _get_client()
    result = client._request("PATCH", f"/v1/research/sessions/{session_id}/extend",
                             json={"additional_minutes": minutes})
    click.echo(f"Extended by {minutes}min. New timeout: {result.get('new_timeout_minutes', '?')}min")


# ---------------------------------------------------------------------------
# Notebooks
# ---------------------------------------------------------------------------

@main.group()
def notebook():
    """JupyterLab notebook management."""
    pass


@notebook.command("start")
@click.option("--name", "-n", default="notebook")
@click.option("--gpu", default="rtx-4090")
@click.option("--gpu-count", type=int, default=1)
@click.option("--timeout", "idle_timeout", type=int, default=120, help="Idle timeout in minutes")
@click.option("--image", default="")
@click.option("--compute-target", default=None, help="Compute target (hub or spoke ID)")
@click.option("--provider", default=None, help="Preferred GPU provider (e.g. lambdalabs, runpod)")
@click.option("--region", default=None, help="Preferred region")
@click.option("--sovereignty", is_flag=True, help="Restrict to Australian-sovereign GPUs")
@click.option("--min-vram", type=int, default=None, help="Minimum VRAM per GPU in GB")
def notebook_start(name: str, gpu: str, gpu_count: int, idle_timeout: int, image: str,
                   compute_target: str | None, provider: str | None, region: str | None,
                   sovereignty: bool, min_vram: int | None):
    """Start a JupyterLab notebook with GPU."""
    client = _get_client()
    kwargs: dict = {"gpu": gpu, "gpu_count": gpu_count, "idle_timeout_minutes": idle_timeout}
    if image:
        kwargs["image"] = image
    if compute_target:
        kwargs["compute_target"] = compute_target
    if provider:
        kwargs["provider"] = provider
    if region:
        kwargs["region"] = region
    if sovereignty:
        kwargs["sovereignty"] = True
    if min_vram:
        kwargs["min_vram"] = min_vram
    nb = client.notebooks.create(name, **kwargs)
    click.echo(f"Notebook: {nb.id}")
    click.echo(f"Status:   {nb.status}")
    if nb.jupyter_url:
        click.echo(f"URL:      {nb.jupyter_url}")


@notebook.command("list")
@click.option("--json", "as_json", is_flag=True)
def notebook_list(as_json: bool):
    """List notebooks."""
    client = _get_client()
    notebooks = client.notebooks.list()
    if as_json:
        _json_out(notebooks)
        return
    if not notebooks:
        click.echo("No notebooks.")
        return
    rows = [
        [n.get("id", "?")[:12], n.get("name", "-"), n.get("status", "?"),
         n.get("gpu_type", "-")]
        for n in notebooks
    ]
    click.echo(_table(rows, ["ID", "Name", "Status", "GPU"]))


@notebook.command("stop")
@click.argument("notebook_id")
def notebook_stop(notebook_id: str):
    """Stop a notebook."""
    client = _get_client()
    client.notebooks.delete(notebook_id)
    click.echo(f"Stopped: {notebook_id}")


# ---------------------------------------------------------------------------
# Storage
# ---------------------------------------------------------------------------

@main.command("upload")
@click.argument("local_path")
@click.argument("remote_key")
def upload(local_path: str, remote_key: str):
    """Upload a file to Yarn storage."""
    client = _get_client()
    result = client.storage.upload(remote_key, local_path)
    click.echo(f"Uploaded: {remote_key} ({result.get('size', '?')} bytes)")


@main.command("download")
@click.argument("remote_key")
@click.argument("local_path")
def download(remote_key: str, local_path: str):
    """Download a file from Yarn storage."""
    client = _get_client()
    client.storage.download(remote_key, local_path)
    click.echo(f"Downloaded: {remote_key} -> {local_path}")


@main.command("ls")
@click.argument("prefix", default="")
@click.option("--json", "as_json", is_flag=True)
def ls(prefix: str, as_json: bool):
    """List files in storage."""
    client = _get_client()
    objects = client.storage.list(prefix=prefix)
    if as_json:
        _json_out(objects)
        return
    if not objects:
        click.echo("No files.")
        return
    for obj in objects:
        size = obj.get("size", 0)
        name = obj.get("key", obj.get("name", "?"))
        click.echo(f"  {size:>10}  {name}")


# ---------------------------------------------------------------------------
# Billing
# ---------------------------------------------------------------------------

@main.command()
def balance():
    """Show current credit balance."""
    client = _get_client()
    data = client.billing.balance()
    click.echo(f"Balance:   ${data.get('available', data.get('balance', '?')):.2f} AUD")
    held = data.get("held", 0)
    if held:
        click.echo(f"Held:      ${held:.2f} AUD")
        click.echo(f"Total:     ${data.get('total', '?'):.2f} AUD")


@main.command()
@click.option("--period", type=click.Choice(["daily", "weekly", "monthly"]), default="daily")
@click.option("--days", type=int, default=7)
@click.option("--json", "as_json", is_flag=True)
def usage(period: str, days: int, as_json: bool):
    """Show credit usage history."""
    from datetime import date, timedelta

    client = _get_client()
    to_date = date.today().isoformat()
    from_date = (date.today() - timedelta(days=days)).isoformat()
    data = client.billing.usage(period=period, from_date=from_date, to_date=to_date)
    if as_json:
        _json_out(data)
        return
    entries = data.get("usage", data.get("entries", []))
    if not entries:
        click.echo("No usage data.")
        return
    rows = [
        [e.get("date", e.get("period", "?")), f"${e.get('total', e.get('amount', 0)):.4f}"]
        for e in entries
    ]
    click.echo(_table(rows, ["Period", "Cost (AUD)"]))


@main.command()
@click.option("--json", "as_json", is_flag=True)
def pricing(as_json: bool):
    """Show GPU and inference pricing."""
    client = _get_client()
    data = client.gpus()
    p = data.get("pricing", {})
    if as_json:
        _json_out(p)
        return
    model_prices = p.get("models", {})
    default = p.get("default", {})
    rows = [
        [mid, f"${rates.get('input_per_1m', '?')}", f"${rates.get('output_per_1m', '?')}"]
        for mid, rates in model_prices.items()
    ]
    rows.append(["(default)", f"${default.get('input_per_1m', '?')}", f"${default.get('output_per_1m', '?')}"])
    click.echo(_table(rows, ["Model", "Input/1M tok (AUD)", "Output/1M tok (AUD)"]))


# ---------------------------------------------------------------------------
# Batch inference
# ---------------------------------------------------------------------------

@main.group()
def batch():
    """Batch inference — run prompts across a dataset."""
    pass


@batch.command("run")
@click.argument("input_file")
@click.option("--model", "-m", required=True, help="Model ID")
@click.option("--prompt", "-p", required=True, help="Prompt template (use {column_name} for variables)")
@click.option("--output", "-o", default=None, help="Output CSV file")
@click.option("--system", "system_prompt", default="")
@click.option("--max-tokens", type=int, default=512)
@click.option("--temperature", type=float, default=0.7)
@click.option("--concurrency", type=int, default=5)
def batch_run(input_file: str, model: str, prompt: str, output: str | None,
              system_prompt: str, max_tokens: int, temperature: float, concurrency: int):
    """Run batch inference on a CSV/JSONL file."""
    import csv

    client = _get_client()
    kwargs: dict = {
        "model": model,
        "prompt_template": prompt,
        "system_prompt": system_prompt,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "concurrency": concurrency,
    }

    if input_file.endswith(".jsonl"):
        kwargs["jsonl_file"] = input_file
    else:
        kwargs["csv_file"] = input_file

    name = Path(input_file).stem
    click.echo(f"Starting batch: {name} ({model})")
    results = client.batch.run(name=name, **kwargs)
    click.echo(f"Complete: {len(results)} results")

    if output:
        if results:
            with open(output, "w", newline="") as f:
                writer = csv.DictWriter(f, fieldnames=results[0].keys())
                writer.writeheader()
                writer.writerows(results)
            click.echo(f"Saved: {output}")
    else:
        _json_out(results)


@batch.command("list")
@click.option("--json", "as_json", is_flag=True)
def batch_list(as_json: bool):
    """List batch jobs."""
    client = _get_client()
    batches = client.batch.list()
    if as_json:
        _json_out(batches)
        return
    if not batches:
        click.echo("No batch jobs.")
        return
    rows = [
        [b.get("id", "?")[:12], b.get("name", "-"), b.get("status", "?"),
         str(b.get("total_rows", "?"))]
        for b in batches
    ]
    click.echo(_table(rows, ["ID", "Name", "Status", "Rows"]))


@batch.command("status")
@click.argument("batch_id")
def batch_status(batch_id: str):
    """Get batch job status."""
    client = _get_client()
    result = client.batch.get(batch_id)
    click.echo(f"Batch:    {result.get('id', '?')}")
    click.echo(f"Name:     {result.get('name', '-')}")
    click.echo(f"Status:   {result.get('status', '?')}")
    click.echo(f"Progress: {result.get('completed_rows', 0)}/{result.get('total_rows', '?')}")


@batch.command("cancel")
@click.argument("batch_id")
def batch_cancel(batch_id: str):
    """Cancel a running batch job."""
    client = _get_client()
    client.batch.cancel(batch_id)
    click.echo(f"Cancelled: {batch_id}")


# ---------------------------------------------------------------------------
# Secrets
# ---------------------------------------------------------------------------

@main.group()
def secret():
    """Manage research secrets (env vars for jobs/sessions)."""
    pass


@secret.command("set")
@click.argument("name")
@click.argument("value")
@click.option("--description", "-d", default="")
def secret_set(name: str, value: str, description: str):
    """Create or update a secret."""
    client = _get_client()
    client.secrets.create(name, value, description=description)
    click.echo(f"Secret set: {name}")


@secret.command("list")
def secret_list():
    """List secrets (names only)."""
    client = _get_client()
    secrets = client.secrets.list()
    if not secrets:
        click.echo("No secrets.")
        return
    for s in secrets:
        desc = f"  ({s['description']})" if s.get("description") else ""
        click.echo(f"  {s.get('name', '?')}{desc}")


@secret.command("delete")
@click.argument("name")
def secret_delete(name: str):
    """Delete a secret."""
    client = _get_client()
    client.secrets.delete(name)
    click.echo(f"Deleted: {name}")


# ---------------------------------------------------------------------------
# Models
# ---------------------------------------------------------------------------

@main.command("models")
@click.option("--json", "as_json", is_flag=True)
def list_models(as_json: bool):
    """List available inference models."""
    client = _get_client()
    models = client.chat.models()
    if as_json:
        _json_out(models)
        return
    if not models:
        click.echo("No models available.")
        return
    rows = [
        [m.get("id", m.get("model_id", "?")), m.get("owned_by", m.get("owner", "-"))]
        for m in models
    ]
    click.echo(_table(rows, ["Model", "Owner"]))
