#!/usr/bin/env python3
"""Compare one Ollama request with the model unloaded versus resident.

This is a bounded latency test, not a cold OS/file-cache benchmark and not a
memory-capacity test. It requires an already-installed model and a local Ollama
server.
"""
from __future__ import annotations

import argparse
import json
import os
import platform
import statistics
import subprocess
import time
import urllib.request
from collections.abc import Sequence
from datetime import datetime, timezone
from pathlib import Path
from urllib.parse import urlsplit

PROMPT = "Return a JSON object with status set to ok and items set to 3. Do not add any other keys."
SCHEMA = {
    "type": "object",
    "properties": {
        "status": {"type": "string", "const": "ok"},
        "items": {"type": "integer", "const": 3},
    },
    "required": ["status", "items"],
    "additionalProperties": False,
}
EXPECTED = {"status": "ok", "items": 3}


def ollama(cli_env: dict[str, str], *args: str) -> str:
    return subprocess.run(
        ["ollama", *args],
        check=True,
        capture_output=True,
        text=True,
        env=cli_env,
    ).stdout


def ollama_cli_environment(endpoint: str) -> dict[str, str]:
    """Bind CLI inspection and cleanup to the same local server as the API."""
    parsed = urlsplit(endpoint)
    if (
        parsed.scheme != "http"
        or parsed.hostname not in {"127.0.0.1", "localhost", "::1"}
        or parsed.username is not None
        or parsed.password is not None
        or parsed.path != "/api/generate"
        or parsed.query
        or parsed.fragment
    ):
        raise SystemExit(
            "--endpoint must be an HTTP loopback Ollama generate endpoint, "
            "for example http://127.0.0.1:11434/api/generate"
        )
    try:
        parsed.port
    except ValueError as exc:
        raise SystemExit(f"Invalid --endpoint port: {exc}") from exc

    cli_env = os.environ.copy()
    cli_env["OLLAMA_HOST"] = f"{parsed.scheme}://{parsed.netloc}"
    return cli_env


def loaded_models(cli_env: dict[str, str]) -> set[str]:
    models: set[str] = set()
    for line in ollama(cli_env, "ps").splitlines():
        fields = line.split()
        if fields and fields[0] != "NAME":
            models.add(fields[0])
    return models


def model_is_loaded(model: str, cli_env: dict[str, str]) -> bool:
    return model in loaded_models(cli_env)


def stop_and_verify(model: str, cli_env: dict[str, str]) -> None:
    ollama(cli_env, "stop", model)
    if model_is_loaded(model, cli_env):
        raise RuntimeError(f"Cleanup failed: {model} is still listed by ollama ps")


def request(endpoint: str, model: str, keep_alive: int | str) -> dict:
    payload = {
        "model": model,
        "prompt": PROMPT,
        "stream": False,
        "keep_alive": keep_alive,
        "format": SCHEMA,
        "options": {"temperature": 0, "num_predict": 30},
    }
    body = json.dumps(payload).encode()
    started = time.perf_counter()
    with urllib.request.urlopen(
        urllib.request.Request(endpoint, body, {"Content-Type": "application/json"}),
        timeout=120,
    ) as response:
        result = json.load(response)
    wall = time.perf_counter() - started
    try:
        accepted = json.loads(result["response"]) == EXPECTED
    except (KeyError, json.JSONDecodeError):
        accepted = False
    return {
        "wall_seconds": round(wall, 3),
        "total_seconds": round(result["total_duration"] / 1e9, 3),
        "load_seconds": round(result["load_duration"] / 1e9, 3),
        "prompt_eval_seconds": round(result["prompt_eval_duration"] / 1e9, 3),
        "eval_seconds": round(result["eval_duration"] / 1e9, 3),
        "response": result["response"],
        "accepted": accepted,
    }


def main(argv: Sequence[str] | None = None) -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="qwen2.5-coder:14b")
    parser.add_argument("--runs", type=int, default=3)
    parser.add_argument("--endpoint", default="http://127.0.0.1:11434/api/generate")
    parser.add_argument("--output", type=Path, default=Path("ollama-unloaded-vs-resident-results.json"))
    args = parser.parse_args(argv)
    if args.runs < 1:
        raise SystemExit("--runs must be at least 1")

    cli_env = ollama_cli_environment(args.endpoint)
    version = ollama(cli_env, "--version").strip()
    stop_and_verify(args.model, cli_env)

    try:
        unloaded = []
        for _ in range(args.runs):
            unloaded.append(request(args.endpoint, args.model, 0))
            if model_is_loaded(args.model, cli_env):
                raise RuntimeError(
                    f"Unloaded-model run failed: {args.model} is still listed by ollama ps"
                )

        # This setup request loads the model but is not counted in resident results.
        resident_setup = request(args.endpoint, args.model, "10m")
        if not model_is_loaded(args.model, cli_env):
            raise RuntimeError(f"Setup failed: {args.model} is not listed by ollama ps")
        resident = [request(args.endpoint, args.model, "10m") for _ in range(args.runs)]
    finally:
        # Always attempt to stop and verify the model, including after request errors.
        stop_and_verify(args.model, cli_env)

    output = {
        "captured_at_utc": datetime.now(timezone.utc).isoformat(),
        "host": {"os": platform.platform(), "machine": platform.machine()},
        "ollama_version": version,
        "endpoint": args.endpoint,
        "model": args.model,
        "prompt": PROMPT,
        "format_schema": SCHEMA,
        "unloaded_definition": "The model was absent from ollama ps before each request because keep_alive was 0. OS and file-system caches were not cleared.",
        "resident_definition": "One uncounted setup request loaded the model with keep_alive 10m; measured requests used the same value.",
        "unloaded_model_runs": unloaded,
        "resident_setup_run": resident_setup,
        "resident_model_runs": resident,
        "accepted_unloaded": sum(run["accepted"] for run in unloaded),
        "accepted_resident": sum(run["accepted"] for run in resident),
        "median_unloaded_wall_seconds": round(statistics.median(run["wall_seconds"] for run in unloaded), 3),
        "median_resident_wall_seconds": round(statistics.median(run["wall_seconds"] for run in resident), 3),
        "range_unloaded_wall_seconds": [min(run["wall_seconds"] for run in unloaded), max(run["wall_seconds"] for run in unloaded)],
        "range_resident_wall_seconds": [min(run["wall_seconds"] for run in resident), max(run["wall_seconds"] for run in resident)],
        "cleanup_verified": True,
        "limit": "This isolates Ollama model residency for one short local request. It does not measure resident memory, control OS caches, establish the constraint in other screenshots, or predict other models and workloads.",
    }
    args.output.write_text(json.dumps(output, indent=2) + "\n")
    print(json.dumps(output, indent=2))


if __name__ == "__main__":
    main()