import os
import sys
import tempfile
import subprocess
import math
from pathlib import Path

import gradio as gr
import spaces


# ============================================================
# Configuration
# ============================================================

REPO_DIR = Path("/tmp/EraserDiT")
MODEL_ID = "jieeliu/EraserDiT"

os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault(
    "PYTORCH_CUDA_ALLOC_CONF",
    "expandable_segments:True",
)


# ============================================================
# Clone EraserDiT
# ============================================================

def setup_repo():
    if REPO_DIR.exists():
        return

    print("Cloning EraserDiT...")

    subprocess.run(
        [
            "git",
            "clone",
            "--depth",
            "1",
            "https://github.com/JieLiu95/EraserDiT.git",
            str(REPO_DIR),
        ],
        check=True,
    )


setup_repo()

if str(REPO_DIR) not in sys.path:
    sys.path.insert(0, str(REPO_DIR))


# ============================================================
# Global model
# ============================================================

pipeline = None


# ============================================================
# Negative prompt
# ============================================================

NEGATIVE_PROMPT = (
    "Colorful color tone, overexposure, static, blurry details, "
    "subtitles, style, artwork, picture, static, overall graying, "
    "worst quality, low-quality, JPEG compression residue, ugly, "
    "incomplete, extra fingers, poorly painted hands, poorly painted "
    "faces, deformed, disfigured, deformed limbs, finger fusion, "
    "still image, cluttered background, three legs, many people in "
    "the background, walking backwards, no noise"
)


# ============================================================
# Load model
# ============================================================

def load_model():

    global pipeline

    if pipeline is not None:
        return pipeline

    import torch

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA GPU is required.")

    print()
    print("=" * 60)
    print("Loading EraserDiT")
    print("=" * 60)

    print("GPU:", torch.cuda.get_device_name(0))

    props = torch.cuda.get_device_properties(0)

    print(
        "VRAM:",
        round(
            props.total_memory / (1024 ** 3),
            2,
        ),
        "GB",
    )

    print("Torch:", torch.__version__)

    print("=" * 60)

    # Lazy imports
    from utils.common import GlobalValues
    from utils.inference_utils import init

    GlobalValues.DEBUG = False

    device = torch.device("cuda")
    weight_dtype = torch.bfloat16

    pipeline = init(
        device=device,
        weight_dtype=weight_dtype,
        pre_dir=MODEL_ID,
    )

    print("EraserDiT loaded successfully.")

    return pipeline


# ============================================================
# Get video information
# ============================================================

def get_video_info(path):

    import ffmpeg

    info = ffmpeg.probe(path)

    video_stream = next(
        (
            stream
            for stream in info["streams"]
            if stream.get("codec_type") == "video"
        ),
        None,
    )

    if video_stream is None:
        raise ValueError("No video stream found.")

    width = int(video_stream["width"])
    height = int(video_stream["height"])

    frame_count = None

    if video_stream.get("nb_frames"):
        try:
            frame_count = int(
                video_stream["nb_frames"]
            )
        except Exception:
            pass

    fps_string = video_stream.get(
        "r_frame_rate",
        "30/1",
    )

    try:
        numerator, denominator = fps_string.split("/")
        fps = float(numerator) / float(denominator)
    except Exception:
        fps = 30.0

    bitrate = video_stream.get("bit_rate")

    try:
        bitrate_mbps = max(
            1,
            int(bitrate) // 1_000_000,
        )
    except Exception:
        bitrate_mbps = 10

    return {
        "width": width,
        "height": height,
        "fps": fps,
        "frame_count": frame_count,
        "bitrate_mbps": bitrate_mbps,
    }


# ============================================================
# Main inference
# ============================================================

