"""Bria's /v2/vgl/edit: the JSON instruction for a Fibo 2.0 edit, written with Fibo 2.0's VGL prompts.

It answers with the caption of the result and its `edit_instruction`, as a JSON object.
"""

import base64
import io
import os
import time

import requests
from PIL import Image

VGL_EDIT_URL = "https://engine.prod.bria-api.com/v2/vgl/edit"

MAX_REFERENCES = 4
REQUEST_TIMEOUT = 180
RETRIES = 3
RETRY_DELAY = 3


class BriaAPIError(RuntimeError):
    """Raised with a message safe to show a user."""


def _encode(image: Image.Image) -> str:
    buffer = io.BytesIO()
    image.convert("RGB").save(buffer, format="PNG")
    return base64.b64encode(buffer.getvalue()).decode("utf-8")


def _post(url: str, payload: dict, field: str) -> dict:
    token = os.environ.get("BRIA_API_TOKEN")
    if not token:
        raise BriaAPIError("BRIA_API_TOKEN is not set on this Space.")
    headers = {"Content-Type": "application/json", "api_token": token, "User-Agent": "BriaHFSpace/fibo-2.0-turbo"}

    # The /v2/vgl endpoints reject seed 0 with a 422; never send it.
    payload = {**payload, "seed": max(int(payload.get("seed", 1)), 1)}
    last_error = None
    for attempt in range(1, RETRIES + 1):
        try:
            response = requests.post(url, json=payload, headers=headers, timeout=REQUEST_TIMEOUT)
            response.raise_for_status()
            result = response.json().get("result", {})
            if field not in result:
                raise BriaAPIError(f"Prompt service returned no {field}.")
            return result[field]
        except requests.HTTPError as exc:
            # 4xx other than rate limiting will not change on a retry.
            status = exc.response.status_code if exc.response is not None else None
            if status is not None and 400 <= status < 500 and status != 429:
                raise BriaAPIError(f"Prompt service rejected the request ({status}).") from exc
            last_error = exc
        except requests.RequestException as exc:
            last_error = exc
        if attempt < RETRIES:
            time.sleep(RETRY_DELAY)
    raise BriaAPIError("Prompt service is unreachable. Try again in a moment.") from last_error


def vgl_edit(instruction, images, mask=None, seed=42) -> dict:
    """Instruction + references (+ mask) -> the edit's JSON. images[0] is the source, and they are
    named <image_1>..<image_N> in this order. A mask, white where the edit goes, needs one image."""
    if not images:
        raise BriaAPIError("Upload a source image first.")
    if len(images) > MAX_REFERENCES:
        raise BriaAPIError(f"At most {MAX_REFERENCES} images are supported.")
    if not instruction:
        raise BriaAPIError("Describe the edit you want.")

    payload = {
        "instruction": instruction,
        "images": [_encode(image) for image in images],
        "seed": seed,
        "sync": True,
    }
    if mask is not None:
        payload["mask"] = _encode(mask)
    return _post(VGL_EDIT_URL, payload, "vgl")