"""
TextRIRO Gradio Demo
====================
Run with:
    python app.py
By default this launches a Gradio web UI at http://0.0.0.0:7860.

Given a style reference image (i_s, a scene-text image) and a target text
(i_t), the model generates a new image that keeps the original style
(font, color, background texture, etc.) while rendering the target text.
"""
import os

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

try:
    import spaces  # MUST come before torch / any CUDA-touching import

    _HAS_ZEROGPU = True
except ImportError:
    _HAS_ZEROGPU = False

    class _NoOpSpaces:
        """Fallback no-op decorator so the app also runs outside HF ZeroGPU Spaces."""

        @staticmethod
        def GPU(*args, **kwargs):
            if len(args) == 1 and callable(args[0]):
                return args[0]

            def decorator(fn):
                return fn

            return decorator

    spaces = _NoOpSpaces()

import random

import gradio as gr
import numpy as np
import torch
from PIL import Image

from infer_pipeline import InferPipeline, load_style_image
from pipeline_builder import TextRIROModel


def seed_everything(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_REPO = "Yesianrohn/TextRIRO"
CKPT_FILENAME = "textriro.pth"
LOCAL_CKPT_PATH = os.path.join(CURRENT_DIR, "weights", CKPT_FILENAME)

# NOTE: do NOT rely on `_HAS_ZEROGPU` (whether `import spaces` succeeded) to
# decide the device. `spaces` can be installed and imported successfully even
# on plain CPU hardware, in which case there is no GPU at all. On real
# ZeroGPU hardware, `spaces` patches `torch.cuda.is_available()` to report
# `True` even before a physical GPU is attached, so checking it directly is
# the correct and safe way to detect ZeroGPU vs. CPU-only Spaces.
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# `TEXTRIRO_CKPT` may point to a local file/dir; otherwise fall back to the
# bundled `weights/` folder, otherwise download from the Hub automatically.
CKPT_SOURCE = os.environ.get("TEXTRIRO_CKPT") or (LOCAL_CKPT_PATH if os.path.exists(LOCAL_CKPT_PATH) else MODEL_REPO)
print(f"Loading TextRIRO model from: {CKPT_SOURCE}")

# ---------------------------------------------------------------------------
# Model loading at module scope (ZeroGPU: .to("cuda") is intercepted here and
# the real transfer happens lazily inside the @spaces.GPU-decorated function)
# ---------------------------------------------------------------------------
_model = TextRIROModel.from_pretrained(CKPT_SOURCE, filename=CKPT_FILENAME, device="cpu", dtype=torch.float32)
_model.to(DEVICE)
print(f"Model loaded and moved to {DEVICE}.")

_pipeline = InferPipeline(_model)
SUPPORTED_CHARS = set(_model.text_encoder.character)


def find_unsupported_chars(text: str):
    """Return the sorted list of characters in `text` that are not in the model's vocabulary."""
    return sorted({c for c in text if c not in SUPPORTED_CHARS})


@spaces.GPU(duration=60)
def run_inference(
    style_image: Image.Image,
    target_text: str,
    num_inference_steps: int,
    guidance_scale: float,
    seed: int,
):
    if style_image is None:
        raise gr.Error("Please upload a style reference image (i_s).")
    if not target_text or len(target_text.strip()) == 0:
        raise gr.Error("Please enter the target text (i_t).")

    unsupported = find_unsupported_chars(target_text)
    if unsupported:
        allowed = "".join(sorted(SUPPORTED_CHARS - {" "}))
        raise gr.Error(
            f"Unsupported character(s): {' '.join(unsupported)}. "
            f"Please re-enter your text using only the supported characters: "
            f"{allowed} (and space)."
        )

    seed_everything(int(seed))

    orig_w, orig_h = style_image.size
    hint = load_style_image(style_image, image_height=64, image_width=256, device=DEVICE)

    result = _pipeline(
        hint=hint,
        cond=[target_text],
        num_inference_steps=int(num_inference_steps),
        guidance_scale=float(guidance_scale),
    )[0]

    result = result.clamp(0, 1).cpu().numpy()
    result = np.transpose(result, (1, 2, 0))
    result_img = Image.fromarray((result * 255).astype(np.uint8)).resize((orig_w, orig_h))
    return result_img


EXAMPLE_DIR = os.path.join(CURRENT_DIR, "examples")
examples = []
if os.path.isdir(EXAMPLE_DIR):
    for fname in sorted(os.listdir(EXAMPLE_DIR)):
        if fname.lower().endswith((".png", ".jpg", ".jpeg")):
            examples.append([os.path.join(EXAMPLE_DIR, fname), "Hello World"])

with gr.Blocks(title="TextRIRO Demo") as demo:
    gr.Markdown(
        """
        # TextRIRO: Scene Text Image Editing Demo
        Upload a scene-text image and enter the text you want it to display.
        The model keeps the original style (font, color, background texture,
        etc.) while replacing the text.

        Model weights: [Yesianrohn/TextRIRO](https://huggingface.co/Yesianrohn/TextRIRO).
        """
    )
    with gr.Row():
        with gr.Column():
            style_image = gr.Image(type="pil", label="Style reference image (i_s)")
            target_text = gr.Textbox(label="Target text (i_t)", placeholder="e.g. Hello World")
            with gr.Accordion("Advanced settings", open=False):
                num_inference_steps = gr.Slider(5, 50, value=20, step=1, label="Number of inference steps")
                guidance_scale = gr.Slider(1.0, 10.0, value=2.2, step=0.1, label="Guidance scale")
                seed = gr.Number(value=42, precision=0, label="Seed")
            run_btn = gr.Button("Generate", variant="primary")
        with gr.Column():
            output_image = gr.Image(type="pil", label="Generated result")

    run_btn.click(
        fn=run_inference,
        inputs=[style_image, target_text, num_inference_steps, guidance_scale, seed],
        outputs=output_image,
        api_name="generate",
    )

    if len(examples) > 0:
        gr.Examples(examples=examples, inputs=[style_image, target_text])

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860)