import json
import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # must be imported before torch / CUDA-touching libs

import random
import re
import tempfile
import time
from types import SimpleNamespace

import einops
import gradio as gr
import imageio
import numpy as np
import torch
from huggingface_hub import snapshot_download
from PIL import Image
from torchvision import transforms

import hyvideo.commons
import hyvideo.pipelines.hunyuan_video_pipeline as hv_pipeline_module
from hyvideo.commons import auto_offload_model
from hyvideo.commons.infer_state import initialize_infer_state
from hyvideo.pipelines.hunyuan_video_pipeline import HunyuanVideo_1_5_Pipeline
from univedit.data.preprocessing import IMAGE_CLIP_FRAMES, resize_and_center_crop_image

# ---------------------------------------------------------------------------
# Weights (downloaded at startup)
# ---------------------------------------------------------------------------
CKPT_DIR = os.path.abspath("ckpts/omniweaving")
BASE_REPO = "tencent/HY-OmniWeaving"
UNIVEDIT_REPO = "SanRiiiii/UniVEdit"
SIGLIP_VISION_CONFIG = {
    "architectures": ["SiglipVisionModel"],
    "attention_dropout": 0.0,
    "hidden_act": "gelu_pytorch_tanh",
    "hidden_size": 1152,
    "image_size": 384,
    "intermediate_size": 4304,
    "layer_norm_eps": 1e-06,
    "model_type": "siglip_vision_model",
    "num_attention_heads": 16,
    "num_channels": 3,
    "num_hidden_layers": 27,
    "patch_size": 14,
}


def _link(src, dst):
    os.makedirs(os.path.dirname(dst), exist_ok=True)
    if not os.path.lexists(dst):
        os.symlink(src, dst)


def prepare_weights() -> str:
    # OmniWeaving base: transformer, VAE, scheduler and the fine-tuned MLLM text encoder.
    base = snapshot_download(
        BASE_REPO,
        allow_patterns=[
            "config.json",
            "scheduler/*",
            "transformer/*",
            "vae/*",
            "text_encoder/ckpt/*",
        ],
    )
    _link(os.path.join(base, "config.json"), os.path.join(CKPT_DIR, "config.json"))
    for sub in ("scheduler", "transformer", "vae"):
        _link(os.path.join(base, sub), os.path.join(CKPT_DIR, sub))
    te_ckpt = os.path.join(base, "text_encoder", "ckpt", "text_encoder_model.safetensors")
    _link(te_ckpt, os.path.join(CKPT_DIR, "text_encoder", "ckpt", "text_encoder_model.safetensors"))

    # Qwen2.5-VL-7B-Instruct: config / tokenizer / processor only. The OmniWeaving
    # fine-tuned MLLM checkpoint holds the full weights in HF key format, so it is
    # used directly as the model weights (no need to fetch the 16 GB base model).
    qwen = snapshot_download(
        "Qwen/Qwen2.5-VL-7B-Instruct",
        allow_patterns=["*.json", "*.txt"],
        ignore_patterns=["model.safetensors.index.json"],
    )
    llm_dir = os.path.join(CKPT_DIR, "text_encoder", "llm")
    os.makedirs(llm_dir, exist_ok=True)
    for f in os.listdir(qwen):
        if f != "model.safetensors.index.json":
            _link(os.path.join(qwen, f), os.path.join(llm_dir, f))
    _link(te_ckpt, os.path.join(llm_dir, "model.safetensors"))

    # Glyph-SDXL-v2 byT5 encoder.
    glyph = snapshot_download("multimodalart/glyph-sdxl-v2-byt5-small", allow_patterns=["assets/*", "pytorch_model.bin"])
    glyph_root = os.path.join(CKPT_DIR, "text_encoder", "Glyph-SDXL-v2")
    for f in ("color_idx.json", "multilingual_10-lang_idx.json"):
        _link(os.path.join(glyph, "assets", f), os.path.join(glyph_root, "assets", f))
    _link(os.path.join(glyph, "pytorch_model.bin"), os.path.join(glyph_root, "checkpoints", "byt5_model.pt"))

    # SigLIP vision encoder. OmniWeaving points at FLUX.1-Redux-dev/image_encoder, which is
    # the (gated) bf16 copy of google/siglip-so400m-patch14-384's vision tower; we use the
    # ungated original with the same vision config + preprocessor.
    siglip = snapshot_download(
        "google/siglip-so400m-patch14-384", allow_patterns=["model.safetensors", "preprocessor_config.json"]
    )
    siglip_dir = os.path.join(CKPT_DIR, "vision_encoder", "siglip")
    enc_dir = os.path.join(siglip_dir, "image_encoder")
    os.makedirs(enc_dir, exist_ok=True)
    with open(os.path.join(enc_dir, "config.json"), "w") as f:
        json.dump(SIGLIP_VISION_CONFIG, f)
    _link(os.path.join(siglip, "model.safetensors"), os.path.join(enc_dir, "model.safetensors"))
    _link(
        os.path.join(siglip, "preprocessor_config.json"),
        os.path.join(siglip_dir, "feature_extractor", "preprocessor_config.json"),
    )

    # UniVEdit: LoRA + modality token embeddings.
    univedit = snapshot_download(UNIVEDIT_REPO, allow_patterns=["geo3d/*"])
    return os.path.join(univedit, "geo3d")


