from __future__ import annotations

import os
from typing import Any, Callable

import numpy as np
from PIL import Image


def make_flux_fill_predictor(
    model_id: str,
    device: str | None = None,
) -> Callable[[Image.Image, Image.Image, str, str], Image.Image]:
    import torch
    from diffusers import FluxFillPipeline

    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"

    if device == "cuda" and torch.cuda.is_bf16_supported():
        torch_dtype = torch.bfloat16
    elif device == "cuda":
        torch_dtype = torch.float16
    else:
        torch_dtype = torch.float32

    hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")

    pipe = FluxFillPipeline.from_pretrained(
        model_id,
        torch_dtype=torch_dtype,
        token=hf_token,
    )

    if device == "cuda":
        try:
            pipe.enable_model_cpu_offload()
        except Exception:
            try:
                pipe.enable_vae_tiling()
            except Exception:
                pass
            pipe.to("cuda")
    else:
        pipe.to(device)

    def _predict(
        image: Image.Image,
        mask: Image.Image,
        prompt: str,
        negative_prompt: str = "",
        strength: float = 1.0,
        guidance_scale: float = 30.0,
        num_inference_steps: int = 28,
        max_side: int = 1024,
        max_sequence_length: int = 512,
        **_: Any,
    ) -> Image.Image:
        _ = negative_prompt
        image = image.convert("RGB")
        mask = mask.convert("L")

        orig_w, orig_h = image.size
        run_w, run_h = orig_w, orig_h
        if max(orig_w, orig_h) > max_side:
            scale = max_side / float(max(orig_w, orig_h))
            run_w = int(round(orig_w * scale))
            run_h = int(round(orig_h * scale))
        run_w = max(8, run_w - (run_w % 8))
        run_h = max(8, run_h - (run_h % 8))

        if (run_w, run_h) != (orig_w, orig_h):
            image_run = image.resize((run_w, run_h), resample=Image.BICUBIC)
            mask_run = mask.resize((run_w, run_h), resample=Image.NEAREST)
        else:
            image_run = image
            mask_run = mask

        p = (prompt or "").strip()
        if not p:
            p = "Chinese ink painting background, natural texture, seamless fill, no figures"

        st = float(strength)
        st = min(1.0, max(0.05, st))

        if device == "cuda":
            try:
                torch.cuda.empty_cache()
            except Exception:
                pass

        out = pipe(
            prompt=p,
            image=image_run,
            mask_image=mask_run,
            height=run_h,
            width=run_w,
            strength=st,
            guidance_scale=float(guidance_scale),
            num_inference_steps=int(num_inference_steps),
            max_sequence_length=int(max_sequence_length),
        ).images[0]

        out = out.convert("RGB")
        if out.size != (orig_w, orig_h):
            out = out.resize((orig_w, orig_h), resample=Image.BICUBIC)
        return out

    return _predict