"""
AI Room Designer — generation backends (Stage 3).

One contract for every backend:  edit(before_path, prompt, out_path) -> out_path
(raises BackendError on failure). The pipeline and QC retry loop don't care
which backend produced the image.

  GeminiBackend   LIVE — Google Gemini image-edit API (generateContent REST,
                  x-goog-api-key header, key from env via config.api_key(),
                  never hardcoded, never logged).
  MockBackend     RIG/DRY-RUN — deterministic local transforms standing in for
                  the API so the full request/response + QC path can be proven
                  with no key and $0 spend. Behaviours: "grade" (warm colour
                  grade, structure untouched -> should PASS the gate) and
                  "reframe" (crop/shift, the known Budget Refresh v1 failure
                  pattern -> should FAIL the gate). Scriptable via
                  ROOM_DESIGNER_MOCK_SCRIPT or constructor arg.
  StandinBackend  The Stage 1/2 operator-run Muse media pipeline: writes a
                  PENDING marker and raises NeedsOperator — job stays queued
                  for the operator exactly as before (fallback preserved).
"""
import base64
import json
import os
import urllib.request
import urllib.error
from pathlib import Path

import config


class BackendError(Exception):
    pass


class NeedsOperator(BackendError):
    pass


# ---------------------------------------------------------------- Gemini ----
class GeminiBackend:
    name = "gemini"

    def __init__(self, key=None, model=None, transport=None, image_size=None):
        self.key = key if key is not None else config.api_key()
        self.model = model or config.model()
        self.image_size = image_size or config.image_size()
        self.transport = transport or self._http_transport
        self.last_usage = None

    @staticmethod
    def _http_transport(url, headers, payload):
        req = urllib.request.Request(
            url, data=json.dumps(payload).encode(),
            headers={"Content-Type": "application/json", **headers}, method="POST")
        try:
            with urllib.request.urlopen(req, timeout=180) as r:
                return json.loads(r.read().decode())
        except urllib.error.HTTPError as e:
            raise BackendError(f"Gemini API HTTP {e.code}: {e.read().decode()[:300]}")
        except Exception as e:  # network, timeout, DNS
            raise BackendError(f"Gemini API request failed: {e}")

    def build_request(self, before_path, prompt):
        img = Path(before_path).read_bytes()
        payload = {
            "contents": [{
                "role": "user",
                "parts": [
                    {"inline_data": {"mime_type": "image/jpeg",
                                     "data": base64.b64encode(img).decode()}},
                    {"text": prompt},
                ],
            }],
            "generationConfig": {
                "responseModalities": ["TEXT", "IMAGE"],
                "imageConfig": {"imageSize": self.image_size},
            },
        }
        url = config.GEMINI_ENDPOINT.format(model=self.model)
        return url, {"x-goog-api-key": self.key}, payload

    def edit(self, before_path, prompt, out_path):
        if not self.key:
            raise BackendError("no API key configured (ROOM_DESIGNER_API_KEY)")
        url, headers, payload = self.build_request(before_path, prompt)
        resp = self.transport(url, headers, payload)
        self.last_usage = resp.get("usageMetadata")
        for cand in resp.get("candidates", []):
            for part in cand.get("content", {}).get("parts", []):
                data = (part.get("inline_data") or part.get("inlineData") or {}).get("data")
                if data:
                    Path(out_path).write_bytes(base64.b64decode(data))
                    return out_path
        raise BackendError("Gemini response contained no image part")


# ------------------------------------------------------------------ Mock ----
class MockBackend:
    """Deterministic local stand-in. `script` = behaviours cycled per call."""
    name = "mock"

    def __init__(self, script=None):
        if script is None:
            script = os.environ.get("ROOM_DESIGNER_MOCK_SCRIPT", "grade")
        self.script = [s.strip() for s in script.split(",") if s.strip()] or ["grade"]
        self.calls = 0

    def edit(self, before_path, prompt, out_path):
        from PIL import Image, ImageEnhance
        behaviour = self.script[min(self.calls, len(self.script) - 1)]
        self.calls += 1
        img = Image.open(before_path).convert("RGB")
        if behaviour == "reframe":
            # The Stage 2 Budget Refresh v1 failure pattern: camera reframes —
            # centre-crop + shift, then back to frame size. Shell lines move.
            w, h = img.size
            img = img.crop((int(w * 0.10), int(h * 0.06), int(w * 0.94), int(h * 0.90)))
            img = img.resize((w, h))
        # Warm grade on top of either behaviour (the "redesign" signal).
        img = ImageEnhance.Color(img).enhance(1.18)
        img = ImageEnhance.Brightness(img).enhance(1.04)
        r, g, b = img.split()
        r = r.point(lambda v: min(255, int(v * 1.05)))
        img = Image.merge("RGB", (r, g, b))
        img.save(out_path)
        return out_path


# --------------------------------------------------------------- Standin ------
class StandinBackend:
    """Stage 1/2 fallback: operator-run media pipeline (job stays queued)."""
    name = "standin"

    def edit(self, before_path, prompt, out_path):
        raise NeedsOperator(
            "standin backend: generation is operator-run; job stays queued-on-rig")


def get_backend(provider=None, **kw):
    p = provider or config.provider()
    if p == "gemini":
        return GeminiBackend(**kw)
    if p == "mock":
        return MockBackend(**{k: v for k, v in kw.items() if k == "script"})
    if p == "standin":
        return StandinBackend()
    raise BackendError(f"unknown provider '{p}'")