#!/usr/bin/env python3
"""Small, bounded Chat Completions probe. Python 3.11+, standard library only."""

import argparse
from concurrent.futures import ThreadPoolExecutor, as_completed
import http.client
import json
import math
import os
import socket
import time
from urllib.error import HTTPError, URLError
from urllib.parse import urlsplit
from urllib.request import HTTPRedirectHandler, Request, build_opener


MAX_RESPONSE_BYTES = 2 * 1024 * 1024
TOKEN_FIELDS = ("prompt_tokens", "completion_tokens")
FINISH_REASONS = {"stop", "length", "tool_calls", "function_call", "content_filter"}


class NoRedirect(HTTPRedirectHandler):
    """Never forward a bearer credential or repeat a probe through redirects."""

    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def bounded_integer(minimum, maximum):
    def parse(value):
        try:
            number = int(value)
        except ValueError:
            raise argparse.ArgumentTypeError("expected an integer") from None
        if not minimum <= number <= maximum:
            raise argparse.ArgumentTypeError(f"must be between {minimum} and {maximum}")
        return number

    return parse


def timeout_seconds(value):
    try:
        number = float(value)
    except ValueError:
        raise argparse.ArgumentTypeError("expected seconds") from None
    if not math.isfinite(number) or not 0 < number <= 3600:
        raise argparse.ArgumentTypeError("must be greater than 0 and at most 3600 seconds")
    return number


def configuration():
    """Validate without including configuration values in error messages."""
    key = os.environ.get("QDIV0_API_KEY", "")
    if not key or any(ord(char) < 33 or ord(char) > 126 for char in key):
        raise ValueError("QDIV0_API_KEY must contain a nonempty, header-safe API key")
    model = os.environ.get("QDIV0_MODEL", "equipo-codigo").strip()
    if not model or len(model) > 256 or any(ord(char) < 32 for char in model):
        raise ValueError("QDIV0_MODEL must contain a valid model alias")
    origin = os.environ.get("QDIV0_API_URL", "https://api.qdiv0.com")
    try:
        if any(char.isspace() or ord(char) < 32 for char in origin):
            raise ValueError
        parsed = urlsplit(origin)
        if (
            not parsed.hostname or parsed.username is not None or parsed.password is not None
            or parsed.path not in ("", "/") or parsed.query or parsed.fragment
            or parsed.scheme not in ("http", "https")
        ):
            raise ValueError
        if parsed.scheme == "http" and parsed.hostname not in ("127.0.0.1", "localhost", "::1"):
            raise ValueError
        parsed.port  # Validate malformed/out-of-range ports before sending anything.
        parsed.netloc.encode("ascii")
    except (ValueError, UnicodeError):
        raise ValueError("QDIV0_API_URL must be an HTTPS origin without credentials, path, query or fragment; HTTP is allowed only on loopback") from None
    return key, origin.rstrip("/") + "/v1/chat/completions", model


def request_payload(model, max_tokens, tools=False):
    prompt = "Write a short Python function add(a, b) that returns their sum, with one example. Use fictional data only."
    payload = {"model": model, "messages": [{"role": "user", "content": prompt}], "max_tokens": max_tokens, "stream": False}
    if tools:
        payload["messages"][0]["content"] = "Use get_project_status to look up the fictional demo-project, if tools are available. Do not invent its status."
        payload["tools"] = [{
            "type": "function",
            "function": {
                "name": "get_project_status",
                "description": "Read-only lookup of a fictional demo project's status.",
                "parameters": {
                    "type": "object",
                    "properties": {"project_id": {"type": "string", "enum": ["demo-project"]}},
                    "required": ["project_id"],
                    "additionalProperties": False,
                },
            },
        }]
    return json.dumps(payload).encode("utf-8")


def response_metadata(body):
    choices = body.get("choices") if isinstance(body, dict) else None
    if not isinstance(choices, list) or not choices or not isinstance(choices[0], dict):
        raise ValueError
    choice = choices[0]
    message = choice.get("message")
    if not isinstance(message, dict):
        raise ValueError
    content = message.get("content")
    calls = message.get("tool_calls") or []
    if (content is not None and not isinstance(content, str)) or not isinstance(calls, list):
        raise ValueError
    finish = choice.get("finish_reason")
    finish = finish if isinstance(finish, str) and finish in FINISH_REASONS else ("other" if finish is not None else None)
    usage = body.get("usage")
    usage = usage if isinstance(usage, dict) else {}
    counts = {}
    for field in TOKEN_FIELDS:
        value = usage.get(field)
        counts[field] = value if type(value) is int and value >= 0 else None
    return {
        "finish_reason": finish,
        "usage": counts,
        "content_nonempty": bool(content and content.strip()),
        "tool_calls": len(calls),
    }