UNIVEDIT_DIR = prepare_weights()

# ---------------------------------------------------------------------------
# Pipeline
# ---------------------------------------------------------------------------
# ZeroGPU has no GPU in the main process: report a large card so the pipeline
# picks the non-tiled-small VAE config, and never enables offloading.
_fake_mem = lambda device=None: 96 * 1024**3  # noqa: E731
hyvideo.commons.get_gpu_memory = _fake_mem
hv_pipeline_module.get_gpu_memory = _fake_mem

# Single-process inference: no sequence parallelism. The repo's ParallelDims builds a
# torch DeviceMesh (real CUDA init), which ZeroGPU's main process cannot do, so we
# register an equivalent no-parallelism state directly.
import hyvideo.commons.parallel_states as _ps

setattr(
    _ps,
    "__parallel_dims",
    SimpleNamespace(sp=1, world_size=1, dp_replicate=1, sp_rank=0, sp_group=None, sp_enabled=False, dp_enabled=False),
)
initialize_infer_state(
    SimpleNamespace(
        sage_blocks_range="0-53",
        no_cache_block_id="53",
        use_sageattn=False,
        enable_torch_compile=False,
        enable_cache=False,
        cache_type="deepcache",
        cache_start_step=11,
        cache_end_step=45,
        total_steps=50,
        cache_step_interval=4,
        use_fp8_gemm=False,
        quant_type="fp8-per-token-sgl",
        include_patterns="double_blocks",
    )
)

# The MLLM weights are already loaded from the fine-tuned checkpoint (see above).
HunyuanVideo_1_5_Pipeline._load_text_encoder_ckpt = staticmethod(lambda text_encoder, ckpt_path: None)

pipe = HunyuanVideo_1_5_Pipeline.create_pipeline(
    pretrained_model_name_or_path=CKPT_DIR,
    transformer_dtype=torch.bfloat16,
    device=torch.device("cpu"),
    transformer_init_device=torch.device("cpu"),
    pipeline_config="omniweaving",
)
pipe.apply_infer_optimization(infer_state=None, enable_offloading=False, enable_group_offloading=False)
pipe.load_modality_token_embeddings(UNIVEDIT_DIR, strict=True)
pipe.transformer.load_lora_adapter(
    pretrained_model_name_or_path_or_dict=UNIVEDIT_DIR,
    prefix=None,
    adapter_name="default",
    weight_name="pytorch_lora_weights.safetensors",
    use_safetensors=True,
)
# Merge the LoRA into the base weights (same math, no per-layer PEFT overhead at inference).
pipe.transformer.fuse_lora(lora_scale=1.0)
pipe.transformer.unload_lora()
pipe.to("cuda")
pipe.execution_device = torch.device("cuda")
pipe.noise_init_device = torch.device("cuda")
print("UniVEdit pipeline ready.")

