import gradio as gr
from transformers import MarianMTModel, MarianTokenizer
import torch
import random

# ตรวจสอบว่าใช้งาน GPU หรือไม่ (ในกรณีนี้ไม่มี GPU)
device = torch.device("cpu")

# เลือกโมเดลแปลภาษาที่มีขนาดเล็กลง
# คุณสามารถค้นหาโมเดลที่เล็กกว่าได้ที่ Hugging Face Model Hub
# ในตัวอย่างนี้ยังคงใช้ "Helsinki-NLP/opus-mt-th-en" แต่คุณสามารถเปลี่ยนได้ตามต้องการ
model_name = "Helsinki-NLP/opus-mt-th-en"
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name).to(device)
model.eval()  # ตั้งค่าโมเดลเป็นโหมดประเมินผล

# แคชสำหรับการแปลข้อความที่ไม่อยู่ในรายการคงที่
translation_cache = {}

def translate_batch(th_texts):
    """แปลข้อความแบบแบทช์เพื่อเพิ่มความเร็ว"""
    inputs = tokenizer(th_texts, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
    with torch.no_grad():
        translated = model.generate(**inputs, max_length=512)
    en_texts = [tokenizer.decode(t, skip_special_tokens=True) for t in translated]
    return en_texts

def translate_th_to_en(th_text: str) -> str:
    th_text = th_text.strip()
    if not th_text:
        return ""
    # ตรวจสอบแคชก่อน
    if th_text in translation_cache:
        return translation_cache[th_text]
    # แปลข้อความ
    en_text = translate_batch([th_text])[0]
    # เก็บในแคช
    translation_cache[th_text] = en_text
    return en_text

# ตัวเลือก (ภาษาไทย) สำหรับ Dropdown
body_shapes_th = [
    "ตัวเหนียวขยุกขยิก", "ตัวปุกปุยกลม", "ตัวเต็มไปด้วยหนาม",
    "ตัวเล็กเหมือนแมลง", "ตัวโปร่งแสงลอยได้", "ตัวโคลนยืดหยุ่น",
    "ตัวหินแข็งแรง", "ตัวทรงพีระมิด", "ตัวระยิบระยับ", "ตัวเหมือนฟองสบู่"
]
heads_th = [
    "หัวตาเดียว", "หัวสองเขาเหมือนมังกร", "หัวมีเขาใหญ่มาก",
    "หัวผีแสยะยิ้ม", "หัวกระต่ายน่ารัก", "หัวปลาหมึกสุดแปลก",
    "หัวเพลิงลุก", "หัวงูพ่นพิษ", "หัววุ้นใส", "หัวเมฆหมอก"
]
arms_legs_th = [
    "แขนหนวดปลาหมึกและขาปลาหมึก", "แขนกลและเท้าสเก็ต",
    "ปีกขนนกและอุ้งเท้าสิงโต", "แขนเรืองแสงและขายานโฮเวอร์",
    "แขนไม้หุ่นเชิดและขายักษ์แมงมุม", "แขนโลหะ",
    "แขนเหล็กข้อต่อ", "ปีกค้างคาว", "ขาไก่ยักษ์", "แขนขนนุ่ม"
]
skin_patterns_th = [
    "ลายสีรุ้ง", "จุดม่วงเรืองแสง", "เกล็ดเหล็ก",
    "ขนนกวิบวับ", "จุดกลมกากเพชร", "ลายหินอ่อน",
    "ลายไฟลุก", "ลายฟ้าแลบ", "ลายสเกลปลา", "ลายหมอก"
]
abilities_th = [
    "พ่นไฟ", "ปล่อยฟองสบู่", "ควบคุมสายฟ้า",
    "ร้องเพลงกล่อม", "วาร์ประยะสั้น", "เสกสิ่งของ",
    "ปล่อยพลังงานแสง", "ควบคุมแรงโน้มถ่วง", "เปลี่ยนรูปร่าง", "เสกภาพลวงตา"
]
moods_th = [
    "เป็นมิตร", "ขี้หงุดหงิด", "ตลกโปกฮา",
    "ซนเป็นลิง", "ขี้อาย", "ขี้เล่น",
    "เศร้าสร้อย", "จริงจัง", "มีความสุข", "ขี้ตกใจ"
]

# แปลล่วงหน้าสำหรับรายการคงที่
# คุณสามารถแปลข้อความเหล่านี้ล่วงหน้านอกโปรแกรมและใส่ค่าลงไปตรงนี้เพื่อความรวดเร็ว
body_shapes_en = {
    "ตัวเหนียวขยุกขยิก": "a squishy and wobbly body",
    "ตัวปุกปุยกลม": "a round and fluffy body",
    "ตัวเต็มไปด้วยหนาม": "a spiky body",
    "ตัวเล็กเหมือนแมลง": "a small, insect-like body",
    "ตัวโปร่งแสงลอยได้": "a translucent, floating body",
    "ตัวโคลนยืดหยุ่น": "a flexible, slime-like body",
    "ตัวหินแข็งแรง": "a strong, rocky body",
    "ตัวทรงพีระมิด": "a pyramidal body",
    "ตัวระยิบระยับ": "a sparkling body",
    "ตัวเหมือนฟองสบู่": "a bubble-like body"
}

heads_en = {
    "หัวตาเดียว": "a single-eyed head",
    "หัวสองเขาเหมือนมังกร": "a two-horned dragon-like head",
    "หัวมีเขาใหญ่มาก": "a head with very large horns",
    "หัวผีแสยะยิ้ม": "a smiling ghostly head",
    "หัวกระต่ายน่ารัก": "a cute rabbit head",
    "หัวปลาหมึกสุดแปลก": "a bizarre squid head",
    "หัวเพลิงลุก": "a flaming head",
    "หัวงูพ่นพิษ": "a venomous snake head",
    "หัววุ้นใส": "a transparent jelly head",
    "หัวเมฆหมอก": "a misty cloud head"
}

arms_legs_en = {
    "แขนหนวดปลาหมึกและขาปลาหมึก": "octopus-like tentacles and legs",
    "แขนกลและเท้าสเก็ต": "mechanical arms and skating feet",
    "ปีกขนนกและอุ้งเท้าสิงโต": "feathered wings and lion-like paws",
    "แขนเรืองแสงและขายานโฮเวอร์": "glowing arms and hover jets",
    "แขนไม้หุ่นเชิดและขายักษ์แมงมุม": "puppet wooden arms and giant spider legs",
    "แขนโลหะ": "metal arms",
    "แขนเหล็กข้อต่อ": "jointed iron arms",
    "ปีกค้างคาว": "bat wings",
    "ขาไก่ยักษ์": "giant chicken legs",
    "แขนขนนุ่ม": "soft feathered arms"
}

skin_patterns_en = {
    "ลายสีรุ้ง": "rainbow patterns",
    "จุดม่วงเรืองแสง": "glowing purple spots",
    "เกล็ดเหล็ก": "iron scales",
    "ขนนกวิบวับ": "shimmering feathers",
    "จุดกลมกากเพชร": "diamond-like round dots",
    "ลายหินอ่อน": "marble patterns",
    "ลายไฟลุก": "flaming patterns",
    "ลายฟ้าแลบ": "lightning patterns",
    "ลายสเกลปลา": "fish scale patterns",
    "ลายหมอก": "misty patterns"
}

abilities_en = {
    "พ่นไฟ": "breathe fire",
    "ปล่อยฟองสบู่": "release soap bubbles",
    "ควบคุมสายฟ้า": "control lightning",
    "ร้องเพลงกล่อม": "sing a lullaby",
    "วาร์ประยะสั้น": "short-range warp",
    "เสกสิ่งของ": "conjure objects",
    "ปล่อยพลังงานแสง": "emit light energy",
    "ควบคุมแรงโน้มถ่วง": "control gravity",
    "เปลี่ยนรูปร่าง": "shape-shift",
    "เสกภาพลวงตา": "create illusions"
}

moods_en = {
    "เป็นมิตร": "friendly",
    "ขี้หงุดหงิด": "irritable",
    "ตลกโปกฮา": "funny",
    "ซนเป็นลิง": "mischievous like a monkey",
    "ขี้อาย": "shy",
    "ขี้เล่น": "playful",
    "เศร้าสร้อย": "sorrowful",
    "จริงจัง": "serious",
    "มีความสุข": "happy",
    "ขี้ตกใจ": "easily startled"
}

def generate_monster_lab(
    body_dd, body_tb,
    head_dd, head_tb,
    arms_dd, arms_tb,
    skin_dd, skin_tb,
    ability_dd, ability_tb,
    mood_dd, mood_tb,
):
    # ถ้า textbox มีข้อความ -> ใช้ค่านั้น, ถ้าว่าง -> ใช้ dropdown
    body_th = body_tb.strip() if body_tb.strip() else body_dd
    head_th = head_tb.strip() if head_tb.strip() else head_dd
    arms_th = arms_tb.strip() if arms_tb.strip() else arms_dd
    skin_th = skin_tb.strip() if skin_tb.strip() else skin_dd
    ability_th = ability_tb.strip() if ability_tb.strip() else ability_dd
    mood_th = mood_tb.strip() if mood_tb.strip() else mood_dd

    desc_th = (
        f"มอนสเตอร์อารมณ์{mood_th} มี{body_th} "
        f"พร้อมด้วย{head_th}, {arms_th}, "
        f"ผิว/ลายแบบ{skin_th}, และพลังพิเศษคือ{ability_th}!"
    )

    # แปลเป็นภาษาอังกฤษ โดยใช้การแปลล่วงหน้าสำหรับรายการคงที่
    body_en = body_shapes_en.get(body_th, translate_th_to_en(body_th))
    head_en = heads_en.get(head_th, translate_th_to_en(head_th))
    arms_en = arms_legs_en.get(arms_th, translate_th_to_en(arms_th))
    skin_en = skin_patterns_en.get(skin_th, translate_th_to_en(skin_th))
    ability_en = abilities_en.get(ability_th, translate_th_to_en(ability_th))
    mood_en = moods_en.get(mood_th, translate_th_to_en(mood_th))

    prompt_en = (
        f"This is a {mood_en} monster with {body_en}, "
        f"a {head_en}, {arms_en}, covered in {skin_en} skin, "
        f"and it can {ability_en}!"
    )

    return desc_th, prompt_en

def random_monster():
    bd = random.choice(body_shapes_th)
    hd = random.choice(heads_th)
    ar = random.choice(arms_legs_th)
    sk = random.choice(skin_patterns_th)
    ab = random.choice(abilities_th)
    md = random.choice(moods_th)
    # Clear textbox
    return (bd, "", hd, "", ar, "", sk, "", ab, "", md, "")

# ---- CSS ----
css_code = """
body {
    background-color: #FFF7EA;
    font-family: "Kanit", sans-serif;
}

#title {
    color: #FF6F91;
    text-align: center;
    font-size: 2rem;
    margin-top: 20px;
    margin-bottom: 10px;
    font-weight: bold;
}
.game-desc {
    margin: 0 auto;
    width: 80%;
    background-color: #FFF2F0;
    border: 2px dashed #FAB1A0;
    border-radius: 10px;
    padding: 15px;
    color: #333;
    margin-bottom: 20px;
}

/* ปุ่มหลัก (สร้างมอนสเตอร์) */
.btn-main {
    background-color: #FFC107;
    border: 2px solid #FFA000;
    font-weight: bold;
    font-size: 1.1rem;
    /* กำหนดความสูงและความยาว (horizontal padding) เยอะๆ */
    padding: 10px 40px;  
    border-radius: 10px;
    margin-right: 10px;
}

/* ปุ่มสุ่ม (รอง) => ตัวเล็กลงในความยาว (horizontal padding) แต่ความสูงเท่ากัน */
.btn-random {
    background-color: #FFE08E; 
    border: 2px solid #FFC107;
    font-weight: normal;
    font-size: 1.1rem;
    /* ใช้ padding vertical เท่ากัน แต่ลด horizontal padding ลง */
    padding: 10px 15px;
    border-radius: 8px;
}
.btn-random:hover {
    background-color: #FFF3C4;
}

#desc-th, #prompt-en {
    background-color: #FFFAE6;
    border: 2px solid #FFE082;
    border-radius: 10px;
    padding: 10px;
    margin-bottom: 20px;
}
"""

def initial_text():
    return "ยังไม่ได้สร้างมอนสเตอร์ ลองเลือกหรือพิมพ์ แล้วกด 'สร้างมอนสเตอร์!' หรือจะกดสุ่มดูก็ได้จ้า"

# ปุ่ม copy
copy_button_html = """
<button style="background-color: #F06292; border: 2px solid #E91E63; font-weight: bold; 
               font-size: 1.1rem; padding: 10px 20px; border-radius: 10px;"
        onclick="copyPromptText()">
  Copy Prompt
</button>
<script>
function copyPromptText() {
  const promptBox = document.querySelector('#prompt-en textarea');
  if (!promptBox) {
    alert('ไม่พบข้อความ Prompt!');
    return;
  }
  const promptText = promptBox.value;
  navigator.clipboard.writeText(promptText);
  alert('คัดลอก Prompt แล้ว!');
}
</script>
"""

with gr.Blocks(css=css_code) as demo:
    gr.Markdown("<h1 id='title'>ZenityX Monster Lab</h1>")
    gr.Markdown("""
    <div class="game-desc">
        <p>หนูน้อยจ๊ะ เลือกได้ว่าจะใช้ค่าใน <strong>Dropdown</strong> หรือจะ <strong>พิมพ์เอง</strong></p>
        <p>หากอยากลุ้นโชค กดปุ่ม (ขวามือ) <strong>"สุ่มมอนสเตอร์สุดเซอร์ไพรส์!"</strong></p>
        <p>เมื่อพร้อมแล้ว กดปุ่ม (ซ้าย) <strong>"สร้างมอนสเตอร์!"</strong> เพื่อดูรายละเอียดและ Prompt</p>
    </div>
    """)

    with gr.Row():
        with gr.Column():
            body_dd = gr.Dropdown(body_shapes_th, label="รูปร่าง", value=body_shapes_th[0])
            body_tb = gr.Textbox(label="หรือพิมพ์เอง", placeholder="ตัวโคลนยืดหยุ่น...")

            head_dd = gr.Dropdown(heads_th, label="หัว/ใบหน้า", value=heads_th[0])
            head_tb = gr.Textbox(label="หรือพิมพ์เอง", placeholder="หัวปลาหมึกสุดแปลก...")

        with gr.Column():
            arms_dd = gr.Dropdown(arms_legs_th, label="แขนขา", value=arms_legs_th[0])
            arms_tb = gr.Textbox(label="หรือพิมพ์เอง", placeholder="แขนออกรากไม้...")

            skin_dd = gr.Dropdown(skin_patterns_th, label="ผิว/ลวดลาย", value=skin_patterns_th[0])
            skin_tb = gr.Textbox(label="หรือพิมพ์เอง", placeholder="ลายทางสีทอง...")

        with gr.Column():
            ability_dd = gr.Dropdown(abilities_th, label="พลังพิเศษ", value=abilities_th[0])
            ability_tb = gr.Textbox(label="หรือพิมพ์เอง", placeholder="สร้างภาพลวงตา...")

            mood_dd = gr.Dropdown(moods_th, label="อารมณ์", value=moods_th[0])
            mood_tb = gr.Textbox(label="หรือพิมพ์เอง", placeholder="ขี้เล่นเป็นพิเศษ...")

    with gr.Row():
        # ปุ่มหลัก (สร้างมอนสเตอร์) -> อยู่ซ้าย
        create_btn = gr.Button("สร้างมอนสเตอร์!", elem_classes="btn-main")
        # ปุ่มรอง (สุ่ม) -> อยู่ขวา, ตัวสั้นลง
        random_btn = gr.Button("สุ่มมอนสเตอร์สุดเซอร์ไพรส์!", elem_classes="btn-random")

    monster_desc_th = gr.Textbox(label="รายละเอียด (ภาษาไทย)", interactive=False, elem_id="desc-th")
    monster_prompt_en = gr.Textbox(label="Prompt (English)", interactive=False, elem_id="prompt-en")

    gr.HTML(copy_button_html)

    # ปุ่มสุ่ม
    random_btn.click(
        fn=random_monster,
        inputs=[],
        outputs=[
            body_dd, body_tb,
            head_dd, head_tb,
            arms_dd, arms_tb,
            skin_dd, skin_tb,
            ability_dd, ability_tb,
            mood_dd, mood_tb
        ]
    )

    # ปุ่มสร้าง
    create_btn.click(
        fn=generate_monster_lab,
        inputs=[
            body_dd, body_tb,
            head_dd, head_tb,
            arms_dd, arms_tb,
            skin_dd, skin_tb,
            ability_dd, ability_tb,
            mood_dd, mood_tb
        ],
        outputs=[monster_desc_th, monster_prompt_en]
    )

    demo.load(fn=initial_text, inputs=None, outputs=monster_desc_th)

demo.launch()