import os
import torch
import gradio as gr
import spaces
import diffusers

from diffusers import ZImagePipeline


# ============================================================
# VERSION CHECK
# ============================================================

print("=" * 60)
print("VERSION CHECK")
print("=" * 60)
print("Torch:", torch.__version__)
print("Diffusers:", diffusers.__version__)
print(
    "Transformers:",
    __import__("transformers").__version__,
)
print(
    "HuggingFace Hub:",
    __import__("huggingface_hub").__version__,
)
print(
    "Gradio:",
    gr.__version__,
)
print(
    "CUDA available:",
    torch.cuda.is_available(),
)

print(
    "Has ZImagePipeline:",
    hasattr(diffusers, "ZImagePipeline"),
)

print("=" * 60)


# ============================================================
# CHECK Z-IMAGE PIPELINE
# ============================================================

if not hasattr(diffusers, "ZImagePipeline"):
    raise RuntimeError(
        "ZImagePipeline is not available. "
        "Please install diffusers==0.36.0 or newer."
    )


# ============================================================
# LOAD Z-IMAGE
# ============================================================

print("Loading Z-Image-Turbo pipeline...")

pipe = ZImagePipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=False,
)

pipe.to("cuda")

print("Z-Image pipeline loaded!")


# ============================================================
# WATERMARK MODULE
# ============================================================

from watermark import (
    preview,
    remove_image,
    process_video,
)

from watermark.video import get_video_info


# ============================================================
# Z-IMAGE GENERATION
# ============================================================

@spaces.GPU(duration=60)
def generate_image(
    prompt: str,
    height: int,
    width: int,
    steps: int,
    seed: int,
    randomize_seed: bool,
    progress=gr.Progress(track_tqdm=True),
):
    """
    Generate an image with Z-Image-Turbo.
    """

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

    if randomize_seed:
        seed = torch.randint(
            0,
            2**32 - 1,
            (1,),
        ).item()

    seed = int(seed)

    generator = torch.Generator(
        device="cuda"
    ).manual_seed(seed)

    print(
        f"Generating image: "
        f"{width}x{height}, "
        f"steps={steps}, "
        f"seed={seed}"
    )

    image = pipe(
        prompt=prompt,
        height=int(height),
        width=int(width),
        num_inference_steps=int(steps),
        guidance_scale=0.0,
        generator=generator,
    ).images[0]

    return image, seed


# ============================================================
# WATERMARK PREVIEW
# ============================================================

@spaces.GPU(duration=30)
def wm_preview(
    image,
    detection_prompt: str = "watermark",
    max_bbox_percent: float = 30,
):
    """
    Detect likely watermark regions.
    """

    if image is None:
        raise gr.Error("Please upload an image.")

    return preview(
        image,
        float(max_bbox_percent),
        detection_prompt or "watermark",
    )


# ============================================================
# WATERMARK IMAGE REMOVAL
# ============================================================

@spaces.GPU(duration=45)
def wm_remove_image(
    image,
    detection_prompt: str = "watermark",
    max_bbox_percent: float = 30,
    transparent: bool = False,
):
    """
    Detect and remove watermark regions from an image.
    """

    if image is None:
        raise gr.Error("Please upload an image.")

    return remove_image(
        image,
        float(max_bbox_percent),
        detection_prompt or "watermark",
        bool(transparent),
    )


# ============================================================
# WATERMARK VIDEO REMOVAL
# ============================================================

@spaces.GPU(duration=60)
def wm_remove_video(
    video_path: str,
    detection_prompt: str = "watermark",
    max_bbox_percent: float = 30,
    detection_skip: int = 5,
    fade_in: int = 2,
    fade_out: int = 2,
    progress=gr.Progress(),
):
    """
    Detect and remove watermarks from a video.
    """

    if not video_path:
        raise gr.Error(
            "Please upload a video."
        )

    # --------------------------------------------------------
    # Inspect video
    # --------------------------------------------------------

    try:

        duration, _, _, _, _ = get_video_info(
            video_path
        )

        max_seconds = float(
            os.getenv(
                "WATERMARK_MAX_VIDEO_SECONDS",
                "30",
            )
        )

        if duration > max_seconds + 0.25:

            raise gr.Error(
                f"Video is {duration:.1f}s. "
                f"Maximum is "
                f"{max_seconds:.0f}s "
                f"for this ZeroGPU build. "
                f"Please trim the video."
            )

    except gr.Error:
        raise

    except Exception as e:

        raise gr.Error(
            f"Unable to inspect video: {e}"
        )

    # --------------------------------------------------------
    # Process video
    # --------------------------------------------------------

    return process_video(
        video_path,
        detection_prompt or "watermark",
        float(max_bbox_percent),
        max(1, int(detection_skip)),
        max(0, int(fade_in)),
        max(0, int(fade_out)),
        progress=progress,
    )


# ============================================================
# GRADIO UI
# ============================================================

