"""NFA Track R — Fun depth ControlNet ZeroGPU (VideoX-Fun / ALIMAMA).

REAL Fun ControlNet Union depth — NOT soft Flux2 image=depth (banned forever).

VRAM choice (documented):
  size=\"large\" (48GB, 1× Pro) + VideoX-Fun model_cpu_offload_and_qfloat8.
  Escalation if OOM: size=\"xlarge\" + model_cpu_offload (2× quota).
"""

from __future__ import annotations

import os
import traceback
from pathlib import Path
from typing import Optional

import gradio as gr
import spaces
import torch
from huggingface_hub import hf_hub_download, login, snapshot_download
from omegaconf import OmegaConf
from PIL import Image

HF_TOKEN = (
    os.environ.get("HF_TOKEN")
    or os.environ.get("HUGGINGFACE_HUB_TOKEN")
    or ""
).strip()
if HF_TOKEN:
    try:
        login(token=HF_TOKEN, add_to_git_credential=False)
    except Exception as exc:  # noqa: BLE001
        print(f"[nfa-fun-cn] HF login warning: {exc}", flush=True)

BASE_MODEL = os.environ.get(
    "NFA_FLUX2_MODEL_ID", "black-forest-labs/FLUX.2-dev"
).strip()
CN_REPO = os.environ.get(
    "NFA_FUN_CN_REPO", "alibaba-pai/FLUX.2-dev-Fun-Controlnet-Union"
).strip()
CN_FILE = os.environ.get(
    "NFA_FUN_CN_FILE", "FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors"
).strip()
# large=48GB 1×; set NFA_FUN_CN_GPU_SIZE=xlarge only after OOM on large
GPU_SIZE = (os.environ.get("NFA_FUN_CN_GPU_SIZE") or "large").strip().lower()
if GPU_SIZE not in ("large", "xlarge"):
    GPU_SIZE = "large"
GPU_DURATION = int(os.environ.get("NFA_FUN_CN_GPU_DURATION") or "300")
WEIGHT_DTYPE = torch.bfloat16
# VideoX-Fun official low-VRAM Fun CN mode for 48GB; offload-only on xlarge
MEM_MODE = (
    os.environ.get("NFA_FUN_CN_MEM_MODE")
    or (
        "model_cpu_offload"
        if GPU_SIZE == "xlarge"
        else "model_cpu_offload_and_qfloat8"
    )
).strip()

APP_DIR = Path(__file__).resolve().parent
CONFIG_PATH = APP_DIR / "config" / "flux2_control.yaml"
CACHE_ROOT = Path(
    os.environ.get("NFA_FUN_CN_CACHE")
    or (Path.home() / ".cache" / "nfa_fun_cn")
)
MODEL_DIR = CACHE_ROOT / "FLUX.2-dev"

_PIPE = None
_CN_FILE_PATH: Path | None = None
_WEIGHTS_READY = False


def _resolve_cn_path() -> Path:
    global _CN_FILE_PATH
    if _CN_FILE_PATH is not None and _CN_FILE_PATH.is_file():
        return _CN_FILE_PATH
    direct = CACHE_ROOT / CN_FILE
    if direct.is_file():
        _CN_FILE_PATH = direct
        return direct
    nested = list(CACHE_ROOT.rglob(CN_FILE))
    if nested:
        _CN_FILE_PATH = nested[0]
        return nested[0]
    raise FileNotFoundError(f"Fun CN weights missing: {CN_FILE}")


def _ensure_weights() -> None:
    """Download on CPU (must run outside @spaces.GPU so quota is not burned)."""
    global _WEIGHTS_READY, _CN_FILE_PATH
    if _WEIGHTS_READY and MODEL_DIR.is_dir():
        try:
            _resolve_cn_path()
            return
        except FileNotFoundError:
            pass
    CACHE_ROOT.mkdir(parents=True, exist_ok=True)
    token = HF_TOKEN or None
    print(f"[nfa-fun-cn] snapshot {BASE_MODEL} -> {MODEL_DIR}", flush=True)
    snapshot_download(
        repo_id=BASE_MODEL,
        local_dir=str(MODEL_DIR),
        local_dir_use_symlinks=False,
        token=token,
    )
    print(f"[nfa-fun-cn] download {CN_REPO}/{CN_FILE}", flush=True)
    path = hf_hub_download(
        repo_id=CN_REPO,
        filename=CN_FILE,
        local_dir=str(CACHE_ROOT),
        local_dir_use_symlinks=False,
        token=token,
    )
    _CN_FILE_PATH = Path(path)
    _WEIGHTS_READY = True
    print(f"[nfa-fun-cn] weights ready cn={_CN_FILE_PATH}", flush=True)


