"""Quality-enhancement pipeline client.

Posts (prompt + input PIL images + output PIL images + generation params) to
an external enhancement endpoint after inference completes. This is the
first stage of a planned pipeline where a downstream service will eventually
return an enhanced version of the output before the user sees it.

For now the call is fire-and-forget — generation must continue normally if
the POST fails or times out. We swallow exceptions on purpose; the only
visible failure mode is a `print()` to the Space logs.

Env vars (set as HF Space secrets):
    QUALITY_ENHANCEMENT_URL   — full URL including the secret path segment
    QUALITY_ENHANCEMENT_TOKEN — value for the X-Debug-Token header
"""

import io
import json
import os
import threading
from typing import Any, Iterable, Optional

try:
    import requests  # ships with gradio
except Exception:  # pragma: no cover — extremely unlikely in HF Spaces
    requests = None  # type: ignore[assignment]

from PIL import Image


_TIMEOUT_SECONDS = 5.0


def _encode_png(img: Image.Image) -> Optional[bytes]:
    """PIL image -> PNG bytes; returns None on failure."""
    try:
        buf = io.BytesIO()
        # PNG so we don't pile JPEG artifacts on top of whatever artifacts
        # the source already had — transfers should be lossless.
        img.convert("RGB").save(buf, format="PNG")
        return buf.getvalue()
    except Exception as exc:
        print(f"[enhancement] PNG encode failed: {exc}")
        return None


def _post_enhancement(url: str, token: str, data: dict[str, str],
                      files: list[tuple[str, tuple[str, bytes, str]]]) -> None:
    try:
        resp = requests.post(
            url,
            data=data,
            files=files if files else None,
            headers={"X-Debug-Token": token},
            timeout=_TIMEOUT_SECONDS,
        )
        if resp.status_code != 200:
            print(f"[enhancement] non-200 from endpoint: {resp.status_code} {resp.text[:200]}")
    except Exception as exc:
        print(f"[enhancement] failed: {exc}")


def _encode_list(images: Iterable[Image.Image], field_name: str,
                 prefix: str) -> list[tuple[str, tuple[str, bytes, str]]]:
    out: list[tuple[str, tuple[str, bytes, str]]] = []
    for idx, img in enumerate(images or []):
        if img is None:
            continue
        png = _encode_png(img)
        if png is None:
            continue
        out.append((field_name, (f"{prefix}_{idx}.png", png, "image/png")))
    return out


def send_for_enhancement(input_images: Iterable[Image.Image],
                         output_images: Iterable[Image.Image],
                         prompt: str,
                         params: dict[str, Any]) -> None:
    """Fire-and-forget POST of (inputs, outputs, prompt, params) to the
    enhancement endpoint.

    Either image list may be empty. Encodes images on the caller's thread
    (fast, lets us fail fast on encode errors), then dispatches the HTTP
    POST to a daemon thread so inference is never blocked by network
    latency or a hung receiver. Losing a transfer if the Space request
    returns before the upload finishes is acceptable for this stage —
    we don't yet read the response.
    """
    url = os.environ.get("QUALITY_ENHANCEMENT_URL", "")
    token = os.environ.get("QUALITY_ENHANCEMENT_TOKEN", "")

    if not url or not token:
        return
    if requests is None:
        print("[enhancement] requests library unavailable; skipping")
        return

    files = _encode_list(input_images, "images[]", "input")
    files += _encode_list(output_images, "output_images[]", "output")

    if not files:
        # Nothing to send — don't bother the server.
        return

    data = {
        "prompt": prompt or "",
        "params": json.dumps(params or {}, default=str),
    }

    threading.Thread(
        target=_post_enhancement,
        args=(url, token, data, files),
        daemon=True,
    ).start()