@spaces.GPU(duration=600)
def process_video(
    video_path,
    mask_path,
    prompt,
    ksize,
    dilate_iter,
    progress=gr.Progress(),
):

    import torch
    import decord

    # --------------------------------------------------------
    # Validate inputs
    # --------------------------------------------------------

    if video_path is None:
        raise gr.Error(
            "Please upload an input video."
        )

    if mask_path is None:
        raise gr.Error(
            "Please upload a mask video."
        )

    # Gradio can return a filepath or a dictionary depending
    # on the component/version.
    if isinstance(video_path, dict):
        video_path = video_path.get("path")

    if isinstance(mask_path, dict):
        mask_path = mask_path.get("path")

    if not video_path:
        raise gr.Error(
            "Could not obtain the input video filepath."
        )

    if not mask_path:
        raise gr.Error(
            "Could not obtain the mask video filepath."
        )

    if not prompt or not prompt.strip():
        raise gr.Error(
            "Please enter a prompt describing the scene."
        )

    video_path = str(video_path)
    mask_path = str(mask_path)

    if not os.path.isfile(video_path):
        raise gr.Error(
            f"Input video does not exist:\n{video_path}"
        )

    if not os.path.isfile(mask_path):
        raise gr.Error(
            f"Mask video does not exist:\n{mask_path}"
        )

    # --------------------------------------------------------
    # Sanitize settings
    # --------------------------------------------------------

    ksize = int(ksize)
    dilate_iter = int(dilate_iter)

    # Kernel size should be positive and odd.
    if ksize < 1:
        ksize = 1

    if ksize % 2 == 0:
        ksize += 1

    if dilate_iter < 0:
        dilate_iter = 0

    print()
    print("=" * 60)
    print("Mask settings")
    print("=" * 60)

    print(
        "Kernel size:",
        f"{ksize}x{ksize}",
    )

    print(
        "Dilation iterations:",
        dilate_iter,
    )

    print("=" * 60)

    # --------------------------------------------------------
    # Video information
    # --------------------------------------------------------

    try:
        video_info = get_video_info(video_path)
        mask_info = get_video_info(mask_path)

    except Exception as e:

        raise gr.Error(
            f"Could not inspect videos:\n{e}"
        )

    width = video_info["width"]
    height = video_info["height"]

    mask_width = mask_info["width"]
    mask_height = mask_info["height"]

    print()
    print("=" * 60)
    print("Input")
    print("=" * 60)

    print(
        f"Video: {width}x{height} "
        f"@ {video_info['fps']:.3f} FPS"
    )

    print(
        f"Mask : {mask_width}x{mask_height} "
        f"@ {mask_info['fps']:.3f} FPS"
    )

    # --------------------------------------------------------
    # Resolution check
    # --------------------------------------------------------

    if width * height > 1920 * 1088:

        raise gr.Error(
            "Videos larger than 1920x1088 are not supported "
            "by this Space version."
        )

    if (
        mask_width != width
        or mask_height != height
    ):

        raise gr.Error(
            "Video and mask must have exactly the same "
            "resolution.\n\n"
            f"Video: {width}x{height}\n"
            f"Mask:  {mask_width}x{mask_height}"
        )

    # --------------------------------------------------------
    # Read frame counts
    # --------------------------------------------------------

    try:

        video_reader = decord.VideoReader(
            video_path,
            ctx=decord.cpu(0),
        )

        mask_reader = decord.VideoReader(
            mask_path,
            ctx=decord.cpu(0),
        )

        video_frames = len(video_reader)
        mask_frames = len(mask_reader)

        del video_reader
        del mask_reader

    except Exception as e:

        raise gr.Error(
            f"Could not read video:\n{e}"
        )

    if video_frames != mask_frames:

        raise gr.Error(
            "Video and mask must contain exactly the "
            "same number of frames.\n\n"
            f"Video frames: {video_frames}\n"
            f"Mask frames:  {mask_frames}"
        )

    print(
        "Frames:",
        video_frames,
    )

    # --------------------------------------------------------
    # FPS check
    # --------------------------------------------------------

    fps_difference = abs(
        video_info["fps"]
        - mask_info["fps"]
    )

    if fps_difference > 0.01:

        raise gr.Error(
            "Video and mask must have the same FPS.\n\n"
            f"Video FPS: {video_info['fps']}\n"
            f"Mask FPS:  {mask_info['fps']}"
        )

    # --------------------------------------------------------
    # Load model
    # --------------------------------------------------------

    pipe = load_model()

    # --------------------------------------------------------
    # Import EraserDiT components
    # --------------------------------------------------------

    from utils.pre import VideoInpaintPre
    from utils.inference_utils import inference_batch
    from utils.post import post_stream_normalized
    from utils.post_pkg import FFmpegWriter

    # --------------------------------------------------------
    # EraserDiT settings
    # --------------------------------------------------------

    device = torch.device("cuda")
    weight_dtype = torch.bfloat16

    preprocessor = VideoInpaintPre(
        device=device,

        align_h=32,
        align_w=32,

        # User-controlled mask processing
        ksize=(ksize, ksize),
        dilate_iter=dilate_iter,

        # Temporal overlap
        shift_alpha=1 * 8 + 1,

        # EraserDiT temporal chunk
        TEMP_INFER_LEN=121,

        crop_flag=False,
    )

    generator = None

    # --------------------------------------------------------
    # Output path
    # --------------------------------------------------------

    output_dir = Path(
        tempfile.mkdtemp(
            prefix="eraserdit_"
        )
    )

    input_name = Path(
        video_path
    ).stem

    output_path = (
        output_dir
        / f"{input_name}_eraserdit.mp4"
    )

    print()
    print("=" * 60)
    print("Output")
    print("=" * 60)

    print(output_path)

    # --------------------------------------------------------
    # Processing state
    # --------------------------------------------------------

    pre_video_shift = None

    video_save_writer = None

    current_batch = 0

    # --------------------------------------------------------
    # Process temporal batches
    # --------------------------------------------------------

    try:

        while True:

            (
                video_ori,
                mask_ori,
                fps,
                videos_input_ori,
                masks_input_ori,
            ) = preprocessor.load_videos(
                video_path=video_path,
                mask_path=mask_path,
                bbox_path=None,
                decord_device=decord.cpu(0),
                sample_rate=1,
                batch_idx=current_batch,
            )

            # ------------------------------------------------
            # End of video
            # ------------------------------------------------

            if video_ori is None:
                break

            # ------------------------------------------------
            # Prepare batch
            # ------------------------------------------------

            (
                video_input,
                mask_input,
                _,
            ) = preprocessor(
                video_ori,
                mask_ori,
                batch_idx=current_batch,
                format="nhwc",
            )

            input_shape = (
                preprocessor.TranslateShape(
                    video_input.shape,
                    src="nchw",
                    dst="nhwc",
                )
            )

            # ------------------------------------------------
            # Create output writer
            # ------------------------------------------------

            if video_save_writer is None:

                video_save_writer = FFmpegWriter(
                    path=str(output_path),
                    width=video_ori.shape[2],
                    height=video_ori.shape[1],
                    fps=fps,
                    bitrate=(
                        f"{video_info['bitrate_mbps']}M"
                    ),
                )

            # ------------------------------------------------
            # Temporal overlap
            # ------------------------------------------------

            if current_batch == 0:

                masks_zero_shift = torch.zeros(
                    (
                        math.ceil(
                            preprocessor.shift_alpha
                            / 8
                        ),
                        mask_input.shape[1],
                        mask_input.shape[2],
                        mask_input.shape[3],
                    ),
                    dtype=mask_input.dtype,
                )

            else:

                video_input = torch.cat(
                    [
                        pre_video_shift,
                        video_input,
                    ],
                    dim=0,
                )

                mask_input = torch.cat(
                    [
                        masks_zero_shift,
                        mask_input,
                    ],
                    dim=0,
                )

            print()
            print(
                f"Batch {current_batch}"
            )

            print(
                "Source frames:",
                video_ori.shape[0],
            )

            print(
                "Model input:",
                tuple(video_input.shape),
            )

            # ------------------------------------------------
            # Run EraserDiT
            # ------------------------------------------------

            output_frames = inference_batch(
                videos=video_input,
                masks_input=mask_input,
                prompt=prompt.strip(),
                negative_prompt=NEGATIVE_PROMPT,
                pipeline=pipe,
                generator=generator,
                device=device,
                weight_dtype=weight_dtype,
            )

            # ------------------------------------------------
            # Save overlap for next batch
            # ------------------------------------------------

            pre_video_shift = (
                output_frames[
                    -preprocessor.shift_alpha:
                ].cpu()
            )

            # ------------------------------------------------
            # Write output
            # ------------------------------------------------

            if current_batch == 0:

                post_stream_normalized(
                    output_frames=output_frames,
                    ori_shape=video_ori.shape,
                    model_video_shape=input_shape,
                    crop_flag=False,
                    videos_input_ori=None,
                    video_ori=video_ori,
                    mask_ori=mask_ori,
                    output_bbox=None,
                    writer=video_save_writer,
                    write_to=True,
                )

            else:

                post_stream_normalized(
                    output_frames=(
                        output_frames[
                            preprocessor.shift_alpha:
                        ]
                    ),
                    ori_shape=video_ori.shape,
                    model_video_shape=input_shape,
                    crop_flag=False,
                    videos_input_ori=None,
                    video_ori=video_ori,
                    mask_ori=mask_ori,
                    output_bbox=None,
                    writer=video_save_writer,
                    write_to=True,
                )

            # ------------------------------------------------
            # Progress
            # ------------------------------------------------

            current_batch += 1

            processed_frames = min(
                121
                + (
                    max(
                        0,
                        current_batch - 1,
                    )
                    * (
                        121
                        - preprocessor.shift_alpha
                    )
                ),
                video_frames,
            )

            fraction = (
                processed_frames
                / video_frames
            )

            progress(
                fraction,
                desc=(
                    f"Processing "
                    f"{processed_frames}/"
                    f"{video_frames} frames"
                ),
            )

            print(
                f"Progress: "
                f"{processed_frames}/"
                f"{video_frames}"
            )

            # ------------------------------------------------
            # Release temporary tensors
            # ------------------------------------------------

            del output_frames
            del video_input
            del mask_input
            del video_ori
            del mask_ori

        # ----------------------------------------------------
        # Close writer
        # ----------------------------------------------------

        if video_save_writer is not None:

            video_save_writer.Close()
            video_save_writer = None

    except Exception:

        if video_save_writer is not None:

            try:
                video_save_writer.Close()
            except Exception:
                pass

            video_save_writer = None

        raise

    # --------------------------------------------------------
    # Validate output
    # --------------------------------------------------------

    if not output_path.exists():

        raise gr.Error(
            "EraserDiT finished without producing "
            "an output video."
        )

    output_size = output_path.stat().st_size

    if output_size <= 0:

        raise gr.Error(
            "The generated output video is empty."
        )

    print()
    print("=" * 60)
    print("Finished")
    print("=" * 60)

    print(
        "Output:",
        output_path,
    )

    print(
        "Size:",
        f"{output_size / (1024 * 1024):.2f} MB",
    )

    return str(output_path)


