import os
import re
import random
import numpy as np

# !!! spaces must be imported before torch/CUDA
import spaces

from huggingface_hub import login
from diffusers import DiffusionPipeline
import gradio as gr
import torch

from utils import QPipeline


device = "cuda" if torch.cuda.is_available() else "cpu"

login(token=os.environ["HF_TOKEN"])
model_repo_id = os.environ["MODEL_ID"]

torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

pipe = QPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype).to(device)

MAX_SEED = 65535
MAX_IMAGE_SIZE = 128


@spaces.GPU  # Enable ZeroGPU if needed
def infer(
    prompt,
    negative_prompt,
    seed,
    randomize_seed,
    num_inference_steps=10,
    progress=gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator().manual_seed(seed)

    image = pipe(
        [prompt],
        batch_size=1,
        generator=generator,
        num_inference_steps=num_inference_steps
    ).images[0]

    return image, seed


examples = [
    "Structure: (LR 文 英). Style: style001",
    "Structure: (TL 广 東). Style: style028",
    "Structure: (TB 艹 (LR 禾 魚)). Style: style015",
    "Structure: (TB 敬 音). Style: style013",
    "Structure: (LR 釒 馬). Style: style018",
    "Structure: (BL 走 羽). Style: style022",
    "Structure: (LR 羊 大). Style: style005",
    "Structure: (LR 鹿 孚). Style: style017",
    "Structure: (OI 口 也). Style: style002",
]

# Map style images to style names (use real image files later)
style_options = {
    "images/style001.png": "style001",
    "images/style002.png": "style002",
    "images/style003.png": "style003",
    "images/style004.png": "style004",
    "images/style005.png": "style005",
    "images/style006.png": "style006",
    "images/style007.png": "style007",
    "images/style008.png": "style008",
    "images/style009.png": "style009",
    "images/style010.png": "style010",
    "images/style011.png": "style011",
    "images/style012.png": "style012",
    "images/style013.png": "style013",
    "images/style014.png": "style014",
    "images/style015.png": "style015",
    # "images/style016.png": "style016", very similar to 002
    "images/style017.png": "style017",
    "images/style018.png": "style018",
    "images/style019.png": "style019",
    "images/style020.png": "style020",
    "images/style021.png": "style021",
    "images/style022.png": "style022",
    "images/style023.png": "style023",
    "images/style024.png": "style024",
    "images/style025.png": "style025",
    "images/style026.png": "style026",
    "images/style027.png": "style027",
    "images/style028.png": "style028",
    "images/style029.png": "style029",
}


def apply_style_on_click(evt: gr.SelectData, prompt_text):
    index = evt.index
    style_label = list(style_options.values())[index]

    if re.search(r"Style: [^\n]+", prompt_text):
        return re.sub(r"Style: [^\n]+", f"Style: {style_label}", prompt_text)
    else:
        return prompt_text.strip() + f" Style: {style_label}"


# CSS for fixing Gallery layout
css = """
#col-container {
    margin: 0 auto;
    max-width: 800px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(" # NeoChar ")
        gr.Markdown(""" - Generate New Chineses Characters (Hanzi/Kanji)
                        - Combine components in a creative way
                        - Write them in style
                        - A Gen-AI's implementation of [Lin Yutang's Ming-Kwai Typewriter](https://thereader.mitpress.mit.edu/the-uncanny-keyboard/)
                        - [README](https://huggingface.co./spaces/lqume/neochar/blob/main/README.md) for more""")
        gr.Markdown(" ## QuickStart: select an example, edit components, pick a style, then 'generate'")

        gr.HTML("""
        <style>
        .gallery-container .gallery-item {
            width: 60px !important;
            height: 60px !important;
            padding: 0 !important;
            margin: 4px !important;
            border-radius: 4px;
            overflow: hidden;
            background: none !important;
            box-shadow: none !important;
        }

        .gallery-container .gallery-item img {
            width: 64px !important;
            height: 64px !important;
            object-fit: cover;
            display: block;
            margin: auto;
        }

        .gallery-container button {
            all: unset !important;
            padding: 0 !important;
            margin: 0 !important;
            border: none !important;
            background: none !important;
            box-shadow: none !important;
        }

        .gallery__modal,
        .gallery-container .preview,
        .gallery-container .gallery-item:focus-visible {
            display: none !important;
            pointer-events: none !important;
        }
        </style>
        """)

        gallery = gr.Gallery(
            value=list(style_options.keys()),
            label="Click any image",
            columns=7,
            allow_preview=False,
            height=None,
            elem_classes=["gallery-container"]
        )

        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            run_button = gr.Button("Generate", scale=0, variant="primary")

        gallery.select(
            fn=apply_style_on_click,
            inputs=[prompt],
            outputs=prompt
        )

        result = gr.Image(label="Result", show_label=False)

        with gr.Accordion("Advanced Settings", open=False):
            negative_prompt = gr.Text(
                label="Negative prompt",
                max_lines=1,
                placeholder="Enter a negative prompt",
                visible=False,
            )

            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )

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

            with gr.Row():
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=20,
                    step=1,
                    value=10,
                )

        gr.Examples(examples=examples, inputs=[prompt])

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            prompt,
            negative_prompt,
            seed,
            randomize_seed,
            num_inference_steps,
        ],
        outputs=[result, seed],
    )

if __name__ == "__main__":
    demo.launch()