"""Server-side HTTPS client. Credentials and Runpod URLs never reach the browser."""

import os
import ssl
import time
import uuid
from urllib.parse import urlparse

import httpx
from dev_protocol import JobRequest, decode_image


class DevError(RuntimeError):
    pass


def configured():
    return bool(os.environ.get("IRIS_Z_DEV_URL") and os.environ.get("IRIS_Z_DEV_API_KEY"))


def expected_seconds(request):
    """Expected GPU time of one Iris-Z big request on an H200, from measured timings (2026-11-11): 0.66 s per
    guided step at 512x512 and 2.45 s at 1024x1024 (time grows with pixels^0.94); depth steps run one pass, not a
    guided pair. The worker reports no step progress, so the Space shows elapsed / expected."""
    if request.task in {"text_to_image", "text_to_rgbd"}:
        pixels = request.width * request.height
    else:
        pixels = request.resolution * request.resolution
    seconds_per_step = 0.66 * (pixels / (512 * 512)) ** 0.94
    if request.task == "image_to_depth":
        seconds_per_step *= 0.5
    return 3.0 + request.steps * seconds_per_step


def run(payload, progress=None):
    if not configured():
        raise DevError("Iris-Z big is not connected yet.")
    try:
        request = JobRequest.model_validate(payload)
    except ValueError:
        raise DevError("Choose supported Iris-Z big parameters and a valid image.") from None
    url = os.environ["IRIS_Z_DEV_URL"].rstrip("/")
    parsed = urlparse(url)
    if (
        parsed.scheme != "https"
        or not parsed.hostname
        or parsed.username
        or parsed.password
        or parsed.query
        or parsed.fragment
        or parsed.path
    ):
        raise DevError("Iris-Z big has an invalid service configuration.")
    headers = {
        "Authorization": "Bearer " + os.environ["IRIS_Z_DEV_API_KEY"],
        "X-Request-ID": str(uuid.uuid4()),
    }
    deadline = time.monotonic() + 1800

    def checked(response):
        if response.status_code >= 400:
            messages = {
                401: "Iris-Z big authentication failed.",
                403: "Iris-Z big access denied.",
                429: "Iris-Z big is busy. Please try again shortly.",
                503: "Iris-Z big is starting or unavailable. Please try again shortly.",
                422: "Iris-Z big rejected the image or settings.",
                404: "Iris-Z big job expired or the service restarted.",
            }
            raise DevError(messages.get(response.status_code, "Iris-Z big could not complete this request."))
        return response.json()

    try:
        # Provisional AWS endpoints use our private CA; never disable TLS checks.
        certificate = os.environ.get("IRIS_Z_DEV_CA_PEM")
        verification = ssl.create_default_context(cadata=certificate) if certificate else True
        with httpx.Client(
            base_url=url,
            headers=headers,
            timeout=45,
            follow_redirects=False,
            trust_env=False,
            verify=verification,
        ) as client:
            # An ambiguous submission is retried with the same ID, never as a second GPU job.
            for attempt in range(3):
                try:
                    job = checked(client.post("/jobs", json=request.model_dump()))
                    break
                except httpx.TransportError:
                    if attempt == 2:
                        raise
                    time.sleep(1)
            job_id = str(uuid.UUID(job["id"]))
            expected = expected_seconds(request)
            running_since = None
            while time.monotonic() < deadline:
                try:
                    job_state = checked(client.get(f"/jobs/{job_id}"))
                    status = job_state["status"]
                except httpx.TransportError:
                    time.sleep(2)
                    continue
                if status == "completed":
                    result = checked(client.get(f"/jobs/{job_id}/result"))
                    for key in ["image", "depth"]:
                        if key in result:
                            result[key] = decode_image(result[key])
                    return result
                if status == "failed":
                    raise DevError(
                        "Iris-Z big could not complete this request. Try a smaller image or fewer steps."
                    )
                if progress is not None:
                    position = job_state.get("queue_position")
                    waiting = (
                        f"Iris-Z big queue · approximately {position} in line"
                        if isinstance(position, int) and position > 0
                        else "Waiting for Iris-Z big"
                    )
                    if status == "queued":
                        running_since = None
                        progress(0, desc=waiting)
                    else:
                        running_since = running_since or time.monotonic()
                        fraction = min(0.97, (time.monotonic() - running_since) / expected)
                        progress(fraction, desc="Iris-Z big is generating")
                time.sleep(1)
    except (httpx.HTTPError, ValueError, KeyError):
        raise DevError("The connection to Iris-Z big failed. Please try again shortly.") from None
    raise DevError("Iris-Z big exceeded the request time limit.")