src/compliment_forest/backends/image.py

from __future__ import annotations

import base64
import hashlib
import hmac
import io
import random
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Protocol

import httpx
from huggingface_hub import InferenceClient
from PIL import Image, ImageDraw, ImageFilter, ImageOps

from compliment_forest.schema import ForestStyle


@dataclass(frozen=True)
class StyleProfile:
    label: str
    trigger: str
    prefix: str
    suffix: str
    negative_prompt: str


_COMMON_NEGATIVE = (
    "text, letters, watermark, logo, frame, collage, duplicate animal, extra limbs, "
    "deformed face, distorted hands, photorealistic stock photo, 3d render"
)

STYLE_PROFILES: dict[str, StyleProfile] = {
    "watercolor": StyleProfile(
        label="Watercolor Storybook",
        trigger="cmprst_watercolor",
        prefix=(
            "cmprst_watercolor, soft watercolor storybook illustration, loose wet-on-wet "
            "washes, visible cold-press paper grain, feathered edges, "
        ),
        suffix=", warm dappled light, muted natural palette, generous negative space",
        negative_prompt=f"hard outlines, neon, oversaturated, {_COMMON_NEGATIVE}",
    ),
    "paper_cut": StyleProfile(
        label="Layered Paper Cut",
        trigger="cmprst_papercut",
        prefix=(
            "cmprst_papercut, handcrafted layered paper-cut illustration, deckled cut "
            "edges, overlapping botanical silhouettes, tactile paper fibers, "
        ),
        suffix=", soft dimensional shadows, moss and ochre palette, clean composition",
        negative_prompt=f"paint splashes, glossy plastic, neon, {_COMMON_NEGATIVE}",
    ),
    "moonlit_gouache": StyleProfile(
        label="Moonlit Gouache",
        trigger="cmprst_moonlit",
        prefix=(
            "cmprst_moonlit, moonlit gouache storybook painting, opaque velvety brushwork, "
            "deep indigo woodland, silver rim light, "
        ),
        suffix=", quiet luminous atmosphere, restrained jewel tones, simple composition",
        negative_prompt=f"daylight, washed-out contrast, digital gradients, {_COMMON_NEGATIVE}",
    ),
    "botanical_ink": StyleProfile(
        label="Botanical Ink Wash",
        trigger="cmprst_inkwash",
        prefix=(
            "cmprst_inkwash, botanical ink-wash illustration, expressive sepia and moss "
            "brush lines, diluted ink blooms, cream washi paper, "
        ),
        suffix=", sparse wildflower details, calm asymmetry, airy negative space",
        negative_prompt=f"thick cartoon outlines, saturated blocks, neon, {_COMMON_NEGATIVE}",
    ),
}

STYLE_LORA_V1_IDS = {
    "watercolor": "thangvip/compliment-forest-watercolor-flux-lora",
    "paper_cut": "thangvip/compliment-forest-paper-cut-flux-lora",
    "moonlit_gouache": "thangvip/compliment-forest-moonlit-gouache-flux-lora",
    "botanical_ink": "thangvip/compliment-forest-botanical-ink-flux-lora",
}

STYLE_LORA_IDS = {
    style: f"{model_id}-v2"
    for style, model_id in STYLE_LORA_V1_IDS.items()
}
STYLE_ADAPTER_WEIGHT = 0.45

STYLE_PREFIX = (
    "cmprst_forest, cmprst_watercolor, soft watercolor storybook illustration, "
    "loose wet-on-wet washes, "
    "visible cold-press paper grain, no hard outlines, soft feathered edges, "
)
STYLE_SUFFIX = (
    ", warm dappled light, kind expression, lots of negative space, muted warm palette"
)
NEGATIVE_PROMPT = (
    "hard outlines, photorealistic, 3d render, neon, oversaturated, busy background, "
    "text, watermark, deformed"
)


class ImageBackend(Protocol):
    def generate(self, prompt: str, seed: int, style: ForestStyle) -> str: ...


def resolve_style(style: ForestStyle, seed: int) -> str:
    if style != "surprise":
        return style
    style_ids = tuple(STYLE_PROFILES)
    return style_ids[seed % len(style_ids)]


def compose_flux_prompt(
    creature_prompt: str,
    style: ForestStyle = "watercolor",
    *,
    seed: int = 3407,
) -> str:
    resolved = resolve_style(style, seed)
    profile = STYLE_PROFILES[resolved]
    subject = creature_prompt.strip()
    style_language = profile.prefix.removeprefix(f"{profile.trigger}, ").rstrip(", ")
    return (
        f"{profile.trigger}, primary subject: {subject}. "
        "Preserve this exact subject, action, and setting as the clear focal point. "
        "No text, lettering, signature, logo, or watermark anywhere. "
        f"{style_language}{profile.suffix}"
    )


def image_to_data_uri(image: Image.Image) -> str:
    buffer = io.BytesIO()
    image.save(buffer, format="PNG", optimize=True)
    encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
    return f"data:image/png;base64,{encoded}"


