from __future__ import annotations

import sys
from pathlib import Path

import gradio as gr


def _patch_gradio_client_schema_bug():
    try:
        import gradio_client.utils as client_utils
    except Exception:
        return
    original = getattr(client_utils, "_json_schema_to_python_type", None)
    if original is None or getattr(original, "_bila_bool_schema_patch", False):
        return

    def patched(schema, defs=None):
        if isinstance(schema, bool):
            return "Any"
        if isinstance(schema, dict) and isinstance(schema.get("additionalProperties"), bool):
            schema = dict(schema)
            schema.pop("additionalProperties", None)
        return original(schema, defs)

    patched._bila_bool_schema_patch = True
    client_utils._json_schema_to_python_type = patched


_patch_gradio_client_schema_bug()

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT / "vendor"))

try:
    import spaces
except ImportError:
    class _SpacesFallback:
        @staticmethod
        def GPU(*args, **kwargs):
            if args and callable(args[0]) and len(args) == 1 and not kwargs:
                return args[0]

            def decorator(fn):
                return fn

            return decorator

    spaces = _SpacesFallback()

from demo_runtime.manager import DemoManager


manager = DemoManager()
DEFAULT_MODEL = manager.default_model
EXAMPLE_DIR = ROOT / "assets" / "examples"
EXAMPLES = [
    [str(EXAMPLE_DIR / "Expert338-0-0-Best-Input.jpg"), "Make the scene more dramatic with a warm sunset vibe and richer colors for a cinematic travel magazine feel.", 42, 1024, 1.0],
    [str(EXAMPLE_DIR / "Expert357-0-0-Best-Input.jpg"), "Make the scene feel like a warm autumn evening with a cozy, inviting travel magazine feel.", 7, 1024, 1.0],
    [str(EXAMPLE_DIR / "Expert352-0-0-Best-Input.jpg"), "Make the scene feel more cinematic with a moody, atmospheric vibe while keeping the person and cityscape prominent.", 123, 1024, 0.9],
    [str(EXAMPLE_DIR / "Expert394-0-0-Best-Input.jpg"), "Make the scene more dramatic and moody, emphasizing the sunset glow and enhancing the lanterns' warmth for a cinematic feel.", 314, 1024, 1.0],
    [str(EXAMPLE_DIR / "Expert388-0-0-Best-Input.jpg"), "Make the flower look timeless and classic with a black-and-white vintage tone, emphasizing its delicate details.", 42, 1024, 1.0],
]


@spaces.GPU(duration=60, size="xlarge")
def run_demo(image, instruction, seed, max_side, strength):
    try:
        edited, _diff, _input_image, status = manager.generate(
            image=image,
            instruction=instruction,
            model_key=DEFAULT_MODEL,
            seed=int(seed),
            max_side=int(max_side),
            strength=float(strength),
        )
        return edited, status
    except Exception as exc:
        raise gr.Error(str(exc))


with gr.Blocks(title="InstantRetouch") as demo:
    gr.Markdown(
        """
        # InstantRetouch
        Instruction-guided photo retouching that is fast and high-fidelity. Upload an image, enter an optional instruction, or click one of the examples below.

        Please note that running on HF space requires extra time for GPU allocation, environment preperation and startup, so the processing time will be longer that the pure inference time reported in our paper. The strength slider allows continous control on editing strength.
        """
    )
    with gr.Row():
        with gr.Column(scale=1):
            image = gr.Image(type="pil", label="Input image")
            instruction = gr.Textbox(label="Instruction", lines=3, placeholder="Optional. Leave empty for prompt=\"\".")
            with gr.Row():
                seed = gr.Number(value=42, precision=0, label="Seed")
                max_side = gr.Slider(512, 2048, value=1024, step=64, label="Max side")
            strength = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="Strength")
            button = gr.Button("Run", variant="primary")
        with gr.Column(scale=1):
            edited = gr.Image(type="pil", label="Edited image")
            status = gr.Textbox(label="Status", interactive=False)

    gr.Examples(
        examples=EXAMPLES,
        inputs=[image, instruction, seed, max_side, strength],
        examples_per_page=5,
        cache_examples=False,
    )

    button.click(
        fn=run_demo,
        inputs=[image, instruction, seed, max_side, strength],
        outputs=[edited, status],
    )


if __name__ == "__main__":
    try:
        demo.queue(default_concurrency_limit=1, max_size=8)
    except TypeError:
        demo.queue(concurrency_count=1, max_size=8)
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_api=False,
    )