from __future__ import annotations

import argparse
import json
import sys
import urllib.error
import urllib.request


def post_chat(base_url: str, model: str, prompt: str) -> dict:
    url = f"{base_url.rstrip('/')}/v1/chat/completions"
    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": "Отвечай кратко и по делу."},
            {"role": "user", "content": prompt},
        ],
        "temperature": 0.2,
        "max_tokens": 256,
    }
    request = urllib.request.Request(
        url,
        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=120) as response:
        return json.loads(response.read().decode("utf-8"))


def main() -> int:
    parser = argparse.ArgumentParser(description="Smoke test an OpenAI-compatible chat endpoint.")
    parser.add_argument("--base-url", default="http://docker-gpu.cin.su:8000")
    parser.add_argument("--model", default="qwen3-4b-instruct")
    parser.add_argument(
        "--prompt",
        default="Объясни в одном предложении, зачем нужен реестр локальных моделей.",
    )
    args = parser.parse_args()

    try:
        response = post_chat(args.base_url, args.model, args.prompt)
    except urllib.error.URLError as exc:
        print(f"Request failed: {exc}", file=sys.stderr)
        return 1

    choices = response.get("choices") or []
    if not choices:
        print("Response has no choices.", file=sys.stderr)
        print(json.dumps(response, ensure_ascii=False, indent=2))
        return 1

    message = choices[0].get("message") or {}
    content = (message.get("content") or "").strip()
    if not content:
        print("Response content is empty.", file=sys.stderr)
        print(json.dumps(response, ensure_ascii=False, indent=2))
        return 1

    print(content)
    return 0


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