def clean_generated_image(image: Image.Image) -> Image.Image:
    """Trim common synthetic signature zones while preserving output dimensions."""
    width, height = image.size
    crop = image.crop(
        (
            round(width * 0.035),
            round(height * 0.02),
            round(width * 0.965),
            round(height * 0.935),
        )
    )
    return ImageOps.fit(
        crop,
        image.size,
        method=Image.Resampling.LANCZOS,
    )


def sign_modal_request(key: str, prompt: str, seed: int, style: str) -> str:
    message = f"{prompt}\n{seed}\n{style}".encode()
    return hmac.new(key.encode(), message, hashlib.sha256).hexdigest()


class DemoImageBackend:
    """Creates deterministic procedural placeholders for offline development."""

    def __init__(
        self,
        width: int = 768,
        height: int = 576,
        asset_dir: str | Path | None = None,
    ) -> None:
        self.width = width
        self.height = height
        self.asset_dir = Path(asset_dir) if asset_dir is not None else None

    def generate(
        self,
        prompt: str,
        seed: int,
        style: ForestStyle = "surprise",
    ) -> str:
        resolved_style = resolve_style(style, seed)
        prompt_seed = int.from_bytes(
            hashlib.sha256(f"{resolved_style}:{prompt}".encode()).digest()[:4],
            byteorder="big",
        )
        rng = random.Random(seed ^ prompt_seed)
        image = Image.new("RGB", (self.width, self.height), "#f5f0df")
        wash = Image.new("RGBA", image.size, (0, 0, 0, 0))
        draw = ImageDraw.Draw(wash, "RGBA")
        palette = [
            (124, 145, 112, 28),
            (77, 105, 78, 24),
            (208, 151, 139, 22),
            (225, 184, 101, 22),
            (132, 172, 181, 18),
        ]
        if resolved_style == "moonlit_gouache":
            palette = [(36, 49, 92, 42), (70, 82, 124, 36), (194, 184, 151, 26)]
        elif resolved_style == "paper_cut":
            palette = [(83, 104, 61, 40), (205, 151, 96, 34), (201, 143, 131, 28)]
        elif resolved_style == "botanical_ink":
            palette = [(73, 79, 55, 32), (120, 96, 67, 28), (148, 161, 132, 20)]
        for _ in range(100):
            color = rng.choice(palette)
            x = rng.randint(-100, self.width)
            y = rng.randint(-80, self.height)
            radius_x = rng.randint(35, 160)
            radius_y = rng.randint(25, 110)
            draw.ellipse(
                (x - radius_x, y - radius_y, x + radius_x, y + radius_y),
                fill=color,
            )
        wash = wash.filter(ImageFilter.GaussianBlur(radius=22))
        image = Image.alpha_composite(image.convert("RGBA"), wash)

        creature = Image.new("RGBA", image.size, (0, 0, 0, 0))
        creature_draw = ImageDraw.Draw(creature, "RGBA")
        center_x = self.width // 2 + rng.randint(-30, 30)
        ground_y = int(self.height * 0.72)
        body_color = rng.choice(
            [(169, 105, 69, 205), (128, 105, 78, 205), (154, 135, 96, 205)]
        )
        shadow = (76, 91, 67, 35)
        creature_draw.ellipse(
            (center_x - 145, ground_y - 10, center_x + 145, ground_y + 38),
            fill=shadow,
        )
        creature_draw.ellipse(
            (center_x - 96, ground_y - 170, center_x + 90, ground_y + 5),
            fill=body_color,
        )
        creature_draw.ellipse(
            (center_x - 72, ground_y - 245, center_x + 74, ground_y - 100),
            fill=body_color,
        )
        creature_draw.polygon(
            [
                (center_x - 62, ground_y - 210),
                (center_x - 88, ground_y - 290),
                (center_x - 20, ground_y - 235),
            ],
            fill=body_color,
        )
        creature_draw.polygon(
            [
                (center_x + 58, ground_y - 210),
                (center_x + 88, ground_y - 290),
                (center_x + 20, ground_y - 235),
            ],
            fill=body_color,
        )
        creature_draw.ellipse(
            (center_x - 38, ground_y - 195, center_x - 15, ground_y - 169),
            fill=(51, 53, 44, 220),
        )
        creature_draw.ellipse(
            (center_x + 18, ground_y - 195, center_x + 41, ground_y - 169),
            fill=(51, 53, 44, 220),
        )
        creature_draw.ellipse(
            (center_x - 6, ground_y - 163, center_x + 9, ground_y - 150),
            fill=(72, 57, 47, 190),
        )
        creature = creature.filter(ImageFilter.GaussianBlur(radius=1.2))
        image = Image.alpha_composite(image, creature)

        grain = Image.new("RGBA", image.size, (0, 0, 0, 0))
        grain_pixels = grain.load()
        for y in range(self.height):
            for x in range(self.width):
                value = rng.randint(0, 10)
                grain_pixels[x, y] = (75, 66, 50, value)
        image = Image.alpha_composite(image, grain)
        return image_to_data_uri(image.convert("RGB"))