# ---------------------------------------------------------------------------
# Helpers (ported 1:1 from univedit/inference/batch_infer.py)
# ---------------------------------------------------------------------------
IMAGE_EDIT_BUCKETS = ((352, 624), (432, 576), (480, 480), (576, 432), (624, 352))  # (width, height)
VIDEO_EDIT_360P_BUCKETS = {"1:1": (480, 480), "16:9": (352, 624), "9:16": (624, 352)}  # (height, width)
MAX_SEED = 2**31 - 1
DEFAULT_STEPS = 50
DEFAULT_GUIDANCE = 6.0
DEFAULT_SHIFT = 5.0


def replace_reference_placeholder(text: str) -> str:
    text = re.sub(r"\b(?:a|an|the)\s*\[\s*v\s*\*\s*\]", "the reference image", text, flags=re.IGNORECASE)
    text = re.sub(r"\[\s*v\s*\*\s*\]", "the reference image", text, flags=re.IGNORECASE)
    text = re.sub(r"\s+([,.;:!?])", r"\1", text)
    text = re.sub(r"\s{2,}", " ", text)
    return text.strip()


def image_bucket_size(img: Image.Image):
    ratio = img.width / img.height
    w, h = min(IMAGE_EDIT_BUCKETS, key=lambda s: abs((s[0] / s[1]) - ratio))
    return h, w


def video_bucket_size(width: int, height: int):
    ratio = width / height
    _, (h, w) = min(VIDEO_EDIT_360P_BUCKETS.items(), key=lambda it: abs((it[1][1] / it[1][0]) - ratio))
    return h, w


def encode_frames_to_latents(frames: torch.Tensor, *, scale: bool) -> torch.Tensor:
    pixel_values = transforms.Normalize(mean=[0.5] * 3, std=[0.5] * 3)(frames)
    pixel_values = pixel_values.unsqueeze(0).transpose(1, 2).to(pipe.execution_device)
    with (
        torch.no_grad(),
        torch.autocast(device_type="cuda", dtype=torch.float16, enabled=True),
        auto_offload_model(pipe.vae, pipe.execution_device, enabled=False),
        pipe.vae.memory_efficient_context(),
    ):
        latents = pipe.vae.encode(pixel_values).latent_dist.mode()
        if scale:
            if getattr(pipe.vae.config, "shift_factor", None):
                latents = (latents - pipe.vae.config.shift_factor) * pipe.vae.config.scaling_factor
            else:
                latents = latents * pipe.vae.config.scaling_factor
    return latents


def to_pil_first_frame(video: torch.Tensor) -> Image.Image:
    frame = video[0, :, 0] if video.ndim == 5 else video[0]
    frame = (frame * 255).clamp(0, 255).to(torch.uint8)
    return Image.fromarray(einops.rearrange(frame, "c h w -> h w c").cpu().numpy())


