"""Minimal PixelLab v2 client for style-consistent pixel-art generation."""

import base64
import json
import os
import time
from io import BytesIO
from typing import Any, Optional
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

from PIL import Image


API_URL = "https://api.pixellab.ai/v2"


class PixelLabError(RuntimeError):
    pass


class PixelLabClient:
    def __init__(self, api_key: Optional[str] = None, api_url: str = API_URL):
        self.api_key = api_key or os.environ.get("PIXELLAB_API_KEY")
        self.api_url = api_url.rstrip("/")

    def _request(self, method: str, path: str, payload: Optional[dict] = None) -> dict:
        if not self.api_key:
            raise PixelLabError("PIXELLAB_API_KEY is not configured.")
        request = Request(
            f"{self.api_url}{path}",
            data=json.dumps(payload).encode() if payload else None,
            headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
            method=method,
        )
        try:
            with urlopen(request, timeout=125) as response:
                return json.loads(response.read().decode())
        except HTTPError as exc:
            raise PixelLabError(f"PixelLab API returned HTTP {exc.code}: {exc.read().decode(errors='replace')}") from exc
        except (URLError, TimeoutError) as exc:
            raise PixelLabError(f"Could not reach PixelLab API: {exc}") from exc

    @staticmethod
    def _reference(value: str) -> dict:
        if value.startswith("data:"):
            raw = base64.b64decode(value.split(",", 1)[1])
        elif value.startswith(("https://", "http://")):
            with urlopen(value, timeout=60) as response:
                raw = response.read()
        else:
            with open(os.path.expanduser(value), "rb") as image_file:
                raw = image_file.read()
        with Image.open(BytesIO(raw)) as image:
            image = image.convert("RGBA")
            image.thumbnail((512, 512), Image.Resampling.LANCZOS)
            width, height = image.size
            output = BytesIO()
            image.save(output, format="PNG")
        encoded = "data:image/png;base64," + base64.b64encode(output.getvalue()).decode()
        return {"image": {"base64": encoded}, "width": width, "height": height}

    def create_image(self, description: str, width: int, height: int, *, seed: Optional[int] = None,
                     reference_images: Optional[list[str]] = None, no_background: bool = True) -> dict:
        refs = [self._reference(ref) for ref in (reference_images or [])]
        if refs:
            payload: dict[str, Any] = {
                "description": description, "image_size": {"width": width, "height": height},
                "style_images": refs[:4], "no_background": no_background,
            }
            if seed is not None:
                payload["seed"] = seed
            return self._request("POST", "/generate-with-style-v2", payload)
        payload = {"description": description, "image_size": {"width": width, "height": height},
                   "no_background": no_background}
        if seed is not None:
            payload["seed"] = seed
        return self._request("POST", "/generate-image-v2", payload)

    def get_status(self, job_id: str) -> dict:
        return self._request("GET", f"/background-jobs/{job_id}")

    def wait_for_completion(self, job_id: str, timeout_seconds: int = 115) -> dict:
        deadline = time.monotonic() + timeout_seconds
        while True:
            status = self.get_status(job_id)
            state = status.get("status", "").lower()
            if state == "completed":
                return status
            if state in {"failed", "error", "cancelled"}:
                raise PixelLabError(str(status.get("last_response") or state))
            if time.monotonic() >= deadline:
                return status
            time.sleep(2)

    @staticmethod
    def image_sources(response: dict) -> list[str]:
        sources: list[str] = []
        def visit(value: Any) -> None:
            if isinstance(value, dict):
                for key, item in value.items():
                    if key in {"url", "base64", "b64_json"} and isinstance(item, str):
                        sources.append(item)
                    else:
                        visit(item)
            elif isinstance(value, list):
                for item in value:
                    visit(item)
        visit(response.get("last_response", response))
        return list(dict.fromkeys(sources))