def _prep_depth(depth_image: Image.Image, width: int, height: int) -> Image.Image:
    img = depth_image.convert("RGB")
    if img.size != (width, height):
        img = img.resize((width, height), Image.Resampling.LANCZOS)
    return img


def _compose_prompt(positive: str, negative: str) -> tuple[str, str]:
    pos = (positive or "").strip()
    neg = (negative or "").strip() or " "
    return pos, neg


def get_pipe():
    """Build VideoX-Fun Flux2ControlPipeline once per warm process."""
    global _PIPE
    if _PIPE is not None:
        return _PIPE

    _ensure_weights()
    from diffusers import FlowMatchEulerDiscreteScheduler
    from safetensors.torch import load_file
    from videox_fun.models import (
        AutoencoderKLFlux2,
        Flux2ControlTransformer2DModel,
        Mistral3ForConditionalGeneration,
        PixtralProcessor,
    )
    from videox_fun.pipeline import Flux2ControlPipeline
    from videox_fun.utils.fp8_optimization import (
        convert_model_weight_to_float8,
        convert_weight_dtype_wrapper,
    )
    from videox_fun.utils.utils import get_image_latent

    # stash for generate
    get_pipe._get_image_latent = get_image_latent  # type: ignore[attr-defined]

    model_name = str(MODEL_DIR)
    cn_file = str(_resolve_cn_path())
    config = OmegaConf.load(str(CONFIG_PATH))
    print(
        f"[nfa-fun-cn] load Flux2ControlTransformer + Fun CN "
        f"mem={MEM_MODE} size={GPU_SIZE}",
        flush=True,
    )
    transformer = Flux2ControlTransformer2DModel.from_pretrained(
        model_name,
        subfolder="transformer",
        low_cpu_mem_usage=True,
        torch_dtype=WEIGHT_DTYPE,
        transformer_additional_kwargs=OmegaConf.to_container(
            config["transformer_additional_kwargs"]
        ),
    ).to(WEIGHT_DTYPE)

    state_dict = load_file(cn_file)
    state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
    missing, unexpected = transformer.load_state_dict(state_dict, strict=False)
    print(
        f"[nfa-fun-cn] Fun CN loaded missing={len(missing)} unexpected={len(unexpected)}",
        flush=True,
    )

    vae = AutoencoderKLFlux2.from_pretrained(model_name, subfolder="vae").to(
        WEIGHT_DTYPE
    )
    tokenizer = PixtralProcessor.from_pretrained(model_name, subfolder="tokenizer")
    text_encoder = Mistral3ForConditionalGeneration.from_pretrained(
        model_name,
        subfolder="text_encoder",
        torch_dtype=WEIGHT_DTYPE,
        low_cpu_mem_usage=True,
    )
    scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
        model_name, subfolder="scheduler"
    )
    pipeline = Flux2ControlPipeline(
        vae=vae,
        tokenizer=tokenizer,
        text_encoder=text_encoder,
        transformer=transformer,
        scheduler=scheduler,
    )

    device = "cuda" if torch.cuda.is_available() else "cpu"
    if MEM_MODE == "model_cpu_offload_and_qfloat8":
        convert_model_weight_to_float8(
            transformer,
            exclude_module_name=["img_in", "txt_in", "timestep"],
            device=device,
        )
        convert_weight_dtype_wrapper(transformer, WEIGHT_DTYPE)
        pipeline.enable_model_cpu_offload(device=device)
    elif MEM_MODE == "sequential_cpu_offload":
        pipeline.enable_sequential_cpu_offload(device=device)
    elif MEM_MODE == "model_cpu_offload":
        pipeline.enable_model_cpu_offload(device=device)
    else:
        pipeline.to(device=device)

    _PIPE = pipeline
    print("[nfa-fun-cn] Flux2ControlPipeline ready (REAL Fun CN)", flush=True)
    return _PIPE