with gr.Blocks(
    title="Z-Image + Watermark Remover V5.2"
) as demo:

    gr.Markdown(
        """
# 🎨 Z-Image + Watermark Remover V5.2

Z-Image-Turbo + Florence-2 + LaMa

支持：

- 🎨 AI 文生图
- 🧹 图片去水印
- 🎬 短视频去水印
- 🔌 Gradio API
"""
    )

    # ========================================================
    # TABS
    # ========================================================

    with gr.Tabs():

        # ====================================================
        # Z-IMAGE
        # ====================================================

        with gr.Tab("🎨 Z-Image"):

            with gr.Row():

                with gr.Column():

                    prompt = gr.Textbox(
                        label="Prompt",
                        placeholder=(
                            "Describe the image..."
                        ),
                        lines=5,
                    )

                    with gr.Row():

                        height = gr.Slider(
                            minimum=512,
                            maximum=2048,
                            value=1024,
                            step=64,
                            label="Height",
                        )

                        width = gr.Slider(
                            minimum=512,
                            maximum=2048,
                            value=1024,
                            step=64,
                            label="Width",
                        )

                    with gr.Row():

                        steps = gr.Slider(
                            minimum=1,
                            maximum=30,
                            value=9,
                            step=1,
                            label="Inference Steps",
                        )

                        seed = gr.Number(
                            value=42,
                            precision=0,
                            label="Seed",
                        )

                    randomize = gr.Checkbox(
                        value=True,
                        label="Randomize seed",
                    )

                    generate_btn = gr.Button(
                        "Generate",
                        variant="primary",
                    )

                with gr.Column():

                    output_image = gr.Image(
                        label="Generated Image",
                        type="pil",
                    )

                    output_seed = gr.Number(
                        label="Seed used",
                        precision=0,
                    )

            generate_btn.click(
                fn=generate_image,
                inputs=[
                    prompt,
                    height,
                    width,
                    steps,
                    seed,
                    randomize,
                ],
                outputs=[
                    output_image,
                    output_seed,
                ],
                api_name="generate_image",
            )

        # ====================================================
        # IMAGE WATERMARK
        # ====================================================

        with gr.Tab("🧹 图片去水印"):

            with gr.Row():

                with gr.Column():

                    wm_input = gr.Image(
                        type="pil",
                        label="Input image",
                    )

                    wm_prompt = gr.Textbox(
                        value="watermark",
                        label="Detection prompt",
                    )

                    wm_max = gr.Slider(
                        minimum=1,
                        maximum=90,
                        value=30,
                        step=1,
                        label="Max bbox area (%)",
                    )

                    wm_transparent = gr.Checkbox(
                        value=False,
                        label="Transparent output",
                    )

                    with gr.Row():

                        detect_btn = gr.Button(
                            "Detect / Preview"
                        )

                        remove_btn = gr.Button(
                            "Remove watermark",
                            variant="primary",
                        )

                with gr.Column():

                    wm_preview_out = gr.Image(
                        label="Detection preview",
                        type="pil",
                    )

                    wm_result = gr.Image(
                        label="Result",
                        type="pil",
                    )

                    wm_mask = gr.Image(
                        label="Mask",
                        type="pil",
                    )

            detect_btn.click(
                fn=wm_preview,
                inputs=[
                    wm_input,
                    wm_prompt,
                    wm_max,
                ],
                outputs=wm_preview_out,
                api_name="watermark_preview",
            )

            remove_btn.click(
                fn=wm_remove_image,
                inputs=[
                    wm_input,
                    wm_prompt,
                    wm_max,
                    wm_transparent,
                ],
                outputs=[
                    wm_result,
                    wm_mask,
                ],
                api_name="remove_image",
            )

        # ====================================================
        # VIDEO WATERMARK
        # ====================================================

        with gr.Tab("🎬 视频去水印"):

            gr.Markdown(
                """
### 视频去水印

ZeroGPU 版本：

- 默认最多 30 秒
- 每 N 帧进行一次 Florence-2 检测
- LaMa 对检测到的水印区域进行局部修复
- 尽量保留原视频音频
"""
            )

            # ------------------------------------------------
            # IMPORTANT:
            # Gradio 5.23.1 does NOT use type="filepath"
            # ------------------------------------------------

            video_input = gr.Video(
                label="Input video (max 30s)",
                sources=["upload"],
                format="mp4",
            )

            video_prompt = gr.Textbox(
                value="watermark",
                label="Detection prompt",
            )

            with gr.Row():

                video_max = gr.Slider(
                    minimum=1,
                    maximum=90,
                    value=30,
                    step=1,
                    label="Max bbox area (%)",
                )

                video_skip = gr.Slider(
                    minimum=1,
                    maximum=30,
                    value=15,
                    step=1,
                    label="Detection interval",
                )

            with gr.Row():

                fade_in = gr.Slider(
                    minimum=0,
                    maximum=30,
                    value=2,
                    step=1,
                    label="Fade in frames",
                )

                fade_out = gr.Slider(
                    minimum=0,
                    maximum=30,
                    value=2,
                    step=1,
                    label="Fade out frames",
                )

            video_btn = gr.Button(
                "Remove watermark from video",
                variant="primary",
            )

            video_output = gr.Video(
                label="Output video",
                format="mp4",
            )

            video_btn.click(
                fn=wm_remove_video,
                inputs=[
                    video_input,
                    video_prompt,
                    video_max,
                    video_skip,
                    fade_in,
                    fade_out,
                ],
                outputs=video_output,
                api_name="remove_video",
            )


# ============================================================
# START
# ============================================================

if __name__ == "__main__":

    print(
        "Starting Gradio Space..."
    )

    print(
        f"CUDA available: "
        f"{torch.cuda.is_available()}"
    )

    if torch.cuda.is_available():

        print(
            "CUDA device:",
            torch.cuda.get_device_name(
                torch.cuda.current_device()
            ),
        )

    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
    )