def resolve_seed(seed, randomize):
    return random.randint(0, MAX_SEED) if randomize else int(seed)


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------
@spaces.GPU(duration=70)  # measured ~60s for 50 steps (+~15%)
def edit_image(
    source_image: Image.Image,
    instruction: str,
    reference_image: Image.Image | None = None,
    steps: int = DEFAULT_STEPS,
    guidance_scale: float = DEFAULT_GUIDANCE,
    flow_shift: float = DEFAULT_SHIFT,
    seed: int = 123,
    randomize_seed: bool = False,
    progress=gr.Progress(track_tqdm=True),
):
    """Edit an image with a text instruction (optionally guided by a reference image).

    Args:
        source_image: the image to edit.
        instruction: the editing instruction, e.g. "replace the scarf with a hat". With a reference
            image, refer to it as "[V*]" (e.g. "change the dog to a [V*]").
        reference_image: optional reference image (switches to reference-guided editing, ti_i2i).
        steps: number of denoising steps.
        guidance_scale: classifier-free guidance scale.
        flow_shift: timestep shift of the sampling schedule.
        seed: random seed.
        randomize_seed: pick a random seed instead of `seed`.

    Returns:
        The edited image and the seed that was used.
    """
    t0 = time.time()
    if source_image is None:
        raise gr.Error("Please upload an image to edit.")
    if not instruction or not instruction.strip():
        raise gr.Error("Please enter an editing instruction.")
    seed = resolve_seed(seed, randomize_seed)
    source_image = source_image.convert("RGB")
    task = "ti_i2i" if reference_image is not None else "i2i"
    prompt = replace_reference_placeholder(instruction) if task == "ti_i2i" else instruction.strip()

    h, w = image_bucket_size(source_image)
    resized = resize_and_center_crop_image(source_image, h, w)
    clip = transforms.ToTensor()(resized).unsqueeze(0).repeat(IMAGE_CLIP_FRAMES, 1, 1, 1)
    latents = encode_frames_to_latents(clip, scale=True)

    kwargs = dict(
        prompt=prompt,
        negative_prompt="",
        aspect_ratio=f"{w}:{h}",
        num_inference_steps=int(steps),
        guidance_scale=float(guidance_scale),
        flow_shift=float(flow_shift),
        video_length=IMAGE_CLIP_FRAMES,
        seed=seed,
        output_type="pt",
        return_dict=True,
        deepstack=[8, 16, 24],
        setclip=True,
        enable_vae_tile_parallelism=False,  # single process, no torch.distributed
        task_type=task,
        source_image=source_image,
        source_image_latents=latents,
        target_height=h,
        target_width=w,
    )
    if task == "ti_i2i":
        kwargs["reference_image"] = reference_image.convert("RGB")
    t1 = time.time()
    with torch.no_grad():
        out = pipe(**kwargs)
    print(f"[timing] edit_image: prep {t1 - t0:.1f}s, pipe {time.time() - t1:.1f}s")
    return to_pil_first_frame(out.videos), seed


def _video_duration(source_video, instruction, reference_image=None, num_frames=33, steps=DEFAULT_STEPS, *args, **kwargs):
    # Calibrated on the live Space: ~1.15 s/step for a 5-frame image clip and ~2.9 s/step
    # for 33 frames at 16:9 (50 steps -> ~148 s); ~10% margin on top.
    lat_frames = (max(5, min(65, int(num_frames))) - 1) // 4 + 1
    tokens = lat_frames * 900
    per_step = 0.55 + 0.000306 * tokens + 1.5e-9 * tokens**2
    return int(5 + 1.1 * int(steps) * per_step)