def probe(index, url, key, payload, timeout):
    result = {
        "type": "request", "request": index, "http_status": None,
        "finish_reason": None, "usage": dict.fromkeys(TOKEN_FIELDS),
        "content_nonempty": False, "tool_calls": 0,
    }
    start = time.perf_counter()
    try:
        request = Request(url, data=payload, headers={
            "Authorization": f"Bearer {key}", "Content-Type": "application/json", "Accept": "application/json",
        }, method="POST")
        # A separate opener per worker avoids shared mutable client state.
        with build_opener(NoRedirect()).open(request, timeout=timeout) as response:
            result["http_status"] = response.status
            raw = response.read(MAX_RESPONSE_BYTES + 1)
        if not 200 <= result["http_status"] < 300:
            result["error"] = "http_error"
        elif len(raw) > MAX_RESPONSE_BYTES:
            result["error"] = "response_too_large"
        else:
            try:
                body = json.loads(raw)
            except (ValueError, UnicodeError):
                result["error"] = "invalid_json"
            else:
                try:
                    result.update(response_metadata(body))
                except (ValueError, TypeError):
                    result["error"] = "invalid_response"
    except HTTPError as error:
        result["http_status"] = error.code
        result["error"] = "http_error"
        error.close()  # Do not read, echo or retain the server's error body.
    except (TimeoutError, socket.timeout):
        result["error"] = "timeout"
    except URLError as error:
        result["error"] = "timeout" if isinstance(error.reason, TimeoutError) else "network_error"
    except (OSError, http.client.HTTPException):
        result["error"] = "network_error"
    result["seconds"] = time.perf_counter() - start
    return result


def percentile(values, fraction):
    """Linear interpolation at (n - 1) * fraction (the common R-7 method)."""
    if not values:
        return None
    ordered = sorted(values)
    rank = (len(ordered) - 1) * fraction
    lower = math.floor(rank)
    upper = math.ceil(rank)
    return ordered[lower] + (ordered[upper] - ordered[lower]) * (rank - lower)


def summary(results, concurrency, elapsed):
    successful = [row for row in results if "error" not in row]
    times = [row["seconds"] for row in successful]
    totals, observed = {}, {}
    for field in TOKEN_FIELDS:
        values = [row["usage"][field] for row in successful if row["usage"][field] is not None]
        totals[field] = sum(values) if values else None
        observed[field] = len(values)
    return {
        "type": "summary", "requests": len(results), "concurrency": concurrency,
        "successes": len(successful), "failures": len(results) - len(successful),
        "elapsed_seconds": elapsed,
        "latency_seconds": {"p50": percentile(times, 0.5), "p95": percentile(times, 0.95), "samples": len(times)},
        "usage": totals, "usage_responses": observed,
    }


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--requests", type=bounded_integer(1, 64), default=1, help="total requests, 1-64 (default: 1)")
    parser.add_argument("--concurrency", type=bounded_integer(1, 32), default=1, help="parallel workers, 1-32 (default: 1)")
    parser.add_argument("--max-tokens", type=bounded_integer(1, 32768), default=2048, help="output-token cap, 1-32768 (default: 2048)")
    parser.add_argument("--timeout", type=timeout_seconds, default=900, help="network-operation timeout in seconds (default: 900)")
    parser.add_argument("--tools", action="store_true", help="offer a fictional read-only function; never execute returned calls")
    args = parser.parse_args(argv)
    try:
        key, url, model = configuration()
    except ValueError as error:
        parser.error(str(error))
    payload = request_payload(model, args.max_tokens, args.tools)
    workers = min(args.concurrency, args.requests)
    results = []
    start = time.perf_counter()
    with ThreadPoolExecutor(max_workers=workers) as pool:
        futures = [pool.submit(probe, index, url, key, payload, args.timeout) for index in range(1, args.requests + 1)]
        for future in as_completed(futures):
            result = future.result()
            results.append(result)
            print(json.dumps(result, allow_nan=False), flush=True)
    report = summary(results, workers, time.perf_counter() - start)
    print(json.dumps(report, allow_nan=False), flush=True)
    return 1 if report["failures"] else 0


if __name__ == "__main__":
    raise SystemExit(main())