class FluxImageBackend:
    """Lazy FLUX.1-dev pipeline that loads a local or Hub LoRA."""

    def __init__(
        self,
        model_id: str = "black-forest-labs/FLUX.1-dev",
        lora_id: str = "build-small-hackathon/compliment-forest-flux-lora",
        *,
        local_files_only: bool = False,
        width: int = 768,
        height: int = 768,
        steps: int = 28,
    ) -> None:
        self.model_id = model_id
        self.lora_id = lora_id
        self.local_files_only = local_files_only
        self.width = width
        self.height = height
        self.steps = steps
        self._pipeline = None

    def _load(self):
        if self._pipeline is not None:
            return self._pipeline
        import torch
        from diffusers import FluxPipeline

        pipeline = FluxPipeline.from_pretrained(
            self.model_id,
            torch_dtype=torch.bfloat16,
            local_files_only=self.local_files_only,
        )
        pipeline.load_lora_weights(
            self.lora_id,
            local_files_only=self.local_files_only,
        )
        pipeline.enable_model_cpu_offload()
        self._pipeline = pipeline
        return pipeline

    def generate(
        self,
        prompt: str,
        seed: int,
        style: ForestStyle = "watercolor",
    ) -> str:
        import torch

        pipeline = self._load()
        generator = torch.Generator(device="cpu").manual_seed(seed)
        image = pipeline(
            prompt=compose_flux_prompt(prompt, style, seed=seed),
            negative_prompt=NEGATIVE_PROMPT,
            width=self.width,
            height=self.height,
            num_inference_steps=self.steps,
            guidance_scale=3.5,
            generator=generator,
        ).images[0]
        return image_to_data_uri(clean_generated_image(image.convert("RGB")))


class HfInferenceImageBackend:
    """Fresh FLUX.1-schnell generation through Hugging Face Inference Providers."""

    def __init__(
        self,
        model: str = "black-forest-labs/FLUX.1-schnell",
        *,
        client: Any | None = None,
        provider: str = "auto",
        width: int = 768,
        height: int = 768,
        steps: int = 4,
        timeout: float = 180,
    ) -> None:
        self.model = model
        self.client = client or InferenceClient(provider=provider, timeout=timeout)
        self.width = width
        self.height = height
        self.steps = steps

    def generate(
        self,
        prompt: str,
        seed: int,
        style: ForestStyle = "surprise",
    ) -> str:
        resolved = resolve_style(style, seed)
        profile = STYLE_PROFILES[resolved]
        image = self.client.text_to_image(
            compose_flux_prompt(prompt, resolved, seed=seed),
            model=self.model,
            negative_prompt=profile.negative_prompt,
            width=self.width,
            height=self.height,
            num_inference_steps=self.steps,
            guidance_scale=0.0,
            seed=seed,
        )
        return image_to_data_uri(clean_generated_image(image.convert("RGB")))


class ZeroGpuImageBackend:
    """Route resolved styles to a Space GPU function with a hosted fallback."""

    def __init__(
        self,
        generator: Callable[[str, int, str], str],
        *,
        fallback: ImageBackend | None = None,
    ) -> None:
        self.generator = generator
        self.fallback = fallback

    def generate(
        self,
        prompt: str,
        seed: int,
        style: ForestStyle = "surprise",
    ) -> str:
        resolved = resolve_style(style, seed)
        try:
            return self.generator(prompt, seed, resolved)
        except Exception:
            if self.fallback is None:
                raise
            return self.fallback.generate(prompt, seed, resolved)


class ModalImageBackend:
    """Call the private-token Modal service that hosts the trained adapters."""

    def __init__(
        self,
        endpoint: str,
        signing_key: str,
        *,
        client: Any | None = None,
        fallback: ImageBackend | None = None,
        timeout: float = 600,
    ) -> None:
        self.endpoint = endpoint
        self.signing_key = signing_key
        self.client = client or httpx.Client(timeout=timeout, follow_redirects=True)
        self.fallback = fallback

    def generate(
        self,
        prompt: str,
        seed: int,
        style: ForestStyle = "surprise",
    ) -> str:
        resolved = resolve_style(style, seed)
        try:
            response = self.client.post(
                self.endpoint,
                json={
                    "prompt": prompt,
                    "seed": seed,
                    "style": resolved,
                    "signature": sign_modal_request(
                        self.signing_key,
                        prompt,
                        seed,
                        resolved,
                    ),
                },
            )
            response.raise_for_status()
            image = response.json()["image"]
            if not isinstance(image, str) or not image.startswith("data:image/png;base64,"):
                raise ValueError("Modal image response is not a PNG data URI")
            return image
        except Exception:
            if self.fallback is None:
                raise
            return self.fallback.generate(prompt, seed, resolved)