@spaces.GPU(duration=_video_duration)
def edit_video(
    source_video: str,
    instruction: str,
    reference_image: Image.Image | None = None,
    num_frames: int = 33,
    steps: int = DEFAULT_STEPS,
    guidance_scale: float = DEFAULT_GUIDANCE,
    flow_shift: float = DEFAULT_SHIFT,
    seed: int = 123,
    randomize_seed: bool = False,
    progress=gr.Progress(track_tqdm=True),
):
    """Edit a short video with a text instruction (optionally guided by a reference image).

    Args:
        source_video: path of the video to edit (the first `num_frames` frames are edited).
        instruction: the editing instruction, e.g. "Replace the brown bear with a panda.".
        reference_image: optional reference image (switches to reference-guided video editing, tiv2v).
        num_frames: number of frames to edit (4n+1, at most 65).
        steps: number of denoising steps.
        guidance_scale: classifier-free guidance scale.
        flow_shift: timestep shift of the sampling schedule.
        seed: random seed.
        randomize_seed: pick a random seed instead of `seed`.

    Returns:
        Path of the edited mp4 video and the seed that was used.
    """
    from decord import VideoReader

    t0 = time.time()
    if not source_video:
        raise gr.Error("Please upload a video to edit.")
    if not instruction or not instruction.strip():
        raise gr.Error("Please enter an editing instruction.")
    seed = resolve_seed(seed, randomize_seed)
    num_frames = int(num_frames)
    num_frames = max(5, min(65, ((num_frames - 1) // 4) * 4 + 1))
    task = "tiv2v" if reference_image is not None else "v2v"

    vr = VideoReader(source_video)
    total = len(vr)
    if total == 0:
        raise gr.Error("The video has no frames.")
    fps = round(vr.get_avg_fps()) or 12
    indices = list(range(min(total, num_frames))) + [total - 1] * max(0, num_frames - total)
    frames = torch.from_numpy(vr.get_batch(indices).asnumpy()).permute(0, 3, 1, 2).contiguous() / 255.0
    del vr

    height, width = frames.shape[-2:]
    th, tw = video_bucket_size(width, height)
    if th / height > tw / width:
        resize_size = th, int(width * th / height)
    else:
        resize_size = int(height * tw / width), tw
    frames = transforms.Compose(
        [transforms.Resize(resize_size, antialias=True), transforms.CenterCrop((th, tw))]
    )(frames)
    latents = encode_frames_to_latents(frames, scale=False)

    kwargs = dict(
        prompt=instruction.strip(),
        negative_prompt="",
        aspect_ratio=f"{tw}:{th}",
        num_inference_steps=int(steps),
        guidance_scale=float(guidance_scale),
        flow_shift=float(flow_shift),
        video_length=num_frames,
        seed=seed,
        output_type="pt",
        return_dict=True,
        deepstack=[8, 16, 24],
        setclip=True,
        enable_vae_tile_parallelism=False,  # single process, no torch.distributed
        task_type=task,
        source_video=[source_video],
        source_video_latents=latents,
    )
    if task == "tiv2v":
        ref_path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
        reference_image.convert("RGB").save(ref_path)
        kwargs["reference_image"] = ref_path
    t1 = time.time()
    with torch.no_grad():
        out = pipe(**kwargs)
    print(f"[timing] edit_video: prep {t1 - t0:.1f}s, pipe {time.time() - t1:.1f}s ({num_frames} frames)")

    video = out.videos
    if video.ndim == 5:
        video = video[0]
    vid = (video * 255).clamp(0, 255).to(torch.uint8)
    vid = einops.rearrange(vid, "c f h w -> f h w c").cpu().numpy()
    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    imageio.mimwrite(out_path, list(vid), fps=fps, quality=8)
    return out_path, seed


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# VitronEdit (UniVEdit): universal visual editing\n"
            "Edit images and videos with a text instruction, and optionally a reference image. "
            "This is a LoRA plus modality tokens on top of "
            "[HY-OmniWeaving](https://huggingface.co/tencent/HY-OmniWeaving) (HunyuanVideo-1.5).\n\n"
            "[Paper](https://huggingface.co/papers/2610.27049) · "
            "[Code](https://github.com/Vitron-LLM/VitronEdit-Code) · "
            "[Model](https://huggingface.co/SanRiiiii/UniVEdit)"
        )
        with gr.Tabs():
            # ---------------- Image editing ----------------
            with gr.Tab("Image editing"):
                with gr.Row():
                    with gr.Column():
                        img_src = gr.Image(label="Image to edit", type="pil", height=380)
                        img_prompt = gr.Textbox(
                            label="Instruction",
                            placeholder='e.g. "replace the scarf with a hat". With a reference image, call it [V*].',
                            lines=2,
                        )
                        with gr.Accordion("Reference image (optional)", open=False):
                            img_ref = gr.Image(label="Reference image", type="pil", height=240)
                        img_btn = gr.Button("Edit image", variant="primary")
                        with gr.Accordion("Advanced settings", open=False):
                            img_steps = gr.Slider(10, 50, value=DEFAULT_STEPS, step=1, label="Steps")
                            img_cfg = gr.Slider(1.0, 10.0, value=DEFAULT_GUIDANCE, step=0.5, label="Guidance scale")
                            img_shift = gr.Slider(1.0, 10.0, value=DEFAULT_SHIFT, step=0.5, label="Flow shift")
                            img_seed = gr.Slider(0, MAX_SEED, value=123, step=1, label="Seed")
                            img_rand = gr.Checkbox(label="Randomize seed", value=False)
                    with gr.Column():
                        img_out = gr.Image(label="Edited image", height=480)
                img_btn.click(
                    edit_image,
                    inputs=[img_src, img_prompt, img_ref, img_steps, img_cfg, img_shift, img_seed, img_rand],
                    outputs=[img_out, img_seed],
                    api_name="edit_image",
                )
                gr.Examples(
                    examples=[
                        ["examples/cat_tabby.jpg", "add a small red bow tie under the cat's chin"],
                        ["examples/living_room_blue_couch.jpg", "replace the blue sofa with a brown leather sofa"],
                        ["examples/rubber_ducks.jpg", "remove the smallest rubber duck on the left"],
                        ["examples/lighthouse_coast.jpg", "Make the style of the image a watercolor painting"],
                    ],
                    inputs=[img_src, img_prompt],
                    outputs=[img_out, img_seed],
                    fn=edit_image,
                    cache_examples=True,
                    cache_mode="lazy",
                    label="Image examples",
                )

            # ---------------- Video editing ----------------
            with gr.Tab("Video editing"):
                with gr.Row():
                    with gr.Column():
                        vid_src = gr.Video(label="Video to edit", height=340)
                        vid_prompt = gr.Textbox(
                            label="Instruction",
                            placeholder='e.g. "Replace the brown bear with a panda."',
                            lines=2,
                        )
                        with gr.Accordion("Reference image (optional)", open=False):
                            vid_ref = gr.Image(label="Reference image", type="pil", height=240)
                        vid_btn = gr.Button("Edit video", variant="primary")
                        with gr.Accordion("Advanced settings", open=False):
                            vid_frames = gr.Slider(5, 65, value=33, step=4, label="Frames (4n+1)")
                            vid_steps = gr.Slider(10, 50, value=DEFAULT_STEPS, step=1, label="Steps")
                            vid_cfg = gr.Slider(1.0, 10.0, value=DEFAULT_GUIDANCE, step=0.5, label="Guidance scale")
                            vid_shift = gr.Slider(1.0, 10.0, value=DEFAULT_SHIFT, step=0.5, label="Flow shift")
                            vid_seed = gr.Slider(0, MAX_SEED, value=123, step=1, label="Seed")
                            vid_rand = gr.Checkbox(label="Randomize seed", value=False)
                    with gr.Column():
                        vid_out = gr.Video(label="Edited video", autoplay=True, height=420)
                vid_btn.click(
                    edit_video,
                    inputs=[vid_src, vid_prompt, vid_ref, vid_frames, vid_steps, vid_cfg, vid_shift, vid_seed, vid_rand],
                    outputs=[vid_out, vid_seed],
                    api_name="edit_video",
                )
                gr.Examples(
                    examples=[
                        ["examples/cat.mp4", "Replace the cat with a tiger."],
                        ["examples/border_collie_running.mp4", "Replace the dog with a red fox."],
                    ],
                    inputs=[vid_src, vid_prompt],
                    outputs=[vid_out, vid_seed],
                    fn=edit_video,
                    cache_examples=True,
                    cache_mode="lazy",
                    label="Video examples",
                )

if __name__ == "__main__":
    demo.queue().launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)