#!/usr/bin/env python

import os
import random
import uuid
import json
import re
import gradio as gr
import numpy as np
from PIL import Image
import spaces
import torch
from diffusers import DiffusionPipeline
from typing import Tuple
from transformers import pipeline

# Setup rules for bad words (ensure the prompts are kid-friendly)
bad_words = json.loads(os.getenv('BAD_WORDS', '["violence", "blood", "scary", "death", "ghost"]'))
default_negative = os.getenv("default_negative","")

# Add the translation pipeline
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")

def check_text(prompt, negative=""):
    for i in bad_words:
        if i in prompt:
            return True
    return False

# Kid-friendly styles
style_list = [
    {
        "name": "만화",
        "prompt": "colorful cartoon {prompt}. vibrant, playful, friendly, suitable for children, highly detailed, bright colors",
        "negative_prompt": "scary, dark, violent, ugly, realistic",
    },
    {
        "name": "어린이 일러스트",
        "prompt": "children's illustration {prompt}. cute, colorful, fun, simple shapes, smooth lines, highly detailed, joyful",
        "negative_prompt": "scary, dark, violent, deformed, ugly",
    },   
    {
        "name": "스티커",
        "prompt": "children's sticker of {prompt}. bright colors, playful, high resolution, cartoonish",
        "negative_prompt": "scary, dark, violent, ugly, low resolution",
    },
    {
        "name": "판타지",
        "prompt": "fantasy world for children with {prompt}. magical, vibrant, friendly, beautiful, colorful",
        "negative_prompt": "dark, scary, violent, ugly, realistic",
    },
    {
        "name": "(스타일 없음)",
        "prompt": "{prompt}",
        "negative_prompt": "",
    },
]

styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
STYLE_NAMES = list(styles.keys())
DEFAULT_STYLE_NAME = "스티커"

def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
    return p.replace("{prompt}", positive), n + negative

DESCRIPTION = """## 어린이 스티커 생성기
AI를 사용하여 재미있고 귀여운 어린이 스티커를 생성합니다.
"""

if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>⚠️CPU에서 실행 중입니다. CPU에서는 제대로 작동하지 않을 수 있습니다.</p>"

MAX_SEED = np.iinfo(np.int32).max
CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "1") == "1"

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

# Initialize the DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained(
    "SG161222/RealVisXL_V3.0_Turbo",  # or any model of your choice
    torch_dtype=torch.float16,
    use_safetensors=True,
    variant="fp16"
).to(device)

# Convert mm to pixels for a specific DPI (300) and ensure divisible by 8
def mm_to_pixels(mm, dpi=300):
    """Convert mm to pixels and make the dimensions divisible by 8."""
    pixels = int((mm / 25.4) * dpi)
    return pixels - (pixels % 8)  # Adjust to the nearest lower multiple of 8

# Default sizes for 75mm and 35mm, rounded to nearest multiple of 8
size_map = {
    "75mm": (mm_to_pixels(75), mm_to_pixels(75)),  # 75mm in pixels at 300dpi
    "35mm": (mm_to_pixels(35), mm_to_pixels(35)),  # 35mm in pixels at 300dpi
}

# Function to post-process images (transparent or white background)
def save_image(img, background="transparent"):
    img = img.convert("RGBA")
    data = img.getdata()
    new_data = []

    if background == "transparent":
        for item in data:
            # Replace white with transparent
            if item[0] == 255 and item[1] == 255 and item[2] == 255:
                new_data.append((255, 255, 255, 0))  # Transparent
            else:
                new_data.append(item)
    elif background == "white":
        for item in data:
            new_data.append(item)  # Keep as white

    img.putdata(new_data)
    unique_name = str(uuid.uuid4()) + ".png"
    img.save(unique_name)
    return unique_name

def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed

def translate_if_korean(text):
    # Check if the text contains Korean characters
    if re.search("[\uac00-\ud7a3]", text):
        # Translate Korean to English
        translation = translator(text, max_length=512)
        return translation[0]['translation_text']
    return text

