"""
Minimal PixelLab v2 client used by pixellab_library.py.

- Token from PIXELLAB_API_KEY (environment or .env next to this file).
- Retries rate limits (429/529) with backoff, polls async jobs (edit endpoints return a job id),
  and runs one job at a time across all processes (the account's concurrency limit is 1).
- Appends one line per paid call to pixellab_src/ledger.jsonl so spend can be tracked across
  runs and parallel workers: `python pixellab_library.py cost`.
"""
import base64
import fcntl
import io
import json
import os
import random
import sys
import time
from pathlib import Path

import requests
from PIL import Image

API = "https://api.pixellab.ai/v2"
LEDGER = Path(__file__).with_name("pixellab_src") / "ledger.jsonl"
LOCK = LEDGER.with_name(".api.lock")
RETRIES = 10
JOB_POLL_SECONDS = 4
JOB_TIMEOUT_SECONDS = 900


class ApiError(Exception):
    pass


def load_key():
    key = os.environ.get("PIXELLAB_API_KEY")
    env = Path(__file__).with_name(".env")
    if not key and env.exists():
        for line in env.read_text().splitlines():
            name, sep, value = line.strip().partition("=")
            if sep and name.strip() == "PIXELLAB_API_KEY":
                key = value.strip().strip("'\"")
    if not key:
        sys.exit("PIXELLAB_API_KEY not set (environment or .env)")
    return key


def b64(img):
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return {"type": "base64", "base64": base64.b64encode(buf.getvalue()).decode(), "format": "png"}


def decode(field):
    data = field["base64"]
    if data.startswith("data:"):
        data = data.split(",", 1)[1]
    return Image.open(io.BytesIO(base64.b64decode(data))).convert("RGBA")


def find_image(obj):
    """First image in a response: a Base64Image dict or a data: URL, at any depth."""
    if isinstance(obj, dict):
        if isinstance(obj.get("base64"), str):
            return decode(obj)
        obj = list(obj.values())
    if isinstance(obj, list):
        for v in obj:
            img = find_image(v)
            if img is not None:
                return img
    if isinstance(obj, str) and obj.startswith("data:image"):
        return decode({"base64": obj})
    return None


def redact(payload):
    return {k: ("<image>" if isinstance(v, dict) and v.get("type") == "base64" else v) for k, v in payload.items()}


class Client:
    def __init__(self):
        self.session = requests.Session()
        self.session.headers["Authorization"] = f"Bearer {load_key()}"

    def call(self, method, path, payload=None):
        for attempt in range(RETRIES):
            try:
                r = self.session.request(method, f"{API}/{path}", json=payload, timeout=300)
            except requests.RequestException as e:
                if attempt == RETRIES - 1:
                    raise ApiError(str(e))
                time.sleep(5 + attempt * 5)
                continue
            if r.status_code in (429, 529, 502, 503) and attempt < RETRIES - 1:
                time.sleep(min(60, 6 * (attempt + 1)) + random.uniform(0, 4))
                continue
            if not r.ok:
                raise ApiError(f"{path}: {r.status_code} {r.text[:400]}")
            return r.json()

    def balance(self):
        return self.call("GET", "balance")

    def run(self, endpoint, payload, label):
        """POST, wait for the job if there is one, return (image, usage). Logs the call to the ledger.

        The account allows one job at a time, so jobs from every process (parallel workers, other
        terminals) queue on a file lock instead of failing with 429."""
        LOCK.parent.mkdir(parents=True, exist_ok=True)
        with LOCK.open("a") as lock:
            fcntl.flock(lock, fcntl.LOCK_EX)
            try:
                return self._run(endpoint, payload, label)
            finally:
                fcntl.flock(lock, fcntl.LOCK_UN)

    def _run(self, endpoint, payload, label):
        resp = self.call("POST", endpoint, payload)
        usage = resp.get("usage")
        job = resp.get("background_job_id")
        deadline = time.time() + JOB_TIMEOUT_SECONDS
        while job:
            if time.time() > deadline:
                raise ApiError(f"job {job} still running after {JOB_TIMEOUT_SECONDS}s")
            time.sleep(JOB_POLL_SECONDS)
            status = self.call("GET", f"background-jobs/{job}")
            if status["status"] == "failed":
                raise ApiError(f"job {job} failed: {json.dumps(status.get('last_response'))[:400]}")
            if status["status"] == "completed":
                resp, usage, job = status.get("last_response") or {}, status.get("usage") or usage, None
        img = find_image({k: v for k, v in resp.items() if k != "usage"})
        LEDGER.parent.mkdir(parents=True, exist_ok=True)
        with LEDGER.open("a") as f:
            f.write(json.dumps({"time": round(time.time()), "label": label, "endpoint": endpoint,
                                "usage": usage, "ok": img is not None}) + "\n")
        if img is None:
            raise ApiError(f"no image in response: {json.dumps(redact(resp))[:400]}")
        return img, usage