#!/usr/bin/env python3
"""Spin up an ephemeral Ollama server on Modal.

Launches an Ollama instance in a Modal Sandbox with a public HTTPS
tunnel. The sandbox auto-terminates after idle timeout or a hard
maximum lifetime -- no persistent state is kept (no volumes).

Requires: pip install modal && modal setup

Examples:

    # Start with defaults (llama3.2:3b, L40S GPU, 15 min idle, 2 hr max)
    python scripts/modal_ollama.py

    # Pick a different model and GPU
    python scripts/modal_ollama.py --model qwen2.5:7b --gpu A10G

    # CPU-only (small models only)
    python scripts/modal_ollama.py --model qwen2.5:0.5b --gpu none

    # Shorter leash -- 30 min max, 5 min idle
    python scripts/modal_ollama.py --timeout 1800 --idle-timeout 300

    # Use the printed URL with looselips
    OLLAMA_HOST=https://<printed-url> \\
        looselips export.zip -c config.toml -m ollama/llama3.2:3b

    # List running sandboxes
    python scripts/modal_ollama.py --list

    # Tear down by ID
    python scripts/modal_ollama.py --kill sb-XXXXX
"""

from __future__ import annotations

import argparse
import sys
import time
import urllib.error
import urllib.request

OLLAMA_PORT = 11434
OLLAMA_VERSION = "0.6.5"
APP_NAME = "looselips-ollama"


def _fmt_duration(seconds: int) -> str:
    if seconds < 60:
        return f"{seconds}s"
    m, s = divmod(seconds, 60)
    if m < 60:
        return f"{m}m{s}s" if s else f"{m}m"
    h, m = divmod(m, 60)
    parts = [f"{h}h"]
    if m:
        parts.append(f"{m}m")
    return "".join(parts)


def _wait_healthy(url: str, tries: int = 30, interval: float = 1.0) -> bool:
    """Poll the Ollama API until it responds or we run out of tries."""
    for i in range(tries):
        try:
            with urllib.request.urlopen(f"{url}/api/version", timeout=3) as r:
                if r.status == 200:
                    return True
        except (urllib.error.URLError, OSError):
            pass
        if i < tries - 1:
            time.sleep(interval)
    return False


def _build_image():  # -> modal.Image
    import modal

    return (
        modal.Image.debian_slim(python_version="3.11")
        .apt_install("curl", "ca-certificates")
        .run_commands(
            f"OLLAMA_VERSION={OLLAMA_VERSION}"
            " curl -fsSL https://ollama.com/install.sh | sh",
        )
        .env({"OLLAMA_HOST": f"0.0.0.0:{OLLAMA_PORT}"})
    )


def _start(args: argparse.Namespace) -> None:
    import modal

    print(f"Creating sandbox ({args.gpu}, {args.model})...")

    app = modal.App.lookup(APP_NAME, create_if_missing=True)

    sandbox_kwargs: dict[str, object] = dict(
        app=app,
        image=_build_image(),
        timeout=args.timeout,
        idle_timeout=args.idle_timeout,
        h2_ports=[OLLAMA_PORT],
    )
    if args.gpu.lower() != "none":
        sandbox_kwargs["gpu"] = args.gpu

    sb = modal.Sandbox.create(
        "bash", "-lc", "ollama serve",
        **sandbox_kwargs,
    )

    tunnel = sb.tunnels()[OLLAMA_PORT]
    url = tunnel.url

    print(f"Waiting for Ollama to become healthy...")
    if not _wait_healthy(url):
        print("ERROR: Ollama did not become healthy after 30s", file=sys.stderr)
        sb.terminate()
        sys.exit(1)
    print(f"Ollama is up.")

    if not args.no_pull:
        print(f"Pulling {args.model} (this may take a minute)...")
        p = sb.exec("ollama", "pull", args.model)
        for line in p.stdout:
            sys.stdout.write(line)
        rc = p.wait()
        if rc != 0:
            print(f"ERROR: ollama pull exited {rc}", file=sys.stderr)
            sb.terminate()
            sys.exit(1)

    idle = _fmt_duration(args.idle_timeout)
    maxlife = _fmt_duration(args.timeout)

    print()
    print(f"  Sandbox : {sb.object_id}")
    print(f"  URL     : {url}")
    print(f"  Model   : {args.model}")
    print(f"  GPU     : {args.gpu}")
    print(f"  Expires : {idle} idle / {maxlife} max")
    print()
    print(f"  Run looselips:")
    print()
    print(f"    OLLAMA_HOST={url} \\")
    print(f"        looselips export.zip -c config.toml -m ollama/{args.model}")
    print()
    print(f"  Tear down:")
    print()
    print(f"    python {sys.argv[0]} --kill {sb.object_id}")
    print()


def _kill(sandbox_id: str) -> None:
    import modal

    sb = modal.Sandbox.from_id(sandbox_id)
    sb.terminate()
    print(f"Terminated {sandbox_id}")


def _list_sandboxes() -> None:
    import modal

    app = modal.App.lookup(APP_NAME, create_if_missing=False)
    if app is None:
        print("No running sandboxes.")
        return

    found = False
    for sb in modal.Sandbox.list(app_id=app.app_id):
        print(f"  {sb.object_id}")
        found = True
    if not found:
        print("No running sandboxes.")


def main() -> None:
    ap = argparse.ArgumentParser(
        description="Ephemeral Ollama on Modal for looselips LLM scanning.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Model suggestions:\n"
            "  llama3.2:3b       good general-purpose small model (default)\n"
            "  qwen2.5:7b        strong general, larger\n"
            "  qwen2.5-coder:7b  code-focused\n"
            "  mistral:7b        solid classic 7B\n"
            "  gemma2:2b         very small, CPU-ok\n"
            "  qwen2.5:0.5b      tiny, CPU-ok\n"
        ),
    )
    ap.add_argument(
        "--model", default="llama3.2:3b",
        help="Ollama model tag (default: llama3.2:3b)",
    )
    ap.add_argument(
        "--gpu", default="L40S",
        help="GPU type: L40S, A10G, A100, H100, or 'none' (default: L40S)",
    )
    ap.add_argument(
        "--idle-timeout", type=int, default=15 * 60, metavar="SEC",
        help="kill after N seconds idle (default: 900)",
    )
    ap.add_argument(
        "--timeout", type=int, default=2 * 60 * 60, metavar="SEC",
        help="hard max lifetime in seconds (default: 7200)",
    )
    ap.add_argument(
        "--no-pull", action="store_true",
        help="don't pre-pull the model",
    )
    ap.add_argument(
        "--kill", metavar="ID",
        help="terminate a sandbox by ID",
    )
    ap.add_argument(
        "--list", action="store_true",
        help="list running sandboxes",
    )
    args = ap.parse_args()

    if args.kill:
        _kill(args.kill)
    elif args.list:
        _list_sandboxes()
    else:
        _start(args)


if __name__ == "__main__":
    main()