@spaces.GPU(enable_queue=True)
def generate(
    prompt: str,
    negative_prompt: str = "",
    use_negative_prompt: bool = False,
    style: str = DEFAULT_STYLE_NAME,
    seed: int = 0,
    size: str = "75mm",
    guidance_scale: float = 3,
    randomize_seed: bool = False,
    background: str = "transparent",
    progress=gr.Progress(track_tqdm=True),
):
    # Translate prompt if it's in Korean
    prompt = translate_if_korean(prompt)
    
    if check_text(prompt, negative_prompt):
        raise ValueError("프롬프트에 제한된 단어가 포함되어 있습니다.")
    
    # Ensure prompt is 2-3 words long
    prompt = " ".join(re.findall(r'\w+', prompt)[:3])

    # Apply style
    prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
    seed = int(randomize_seed_fn(seed, randomize_seed))
    generator = torch.Generator().manual_seed(seed)

    # Ensure we have only white or transparent background options
    width, height = size_map.get(size, (1024, 1024))

    if not use_negative_prompt:
        negative_prompt = ""  # type: ignore

    options = {
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "width": width,
        "height": height,
        "guidance_scale": guidance_scale,
        "num_inference_steps": 50,
        "generator": generator,
        "num_images_per_prompt": 1,  # Max 6 images
        "output_type": "pil",
    }
    
    # Generate images with the pipeline
    images = pipe(**options).images
    image_paths = [save_image(img, background) for img in images]
    
    return image_paths, seed

examples = [
    "귀여운 고양이",
    "행복한 토끼",
    "웃고있는 강아지",
    "춤추는 돌고래",
    "신나는 아기 독수리",
    "즐거운 아기 사자",    
]

css = """
footer {
    visibility: hidden;
}
"""

# Define the Gradio UI for the sticker generator
with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as demo:

    with gr.Group():
        with gr.Row():
            prompt = gr.Text(
                label="프롬프트 입력",
                show_label=False,
                max_lines=1,
                placeholder="프롬프트를 입력하세요.",
                container=False,
            )
            run_button = gr.Button("실행")
        result = gr.Gallery(label="생성된 스티커", columns=2, preview=True)
    with gr.Accordion("고급 옵션", open=False):
        use_negative_prompt = gr.Checkbox(label="네거티브 프롬프트 사용", value=True, visible=True)
        negative_prompt = gr.Text(
            label="네거티브 프롬프트",
            max_lines=1,
            placeholder="네거티브 프롬프트를 입력하세요",
            value="(무서운, 폭력적인, 어두운, 추한)",
            visible=True,
        )
        seed = gr.Slider(
            label="시드",
            minimum=0,
            maximum=MAX_SEED,
            step=1,
            value=0,
            visible=True
        )
        randomize_seed = gr.Checkbox(label="랜덤 시드", value=True)
        size_selection = gr.Radio(
            choices=["75mm", "35mm"],
            value="75mm",
            label="스티커 크기",
        )
        style_selection = gr.Radio(
            choices=STYLE_NAMES,
            value=DEFAULT_STYLE_NAME,
            label="이미지 스타일",
        )
        background_selection = gr.Radio(
            choices=["투명", "흰색"],
            value="흰색",
            label="배경색",
        )
        guidance_scale = gr.Slider(
            label="가이던스 스케일",
            minimum=0.1,
            maximum=20.0,
            step=0.1,
            value=6,
        )

    gr.Examples(
        examples=examples,
        inputs=prompt,
        outputs=[result, seed],
        fn=generate,
        cache_examples=CACHE_EXAMPLES,
    )

    gr.on(
        triggers=[
            prompt.submit,
            negative_prompt.submit,
            run_button.click,
        ],
        fn=generate,
        inputs=[
            prompt,
            negative_prompt,
            use_negative_prompt,
            style_selection,
            seed,
            size_selection,
            guidance_scale,
            randomize_seed,
            background_selection,
        ],
        outputs=[result, seed],
        api_name="run",
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch()