#!/usr/bin/env python3
"""Check a user-supplied OpenAI-compatible endpoint. No generation by default."""

import argparse
import getpass
import ipaddress
import json
import math
import os
import sys
import urllib.error
import urllib.parse
import urllib.request

MAX_BYTES = 2 * 1024 * 1024


class CheckError(Exception):
    """A connection or response failed validation."""


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        raise CheckError("Redirect refused; use the final API base URL.")


def validate_url(value):
    try:
        url = urllib.parse.urlsplit(value)
        port = url.port
        if (url.scheme not in ("http", "https") or not url.hostname
                or url.username is not None or url.password is not None
                or url.query or url.fragment or port == 0
                or any(c.isspace() or ord(c) < 32 for c in value)):
            raise ValueError
        if url.scheme == "http":
            try:
                loopback = ipaddress.ip_address(url.hostname).is_loopback
            except ValueError:
                loopback = False
            if not loopback:
                raise CheckError("HTTP requires a loopback IP, such as 127.0.0.1; otherwise use HTTPS.")
        path = url.path.rstrip("/")
        if not path.endswith("/v1"):
            raise CheckError("Base URL must end in /v1, without /models or /chat/completions.")
        return urllib.parse.urlunsplit((url.scheme, url.netloc, path, "", ""))
    except ValueError:
        raise CheckError("Invalid base URL; omit credentials, query strings and fragments.") from None


def request_json(opener, url, key, timeout, payload=None):
    headers = {"Accept": "application/json"}
    if key:
        if any(ord(c) < 32 or ord(c) > 126 for c in key):
            raise CheckError("API key contains unsupported characters.")
        headers["Authorization"] = "Bearer " + key
    data = None
    if payload is not None:
        headers["Content-Type"] = "application/json"
        data = json.dumps(payload).encode("utf-8")
    try:
        with opener.open(urllib.request.Request(url, data=data, headers=headers), timeout=timeout) as response:
            raw = response.read(MAX_BYTES + 1)
            if len(raw) > MAX_BYTES:
                raise CheckError("Response exceeded the 2 MiB limit.")
    except urllib.error.HTTPError as error:
        status = error.code
        error.close()
        raise CheckError(f"HTTP {status}; verify the endpoint, key, model and server logs.") from None
    except (urllib.error.URLError, TimeoutError, OSError):
        raise CheckError("Connection failed or timed out; check the tunnel, TLS and server. No retry was sent.") from None
    try:
        result = json.loads(raw)
    except (ValueError, UnicodeError):
        raise CheckError("Response was not valid JSON; this may be a web page instead of the API.") from None
    if not isinstance(result, dict) or "error" in result:
        raise CheckError("API returned an error or an unexpected JSON structure.")
    return result


def check(base_url, key="", model=None, generate=False, timeout=60):
    base_url = validate_url(base_url)
    if not math.isfinite(timeout) or not 0 < timeout <= 300:
        raise CheckError("Timeout must be between 0 and 300 seconds, excluding zero.")
    if generate and not model:
        raise CheckError("--generate requires an explicit --model from the model list.")
    # Do not forward credentials through redirects or environment HTTP proxies.
    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}), NoRedirect())
    result = request_json(opener, base_url + "/models", key, timeout)
    rows = result.get("data")
    if not isinstance(rows, list) or not rows:
        raise CheckError("No models returned; check whether a model is installed and served.")
    if any(not isinstance(row, dict) or not isinstance(row.get("id"), str)
           or not row["id"].strip() for row in rows):
        raise CheckError("Model list contains invalid IDs.")
    ids = [row["id"] for row in rows]
    if model is not None and model not in ids:
        raise CheckError("Selected model was not returned by /models; copy the exact served ID.")
    if generate:
        result = request_json(opener, base_url + "/chat/completions", key, timeout, {
            "model": model,
            "messages": [{"role": "user", "content": "Reply with one short greeting."}],
            "max_tokens": 64,
            "temperature": 0,
            "stream": False,
        })
        choices = result.get("choices")
        if not isinstance(choices, list) or not choices or not isinstance(choices[0], dict):
            raise CheckError("Generation returned no valid choices.")
        message = choices[0].get("message")
        if (not isinstance(message, dict) or not isinstance(message.get("content"), str)
                or not message["content"].strip()):
            raise CheckError("Generation returned no text; inspect the model, template and reasoning settings.")
        if choices[0].get("finish_reason") != "stop":
            raise CheckError("Generation did not finish normally; inspect output limits and model settings.")
    return ids


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base-url", required=True, help="API base ending in /v1")
    parser.add_argument("--model", help="Exact model ID returned by /models")
    parser.add_argument("--generate", action="store_true", help="Send one real generation request (may cost money)")
    parser.add_argument("--ask-key", action="store_true", help="Read the key without echo; otherwise use CHAT_API_KEY")
    parser.add_argument("--timeout", type=float, default=60, help="Per network operation timeout, not a billing cap")
    args = parser.parse_args(argv)
    try:
        key = getpass.getpass("API key: ") if args.ask_key else os.environ.get("CHAT_API_KEY", "")
        ids = check(args.base_url, key, args.model, args.generate, args.timeout)
    except (CheckError, EOFError) as error:
        print(f"ERROR: {error or 'No API key input available.'}", file=sys.stderr)
        return 1
    print(json.dumps({"models": ids, "generation_checked": args.generate}, ensure_ascii=True))
    return 0


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