import gradio as gr
from transformers import MarianMTModel, MarianTokenizer

# 1) โหลดโมเดลแปล (MarianMT) 
model_name = "Helsinki-NLP/opus-mt-th-en"
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name)

def translate_th_to_en(th_text: str) -> str:
    th_text = th_text.strip()
    if not th_text:
        return ""
    inputs = tokenizer(th_text, return_tensors="pt", max_length=512, truncation=True)
    translation_tokens = model.generate(**inputs, max_length=512)
    en_text = tokenizer.decode(translation_tokens[0], skip_special_tokens=True)
    return en_text

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

# 3) ฟังก์ชันสร้าง Prompt
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 = translate_th_to_en(body_th)
    head_en = translate_th_to_en(head_th)
    arms_en = translate_th_to_en(arms_th)
    skin_en = translate_th_to_en(skin_th)
    ability_en = translate_th_to_en(ability_th)
    mood_en = 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

# 4) สร้าง UI Gradio (ไม่มีปุ่ม Copy ใน .click() แต่ใช้ HTML ปุ่ม)
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;
}
.monster-btn {
    background-color: #FFC107;
    border: 2px solid #FFA000;
    font-weight: bold;
    font-size: 1.1rem;
    padding: 10px 20px;
    border-radius: 10px;
    margin-right: 10px;
}
.monster-btn:hover {
    background-color: #FFD54F;
}
#desc-th, #prompt-en {
    background-color: #FFFAE6;
    border: 2px solid #FFE082;
    border-radius: 10px;
    padding: 10px;
    margin-bottom: 20px;
}
"""

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

# สร้างปุ่ม Copy ด้วย HTML + JS
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() {
  // หยิบ Textbox 'prompt-en' (ภายใน Gradio) โดย id
  const promptBox = document.querySelector('#prompt-en textarea');
  if (!promptBox) {
    alert('ไม่พบข้อความ Prompt!');
    return;
  }
  // Copy
  const promptText = promptBox.value;
  navigator.clipboard.writeText(promptText);
  alert('คัดลอก Prompt แล้ว!');
}
</script>
"""

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

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

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

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

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

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

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

    create_btn = gr.Button("สร้างมอนสเตอร์!", elem_classes="monster-btn")

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

    # แทรกปุ่ม Copy ด้วย HTML
    gr.HTML(copy_button_html)

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