"""Python 3.10+, standard library only. Key stays in the process environment."""
import json
import os
import sys
import urllib.error
import urllib.parse
import urllib.request
import uuid


def api_base(value):
    parsed = urllib.parse.urlparse(value)
    local = parsed.scheme == "http" and parsed.hostname in ("localhost", "127.0.0.1", "::1")
    if not parsed.hostname or parsed.username or parsed.password or parsed.query or parsed.fragment or parsed.path not in ("", "/") or (parsed.scheme != "https" and not local):
        raise ValueError("Use the HTTPS API origin, or HTTP loopback for local tests")
    return value.rstrip("/")


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def prepare(base, model, path, max_cost_microrub=1_000_000):
    if not model:
        raise ValueError("Set KVANTORA_MODEL_ID to the public API ID")
    if isinstance(max_cost_microrub, bool) or not isinstance(max_cost_microrub, int) or not 0 < max_cost_microrub <= 9_007_199_254_740_991:
        raise ValueError("Set a positive safe integer budget in microrubles")
    request = {"version": 1, "apiBase": api_base(base), "idempotencyKey": str(uuid.uuid4()),
               "body": {"model": model, "messages": [{"role": "user", "content": "Назови столицу Франции одним словом."}], "max_tokens": 64, "max_cost_microrub": max_cost_microrub}}
    with open(path, "x", encoding="utf-8") as file:
        json.dump(request, file, ensure_ascii=False, indent=2)
    return request


def call(command, base, key, path):
    if not key:
        raise ValueError("Set KVANTORA_API_KEY in the server environment")
    headers = {"Authorization": "Bearer " + key}
    target = api_base(base) + "/v1/models"
    body = None
    if command in ("quote", "send"):
        with open(path, encoding="utf-8") as file:
            saved = json.load(file)
        if saved.get("version") != 1 or str(uuid.UUID(saved["idempotencyKey"])) != saved["idempotencyKey"] or not saved.get("body", {}).get("model"):
            raise ValueError("Invalid saved request")
        if api_base(saved["apiBase"]) != api_base(base):
            raise ValueError("API origin differs from saved request")
        target = saved["apiBase"] + "/v1/chat/" + ("quote" if command == "quote" else "completions")
        body = json.dumps(saved["body"], ensure_ascii=False).encode("utf-8")
        headers.update({"Content-Type": "application/json", "Idempotency-Key": saved["idempotencyKey"]})
    elif command != "catalog":
        raise ValueError("Use catalog, prepare, quote or send")
    opener = urllib.request.build_opener(NoRedirect)
    try:
        response = opener.open(urllib.request.Request(target, data=body, headers=headers), timeout=60)
    except urllib.error.HTTPError as error:
        code = "request_failed"
        try:
            detail = json.loads(error.read(65536)).get("error", {}).get("code")
            if isinstance(detail, str):
                code = detail
        except (ValueError, AttributeError):
            pass
        return {"status": error.code, "traceId": error.headers.get("X-Request-Id"), "error": code}
    with response:
        payload = json.load(response)
        if command == "catalog":
            return {"status": response.status, "models": payload.get("data", [])}
        if command == "quote":
            return {"status": response.status, "estimate": payload}
        completed = payload.get("object") == "chat.completion" and payload.get("billing", {}).get("status") == "settled"
        text = payload.get("choices", [{}])[0].get("message", {}).get("content") if payload.get("object") == "chat.completion" else None
        return {"status": response.status, "requestId": payload.get("id"), "traceId": response.headers.get("X-Request-Id"), "completed": completed, "text": text,
                "costMicrorub": payload.get("usage", {}).get("cost_microrub"), "billing": payload.get("billing", {}).get("status")}


def main():
    command = sys.argv[1] if len(sys.argv) > 1 else "catalog"
    path = sys.argv[2] if len(sys.argv) > 2 else "kvantora-request.json"
    base = os.environ["KVANTORA_BASE_URL"]
    if command == "prepare":
        prepare(base, os.environ.get("KVANTORA_MODEL_ID"), path, int(os.environ.get("KVANTORA_MAX_COST_MICRORUB", "1000000")))
        print("Saved request. Run quote to estimate; send explicitly starts a paid request if available.")
    else:
        result = call(command, base, os.environ.get("KVANTORA_API_KEY"), path)
        print(json.dumps(result, ensure_ascii=False, indent=2))
        if result["status"] >= 400:
            sys.exit(1)


if __name__ == "__main__":
    try:
        main()
    except Exception:
        print("Request was not completed. Keep the saved request and check configuration or usage history. No automatic retry was made.", file=sys.stderr)
        sys.exit(1)