#!/usr/bin/env python3
"""End-to-end smoke test for a deployed LiteLLM Proxy stack.

Usage:
  python3 verify-chat.py <host> <model-alias>

  HOST         the subdomain (e.g. lapi.spannerjun.top) — must already be HTTPS-fronted
  MODEL_ALIAS  one of the names from litellm_config.yaml model_list

Reads master key from /root/litellm/.env (or $LITELLM_MASTER_KEY env).

Exits non-zero on any failure with a clear reason. Use after any change to
litellm_config.yaml or after restart.
"""

from __future__ import annotations

import json
import os
import subprocess
import sys
import time
from pathlib import Path

import urllib.request

ENV_PATH = Path(os.environ.get("LITELLM_ENV", "/root/litellm/.env"))


def read_master_key() -> str:
    if os.environ.get("LITELLM_MASTER_KEY"):
        return os.environ["LITELLM_MASTER_KEY"]
    if not ENV_PATH.exists():
        sys.exit(f"missing {ENV_PATH}; pass master key via env or create the file")
    for line in ENV_PATH.read_text().splitlines():
        if line.startswith("LITELLM_MASTER_KEY="):
            return line.split("=", 1)[1].strip()
    sys.exit(f"LITELLM_MASTER_KEY not found in {ENV_PATH}")


def probe(url: str, headers: dict | None = None, expect: int = 200, label: str = "") -> bytes:
    req = urllib.request.Request(url, headers=headers or {})
    try:
        with urllib.request.urlopen(req, timeout=15) as r:
            body = r.read()
            if r.status != expect:
                sys.exit(f"FAIL {label or url}: status={r.status} expected={expect}")
            return body
    except urllib.error.HTTPError as e:
        sys.exit(f"FAIL {label or url}: HTTP {e.code} {e.reason}")
    except Exception as e:
        sys.exit(f"FAIL {label or url}: {type(e).__name__}: {e}")


def main() -> int:
    if len(sys.argv) != 3:
        print(__doc__)
        sys.exit(2)
    host, model_alias = sys.argv[1], sys.argv[2]
    master = read_master_key()

    # 1. Health (no auth required)
    print("[1/5] /health/liveliness ...")
    probe(f"https://{host}/health/liveliness", label="liveness")
    print("  ok")

    # 2. /v1/models with master key (lists configured aliases)
    print("[2/5] /v1/models (master key) ...")
    body = probe(
        f"https://{host}/v1/models",
        headers={"Authorization": f"Bearer {master}"},
        label="models",
    )
    models = json.loads(body)
    aliases = [m["id"] for m in models.get("data", [])]
    print(f"  ok ({len(aliases)} aliases: {aliases})")
    if model_alias not in aliases:
        sys.exit(f"FAIL: alias '{model_alias}' not in litellm /v1/models")

    # 3. End-to-end chat completion
    print(f"[3/5] chat completion via {model_alias} ...")
    body = probe(
        f"https://{host}/v1/chat/completions",
        headers={
            "Authorization": f"Bearer {master}",
            "Content-Type": "application/json",
        },
        expect=200,
        label="chat",
    ) if False else _post_chat(host, master, model_alias)
    # Re-implement because we need to POST; the helper above was GET-only
    print("  ok")

    # 4. Sanity: cert chain is fullchain (no openssl code 21)
    print("[4/5] openssl s_client verify ...")
    out = subprocess.run(
        "echo | openssl s_client -connect {h}:443 -servername {h} 2>&1 | grep Verify".split(),
        capture_output=True, text=True, timeout=10,
    ).stdout
    if "Verify return code: 0 (ok)" in out:
        print("  ok (chain is fullchain)")
    else:
        print(f"  warn: {out.strip()} — see nginx-reverse-proxy-ssl code-21 pitfall")

    # 5. Container health + uptime
    print("[5/5] docker container status ...")
    out = subprocess.run(
        ["docker", "ps", "-a", "--format", "table {{.Names}}\t{{.Status}}"],
        capture_output=True, text=True, timeout=10,
    ).stdout
    for line in out.splitlines():
        if "litellm" in line.lower() or line.startswith("NAMES"):
            print(f"  {line}")
    return 0


def _post_chat(host: str, master: str, alias: str):
    import urllib.request
    body = json.dumps({
        "model": alias,
        "messages": [
            {"role": "system", "content": "Reply with the single word: OK"},
            {"role": "user", "content": "ping"},
        ],
        "max_tokens": 10,
    }).encode()
    req = urllib.request.Request(
        f"https://{host}/v1/chat/completions",
        data=body,
        headers={
            "Authorization": f"Bearer {master}",
            "Content-Type": "application/json",
        },
    )
    with urllib.request.urlopen(req, timeout=60) as r:
        data = json.loads(r.read())
        if "choices" not in data or not data["choices"]:
            err = data.get("error", {}).get("message", data)
            sys.exit(f"FAIL chat: no choices. error={err}")
        reply = data["choices"][0]["message"]["content"]
        print(f"  reply: {reply!r}")
        return data


if __name__ == "__main__":
    sys.exit(main())
