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 / "4920_O_0_5_input.png"), "Make the image feel more serene and add a subtle blue hue.", 42, 1024, 1.0],
    [str(EXAMPLE_DIR / "4933_O_0_21_input.png"), "Improve the exposure and make the colors richer while keeping a natural photo look.", 7, 1024, 1.0],
    [str(EXAMPLE_DIR / "expert48_input.png"), "Brighten the image and enhance clarity with balanced contrast.", 123, 1024, 0.9],
    [str(EXAMPLE_DIR / "expert116_input.png"), "", 314, 1024, 1.0],
]


@spaces.GPU(duration=300, 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 with the selected IP2P/BiLA checkpoint. Upload an image, enter an optional instruction, or click one of the examples below.

        This public demo uses the validation-selected IP2P/BiLA model only. The strength slider blends the model output with the input for gentler or stronger edits.
        """
    )
    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=4,
        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,
    )