import io
import os

import modal


APP_NAME = "virtual-characters-image"
GPU = os.environ.get("VC_IMAGE_GPU", "H100")
MODEL_ID = os.environ.get("VC_IMAGE_MODEL", "black-forest-labs/FLUX.1-schnell")
MODEL_DIR = "/root/.cache/huggingface"
HF_SECRET_NAME = os.environ.get("VC_HF_SECRET_NAME", "hf-token")
HF_SECRETS = [] if os.environ.get("VC_SKIP_HF_SECRET") == "1" else [modal.Secret.from_name(HF_SECRET_NAME)]


image = (
    modal.Image.from_registry("nvidia/cuda:12.9.0-devel-ubuntu22.04", add_python="3.11")
    .entrypoint([])
    .apt_install("git", "libglib2.0-0", "libsm6", "libxrender1", "libxext6", "ffmpeg", "libgl1")
    .uv_pip_install(
        "accelerate>=1.8.0",
        "diffusers>=0.35.0",
        "fastapi[standard]>=0.115.0",
        "huggingface-hub>=0.36.0",
        "safetensors>=0.5.0",
        "sentencepiece>=0.2.0",
        "torch>=2.7.0",
        "transformers>=4.57.0",
    )
    .env({"HF_HUB_CACHE": MODEL_DIR, "HF_XET_HIGH_PERFORMANCE": "1"})
)

hf_cache = modal.Volume.from_name("vc-hf-cache", create_if_missing=True)
app = modal.App(APP_NAME, image=image)


@app.cls(
    gpu=GPU,
    scaledown_window=60,
    timeout=60 * 20,
    secrets=HF_SECRETS,
    volumes={MODEL_DIR: hf_cache},
)
class CharacterImage:
    def _ensure_loaded(self):
        if getattr(self, "pipe", None) is not None:
            return
        import torch
        from diffusers import FluxPipeline

        self.pipe = FluxPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16).to("cuda")

    @modal.method()
    def health(self) -> dict:
        return {"ok": True, "model": MODEL_ID, "gpu": GPU, "loaded": getattr(self, "pipe", None) is not None}

    @modal.method()
    def generate(self, prompt: str, steps: int = 4, seed: int | None = None) -> bytes:
        self._ensure_loaded()
        import torch

        generator = None
        if seed is not None:
            generator = torch.Generator(device="cuda").manual_seed(seed)
        image = self.pipe(prompt, num_inference_steps=steps, output_type="pil", generator=generator).images[0]
        buffer = io.BytesIO()
        image.save(buffer, format="PNG")
        buffer.seek(0)
        return buffer.read()

    @modal.fastapi_endpoint(method="POST")
    async def character_image(self, request):
        from fastapi.responses import Response

        payload = await request.json()
        data = self.generate.local(
            prompt=payload["prompt"],
            steps=int(payload.get("steps", 4)),
            seed=payload.get("seed"),
        )
        return Response(content=data, media_type="image/png")


@app.local_entrypoint()
def main(prompt: str = "original anime virtual character portrait, silver hair, teal eyes, soft sci-fi lighting", output_path: str = "/tmp/vc_character.png"):
    print(CharacterImage().health.remote())
    data = CharacterImage().generate.remote(prompt=prompt, steps=4, seed=42)
    with open(output_path, "wb") as f:
        f.write(data)
    print(f"Wrote {len(data)} bytes to {output_path}")