@spaces.GPU(duration=GPU_DURATION, size=GPU_SIZE)
def _generate_still_gpu(
    positive: str,
    negative: str,
    depth_image: Image.Image,
    seed: int,
    width: int,
    height: int,
    steps: int,
    guidance: float,
    cn_strength: float,
) -> Image.Image:
    """GPU-billed Fun CN infer only — weights must already be on disk."""
    if torch.cuda.is_available():
        free, total = torch.cuda.mem_get_info()
        print(
            f"[nfa-fun-cn] cuda free={free/1e9:.1f}G total={total/1e9:.1f}G "
            f"duration={GPU_DURATION} size={GPU_SIZE} mem={MEM_MODE}",
            flush=True,
        )

    w = int(width) if width else 1216
    h = int(height) if height else 832
    w -= w % 16
    h -= h % 16
    prompt, neg = _compose_prompt(positive, negative)
    if not prompt:
        raise gr.Error("positive prompt is required")

    depth = _prep_depth(depth_image, w, h)
    pipe = get_pipe()
    get_image_latent = get_pipe._get_image_latent  # type: ignore[attr-defined]

    # VideoX-Fun control latents (NOT Flux2Pipeline soft image=)
    control_latent = get_image_latent(depth, sample_size=[h, w])[:, :, 0]
    inpaint_image = torch.zeros([1, 3, h, w])
    mask_image = torch.ones([1, 1, h, w]) * 255

    strength = float(cn_strength)
    if strength <= 0:
        strength = 0.75
    # ALIMAMA recommended band 0.65–0.80; allow Track R packet values
    strength = max(0.05, min(1.5, strength))

    device = "cuda" if torch.cuda.is_available() else "cpu"
    generator = torch.Generator(device=device).manual_seed(int(seed))
    print(
        f"[nfa-fun-cn] REAL Fun CN generate seed={seed} {w}x{h} steps={steps} "
        f"cn={strength} path=videox_fun_flux2_control",
        flush=True,
    )
    with torch.no_grad():
        out = pipe(
            prompt=prompt,
            negative_prompt=neg,
            height=h,
            width=w,
            generator=generator,
            guidance_scale=float(guidance),
            image=None,
            inpaint_image=inpaint_image,
            mask_image=mask_image,
            control_image=control_latent,
            num_inference_steps=int(steps),
            control_context_scale=strength,
        ).images
    return out[0]


def generate_still(
    positive: str,
    negative: str = "",
    depth_image: Optional[Image.Image] = None,
    seed: int = 42,
    width: int = 1216,
    height: int = 832,
    steps: int = 28,
    guidance: float = 4.0,
    cn_strength: float = 0.75,
) -> Image.Image:
    """CPU download + GPU Fun CN. Soft image=depth is never used."""
    try:
        if depth_image is None:
            raise gr.Error(
                "FUN_CN_REQUIRES_DEPTH: depth_image is required for real Fun ControlNet."
            )
        _ensure_weights()
        return _generate_still_gpu(
            positive,
            negative or "",
            depth_image,
            int(seed),
            int(width),
            int(height),
            int(steps),
            float(guidance),
            float(cn_strength),
        )
    except gr.Error:
        raise
    except Exception as exc:  # noqa: BLE001
        tb = traceback.format_exc()
        print(tb, flush=True)
        raise gr.Error(f"{type(exc).__name__}: {exc}\n\n{tb[-2500:]}") from exc


with gr.Blocks(title="NFA Track R FLUX.2 Fun CN ZeroGPU") as demo:
    gr.Markdown(
        "## NFA Track R — **Real Fun depth ControlNet** (ZeroGPU)\n"
        f"- Stack: VideoX-Fun `Flux2ControlPipeline` + `{CN_FILE}`\n"
        f"- Base: `{BASE_MODEL}`\n"
        f"- GPU: `size={GPU_SIZE}` duration={GPU_DURATION}s mem=`{MEM_MODE}`\n"
        "- Soft `image=depth` is **banned** on this Space.\n"
        "- First GPU call loads Fun CN (slow once)."
    )
    with gr.Row():
        with gr.Column():
            positive = gr.Textbox(label="positive", lines=12)
            negative = gr.Textbox(label="negative", lines=3)
            depth_image = gr.Image(label="depth_image (required)", type="pil")
            seed = gr.Number(label="seed", value=42, precision=0)
            width = gr.Number(label="width", value=1216, precision=0)
            height = gr.Number(label="height", value=832, precision=0)
            steps = gr.Number(label="steps", value=28, precision=0)
            guidance = gr.Number(label="guidance", value=4.0)
            cn_strength = gr.Number(label="cn_strength", value=0.75)
            btn = gr.Button("Generate (Fun CN)", variant="primary")
        with gr.Column():
            still = gr.Image(label="still")
    btn.click(
        fn=generate_still,
        inputs=[
            positive,
            negative,
            depth_image,
            seed,
            width,
            height,
            steps,
            guidance,
            cn_strength,
        ],
        outputs=[still],
        api_name="generate_still",
    )

if __name__ == "__main__":
    demo.queue(max_size=4).launch()