# ============================================================
# Gradio UI
# ============================================================

with gr.Blocks(
    title="EraserDiT Video Object Removal"
) as demo:

    gr.Markdown(
        """
# EraserDiT — Video Object Removal

Upload an original video and its corresponding mask video.

The masked region will be removed using EraserDiT.

### Requirements

- Video and mask must have the same resolution
- Video and mask must have the same number of frames
- Video and mask must have the same FPS
- Maximum resolution: **1920×1088**
"""
    )

    with gr.Row():

        # ====================================================
        # Input column
        # ====================================================

        with gr.Column():

            input_video = gr.Video(
                label="Input Video",
                sources=["upload"],
            )

            input_mask = gr.Video(
                label="Mask Video",
                sources=["upload"],
            )

            prompt = gr.Textbox(
                label="Scene Description",
                placeholder=(
                    "Describe what the scene should look "
                    "like after removing the masked object."
                ),
                value=(
                    "A natural continuation of the "
                    "surrounding video scene."
                ),
                lines=4,
            )

            gr.Markdown(
                "### Mask Processing"
            )

            ksize_slider = gr.Slider(
                minimum=1,
                maximum=31,
                value=9,
                step=2,
                label="Kernel Size",
                info=(
                    "Morphological kernel size. "
                    "Must be odd."
                ),
            )

            dilate_slider = gr.Slider(
                minimum=0,
                maximum=30,
                value=9,
                step=1,
                label="Dilation Iterations",
                info=(
                    "0 = no dilation. Higher values "
                    "expand the mask."
                ),
            )

            run_button = gr.Button(
                "Run EraserDiT",
                variant="primary",
                size="lg",
            )

        # ====================================================
        # Output column
        # ====================================================

        with gr.Column():

            output_video = gr.Video(
                label="Output Video",
                interactive=False,
            )

    # ========================================================
    # Run button
    # ========================================================

    run_button.click(
        fn=process_video,
        inputs=[
            input_video,
            input_mask,
            prompt,
            ksize_slider,
            dilate_slider,
        ],
        outputs=[
            output_video,
        ],
    )

    # ========================================================
    # Help
    # ========================================================

    gr.Markdown(
        """
### Mask settings

**Kernel Size**

Controls the size of the morphological kernel used when
processing the mask.

- `1` = effectively no kernel expansion
- `3–7` = small processing area
- `9` = EraserDiT's previous default
- `15+` = increasingly aggressive

**Dilation Iterations**

Controls how many times the mask is dilated.

- `0` = no explicit dilation
- `1–5` = small expansion
- `9` = EraserDiT's previous default
- `10–30` = increasingly large expansion

If you want the smallest possible processing area around your
mask, start with:

**Kernel Size = 1**  
**Dilation Iterations = 0**

### Example prompt

For removing a person from a street:

`A city street with buildings, cars and pedestrians in the background.`
"""
    )


# ============================================================
# Launch
# ============================================================

if __name__ == "__main__":

    demo.queue(
        max_size=1,
        default_concurrency_limit=1,
    ).launch(
        show_error=True,
    )