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;
}
.center-text {
    text-align: center;
}
.radical-section {
    max-width: 400px;
    margin: 10px auto;
    text-align: center;
}
.radical-buttons {
    display: flex;
    flex-wrap: wrap;
    justify-content: center;
    gap: 8px;
    margin-bottom: 10px;
}
.radical-buttons button {
    padding: 8px 10px;
    font-size: 10px;
    border-radius: 50%;
    border: 2px solid #aaa;
    background: radial-gradient(circle at top left, #f5f5dc, #dcdcc4);
    cursor: pointer;
    box-shadow: 2px 2px 5px rgba(0,0,0,0.3), inset 0 1px 2px rgba(255,255,255,0.8);
    transition: transform 0.1s ease-in-out;
}
.radical-buttons button:active {
    transform: scale(0.95);
}
.radical-buttons button.dim {
    opacity: 0.25;
    filter: grayscale(100%);
}
.radical-token {
  display: inline-block;
  margin: 2px;
  padding: 6px 10px;
  border: 1px solid var(--border-color-primary);
  border-radius: 8px;
  /* show text cursor so users know they can select */
  cursor: text;
}
#pinyin-display {
    text-align: center;
    margin: 20px;
    font-size: 18px;
    font-weight: bold;
    color: #555;
    font-family: monospace;
}
"""

A_ALL = [
    "氵","木","扌","釒","口","亻","女","忄","糹","言","虫","土","火","月","王","石",
    "魚","钅","⻊","衤","山","犭","目","日","馬","禾","讠","纟","阝","車","礻","米",
    "酉","飠","鱼","革","舟","巾","貝","彳","牛","耳","歹","冫","马","饣","骨","车",
    "弓","齒","立","豸","白","田","角","身","贝","耒","黑","⺶","鼠","豕","韋","鳥",
    "缶","片","麥","牜","谷","香","子","舌","矢","鼻","羽","爿","方","音","赤","齿",
    "又","青","镸","血","高","夕","句","豆","工","雚","至","束","吉","矛","龺","文",
    "令","合","鬲",
    
    "艹","⺮","宀","⻗","穴","髟","罒","一","大","人","辟","冖","敝","此","⺈","亦",
    "敖","亠","䜌","林","丿","龹","丷","折","非","敄","次","龍","⺷","爫","加","臤",
    "⺊","分","般","殹","丶","厶","亡","奴","殸","執","秋","夂","十","代","覀","不",
    "龶","二","水","尚","比","夗","沙","列","执","其","斬","鼓","吅","小","龸","今",
    "天","將","自","八","士","丆","亼","父","與","任","耂","臼","龙","壯","如","巩",
    "敬","埶","畾","賏","玨",
    
    "囗","又","卩","井"
]

B_ALL = [
    "鳥","阝","刂","頁","力","欠","鸟","攵","見","隹","皮","页","令","各","斤","且",
    "殳","瓦","台","句","分","羽","戈","占","干","夋","召","包","交","合","卑","肖",
    "青","者","乚","甫","毛","勺","可","今","奇","犮","周","扁","需","堯","兼","屯",
    "兆","昷","果","枼","翏","丁","艮","氐","并","同","卒","是","昜","單","圭","昆",
    "曷","票","蒦","俞","乞","吉","良","它","乍","炎","至","軍","粦","番","也","亢",
    "巨","失","其","旁","亥","反","余","工","少","方","古","由","喬","尞","龍","襄",
    "皇","非","及",
    
    "心","土","木","虫","女","皿","金","日","火","口","灬","貝","手","山","衣","一",
    "贝","石","言","糸","月","巾","乙","儿","目","大","子","田","力","水","廾","几",
    "馬","耳","牛","米","寸","毛","小","足","魚","八","示","玉","厶","又","夕","巴",
    "食","马","犬","弓","刀","瓦","用","肉","王","黽","人","二","夂","电","旦","云",
    "辰","立","車","且","鹿","缶","乂","匕","可","氺","共","天","叱","卯","豆","林",
    "羽","酉","角","革","異","鱼","丁","⺀","朩","龰","卜","丂",
    
    "丶","儿","厶","才","寸","屯","亢","不","云","勿","元","化","仑","古","民","台",
    "冬","豕","甫","吾","侖","卷","或","韋","巻","尃","書","專","睘","韦","禾","正",
    "有","员","幸","耑","袁","玉","貟","員","盍","啚","貴","睪","欒"
]

def make_button_grid(items):
    # Wrap buttons in a container with your old CSS classes
    buttons = "\n".join(
        f"<span class='radical-token' data-val='{c}'>{c}</span>"
        for c in items
    )
    return f"<div class='radical-section'><div class='radical-buttons'>{buttons}</div></div>"

with gr.Blocks(css=css) as demo:
    with gr.Row():
        with gr.Column(elem_id="radicals-a", scale=1, min_width=0):
            gr.Markdown("For hard-to-input radicals, copy from below", elem_classes="center-text")
            gr.Markdown("(Left/Top/Outer)", elem_classes="center-text")
            gr.HTML(make_button_grid(A_ALL))
        with gr.Column(elem_id="col-container", scale=0, min_width=800):
            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])

        with gr.Column(elem_id="radicals-b", scale=1, min_width=0):
            gr.Markdown("For hard-to-input radicals, copy from below", elem_classes="center-text")
            gr.Markdown("(Right/Bottom/Inner)", elem_classes="center-text")
            gr.HTML(make_button_grid(B_ALL))


    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()