from __future__ import annotations

from pathlib import Path
from typing import Any, Callable

import numpy as np
from PIL import Image


def _align_spatial_dim_for_flux_fill(n: int, multiple: int = 16) -> int:
    """FLUX Fill 要求 H/W 为 (vae_scale_factor*2) 的倍数，常见为 16。"""
    n = int(max(n, multiple))
    return n - (n % multiple)


def _flux_clip_prompt_and_t5_full(pipe: Any, text_full: str) -> tuple[str, str]:
    """
    CLIP 仅约 77 token：用与管线一致的 tokenizer 截断得到 prompt（池化向量）。
    完整文本走 prompt_2 → T5（max_sequence_length 可达 512），长中文不再整段喂给 CLIP。
    """
    text_full = (text_full or "").strip()
    if not text_full:
        d = "Chinese ink painting background, natural texture, seamless fill, no figures"
        return d, d
    tok = getattr(pipe, "tokenizer", None)
    if tok is None:
        return text_full[:120], text_full
    ml = int(getattr(tok, "model_max_length", 77) or 77)
    enc = tok([text_full], padding="max_length", max_length=ml, truncation=True, return_tensors="pt")
    clip_s = tok.batch_decode(enc.input_ids, skip_special_tokens=True)[0].strip()
    if not clip_s:
        clip_s = text_full[:80]
    return clip_s, text_full


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

    root = Path(model_id.strip()).expanduser()
    # 本机 diffusers 目录（含 model_index.json）时不访问 Hub，避免自动下载。
    local_files_only = root.is_dir() and (root / "model_index.json").is_file()

    pipe = FluxFillPipeline.from_pretrained(
        str(root),
        torch_dtype=torch_dtype,
        token=None,
        local_files_only=local_files_only,
    )

    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 = _align_spatial_dim_for_flux_fill(run_w)
        run_h = _align_spatial_dim_for_flux_fill(run_h)

        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_full = (prompt or "").strip()
        if not p_full:
            p_full = "Chinese ink painting background, natural texture, seamless fill, no figures"
        p_clip, p_t5 = _flux_clip_prompt_and_t5_full(pipe, p_full)

        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_clip,
            prompt_2=p_t5,
            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