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

##############################################
# 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(text_th: str) -> str:
    text_th = text_th.strip()
    if not text_th:
        return ""
    inputs = tokenizer(text_th, 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) สูตรคำนวณ Black Body (Stefan-Boltzmann) + Greenhouse เบื้องต้น
##############################################

def approximate_temp_with_star(star_type, distance_au, albedo=0.3):
    """
    คำนวณอุณหภูมิพื้นผิว (°C) อย่างง่าย โดยอิง Black Body + star luminosity
    - star_type:
      "Red Dwarf" -> luminosity ~0.02 เท่าดวงอาทิตย์
      "Sun-like"  -> luminosity = 1 เท่าดวงอาทิตย์
      "Blue Giant"-> luminosity ~10 เท่าดวงอาทิตย์ (สมมุติ)
    - distance_au: ระยะห่างหน่วย AU
    - albedo: สัดส่วนสะท้อนแสง (สมมุติ 0.3)
    - สุดท้ายบวก greenhouse effect ~ +15 °C (สมมุติ) 
    """
    if star_type == "Red Dwarf":
        lum_ratio = 0.02  # สมมุติ
    elif star_type == "Blue Giant":
        lum_ratio = 10.0
    else:  # Sun-like
        lum_ratio = 1.0
    
    # ความสว่างของดวงอาทิตย์ ~3.828e26 W
    luminosity = 3.828e26 * lum_ratio

    # แปลง AU -> เมตร
    dist_m = distance_au * 1.496e11

    sigma = 5.67e-8  # Stefan-Boltzmann constant
    # T(K) = [ ( (1 - A) * L ) / (16*pi*sigma * d^2 ) ]^(1/4)
    T_k = ((1 - albedo) * luminosity / (16 * math.pi * sigma * dist_m**2))**0.25
    T_c = T_k - 273.15

    # บวก greenhouse เล็กน้อย
    T_c += 15

    return round(T_c)

def approximate_gravity(diameter_factor):
    """
    สมมุติว่า แรงโน้มถ่วง ~ diameter_factor เท่าของโลก
    (จริงๆ เกี่ยวกับความหนาแน่นด้วย แต่เอาแบบง่าย)
    """
    return round(diameter_factor, 2)


##############################################
# 3) ฟังก์ชันแปลงตัวเลข -> คำบรรยายเชิงเปรียบเทียบ
##############################################

def describe_distance(distance_au):
    """
    ตีความระยะห่าง (AU) เป็นคำบรรยาย
    """
    if distance_au < 0.5:
        return "โคจรอยู่ใกล้มากจนดาวเคราะห์ร้อนระอุ"
    elif distance_au < 1.2:
        return "โคจรค่อนข้างใกล้เหมือนโลก หรืออาจอุ่นกว่านิดหน่อย"
    elif distance_au < 2.5:
        return "โคจรห่างพอประมาณ ค่อนข้างเย็น"
    else:
        return "โคจรไกลมาก มีสภาพหนาวเย็นสุดขั้ว"

def describe_temp(temp_c):
    """
    ตีความอุณหภูมิ (°C) เป็นคำบรรยาย
    """
    if temp_c < -30:
        return "อากาศหนาวแข็ง"
    elif temp_c < 10:
        return "อากาศค่อนข้างเย็น"
    elif temp_c < 35:
        return "อากาศกำลังสบาย"
    else:
        return "ร้อนระอุ"

def describe_gravity(g):
    """
    ตีความแรงโน้มถ่วง (g) เป็นคำบรรยาย
    """
    if g < 0.5:
        return "โน้มถ่วงเบามาก (ตัวลอยได้ง่าย)"
    elif g < 1.2:
        return "คล้ายโลก"
    elif g < 2.0:
        return "ค่อนข้างหนัก"
    else:
        return "หนักมากจนอาจกดทับทุกสิ่ง"

def describe_tilt(tilt_deg):
    """
    ตีความแกนเอียง
    """
    tilt = float(tilt_deg)
    if tilt < 5:
        return "แทบไม่เอียง ฤดูกาลเรียบง่าย"
    elif tilt < 25:
        return "เอียงเล็กน้อย มีฤดูกาลพอประมาณ"
    elif tilt < 45:
        return "เอียงปานกลาง ฤดูกาลค่อนข้างเปลี่ยนแปลง"
    else:
        return "เอียงมาก ฤดูกาลสุดขั้ว"

def describe_moons(n):
    """
    ตีความจำนวนดวงจันทร์
    """
    n = int(n)
    if n == 0:
        return "ไม่มีดวงจันทร์"
    elif n == 1:
        return "มีดวงจันทร์หนึ่งดวง"
    else:
        return f"มีดวงจันทร์ {n} ดวง"


##############################################
# 4) สร้าง prompt 3 แบบ (ไม่มีตัวเลขตรงๆ)
##############################################

def build_prompts_en(planet_name_en, star_type_en, distance_desc_en, 
                     temp_desc_en, gravity_desc_en, tilt_desc_en, 
                     moon_desc_en, oxygen_desc_en, life_en):
    """
    สร้าง 3 prompts ภาษาอังกฤษ โดยไม่ใส่ตัวเลขตรงๆ 
    แต่ใช้คำบรรยายเชิงเปรียบเทียบ
    """

    # Prompt 1: ภาพดาวจากอวกาศ
    prompt1 = (
        f"A stunning space illustration of planet '{planet_name_en}' orbiting a {star_type_en} star. "
        f"It is {distance_desc_en}, with {temp_desc_en} climate and {gravity_desc_en} surface gravity. "
        f"{tilt_desc_en}, and {moon_desc_en}. Atmosphere containing {oxygen_desc_en}, truly fascinating."
    )

    # Prompt 2: สิ่งมีชีวิต
    prompt2 = (
        f"Alien life on planet '{planet_name_en}': {life_en}, "
        f"thriving under {temp_desc_en} conditions, {gravity_desc_en} gravity, and {oxygen_desc_en} in the air. "
        f"Highly imaginative, concept art style, fantasy world."
    )

    # Prompt 3: พื้นผิวดาว
    prompt3 = (
        f"The planet '{planet_name_en}' has a surface with {temp_desc_en} weather, {gravity_desc_en}, "
        f"{tilt_desc_en}, and {moon_desc_en}. An epic terrain, atmospheric glow, and {oxygen_desc_en} presence. "
        f"Realistic space painting, cinematic lighting."
    )

    prompt_all = (
        "--- Prompt #1 ---\n" + prompt1 + "\n\n"
        "--- Prompt #2 ---\n" + prompt2 + "\n\n"
        "--- Prompt #3 ---\n" + prompt3 + "\n"
    )
    return prompt_all


##############################################
# 5) ฟังก์ชันหลัก generate_planet_info
##############################################
def generate_planet_info(
    planet_name_th,
    star_type_en_select,  # ("Red Dwarf", "Sun-like", "Blue Giant")
    distance_str,
    diameter_str,
    tilt_str,
    moon_str,
    oxygen_percent,
    planet_type_th,       # ("ดาวหิน", "ดาวก๊าซ", "ดาวน้ำแข็ง")
    life_th
):
    # แปลง string -> float/int
    try:
        distance_au = float(distance_str)
    except:
        distance_au = 1.0

    try:
        diameter_factor = float(diameter_str)
    except:
        diameter_factor = 1.0

    try:
        tilt_deg = float(tilt_str)
    except:
        tilt_deg = 23.5

    try:
        num_moon = int(moon_str)
    except:
        num_moon = 1

    # 1) คำนวณอุณหภูมิ (°C) 
    temp_c = approximate_temp_with_star(star_type_en_select, distance_au)
    # 2) คำนวณแรงโน้มถ่วง
    g_approx = approximate_gravity(diameter_factor)

    # 3) สร้าง “สรุปสำหรับเด็ก” (ภาษาไทย)
    #    โดยยังไม่ได้รวม star_type_en_select เพราะเป็นภาษาอังกฤษ
    #    แต่เราจะใช้ planet_type_th ด้วย
    #    สมมุติ child_friendly_summary: 
    #       => โชว์ temp, g, oxygen, life
    child_summary = (
        f"ดาว {planet_name_th} เป็น{planet_type_th} " 
        f"โคจรรอบดาวฤกษ์ประเภท {star_type_en_select} (อิงวิทยาศาสตร์สมมุติ)\n"
        f"อุณหภูมิประมาณ {temp_c}°C, "
        f"แรงโน้มถ่วง ~{g_approx}g, ออกซิเจน ~{oxygen_percent}%, "
        f"แกนเอียง {tilt_deg}°, {num_moon} ดวงจันทร์,\n"
        f"มีสิ่งมีชีวิต: {life_th}\n"
        f"น่าอัศจรรย์จริง ๆ!"
    )

    # 4) รายละเอียดเชิงเทคนิค (ไทย)
    detail_th = (
        f"ชื่อดาวเคราะห์ (ไทย): {planet_name_th}\n"
        f"ประเภทดาวฤกษ์: {star_type_en_select}\n"
        f"ระยะห่าง: ~{distance_au} AU\n"
        f"ขนาดเส้นผ่านศูนย์กลาง: ~{diameter_factor} เท่าโลก\n"
        f"แกนเอียง: ~{tilt_deg}°\n"
        f"จำนวนดวงจันทร์: {num_moon}\n"
        f"อุณหภูมิคำนวณ: ~{temp_c} °C\n"
        f"แรงโน้มถ่วง: ~{g_approx} g\n"
        f"ออกซิเจน: {oxygen_percent}%\n"
        f"ชนิดดาวเคราะห์ (ไทย): {planet_type_th}\n"
        f"สิ่งมีชีวิต (ไทย): {life_th}\n"
    )

    # 5) แปล planet_name_th -> English
    planet_name_en = translate_th_to_en(planet_name_th)
    # แปล life_th -> อังกฤษ
    life_en = translate_th_to_en(life_th)

    # แมป planet_type_th -> Eng 
    # (ดาวหิน -> rocky planet, ดาวก๊าซ -> gas giant, ดาวน้ำแข็ง -> icy planet)
    type_map = {
        "ดาวหิน": "rocky planet",
        "ดาวก๊าซ": "gas giant",
        "ดาวน้ำแข็ง": "icy planet"
    }
    # (กรณี user พิมพ์ว่า "ดาวหิน" หรือมีวงเล็บ อาจต้อง .replace ให้เคลียร์)
    planet_type_key = planet_type_th.replace("(", "").replace(")", "").replace(" ", "").split("ดาว")[-1] 
    planet_type_key = planet_type_key.strip()
    planet_type_en = type_map.get(planet_type_key, "unknown planet")

    # 6) ตีความ distance, temp, gravity, tilt, moon, oxygen เพื่ออธิบายเป็นอังกฤษ (non-numeric)
    #   - distance
    distance_desc_th = describe_distance(distance_au)
    distance_desc_en = translate_th_to_en(distance_desc_th)
    #   - temp
    temp_desc_th = describe_temp(temp_c)
    temp_desc_en = translate_th_to_en(temp_desc_th)
    #   - gravity
    gravity_desc_th = describe_gravity(g_approx)
    gravity_desc_en = translate_th_to_en(gravity_desc_th)
    #   - tilt
    tilt_desc_th = describe_tilt(tilt_deg)
    tilt_desc_en = translate_th_to_en(tilt_desc_th)
    #   - moon
    moon_desc_th = describe_moons(num_moon)
    moon_desc_en = translate_th_to_en(moon_desc_th)
    #   - oxygen
    #   ถ้า oxygen 0 -> no oxygen
    #      oxygen < 5 -> extremely low oxygen
    #      etc.
    if oxygen_percent < 1:
        oxygen_desc_th = "ไม่มีออกซิเจน"
    elif oxygen_percent < 10:
        oxygen_desc_th = "ออกซิเจนจางมาก"
    elif oxygen_percent < 25:
        oxygen_desc_th = "ออกซิเจนพอเหมาะ"
    else:
        oxygen_desc_th = "ออกซิเจนสูง"
    oxygen_desc_en = translate_th_to_en(oxygen_desc_th)

    # 7) สร้าง 3 Prompts (English)
    prompt_all = build_prompts_en(
        planet_name_en,
        f"{star_type_en_select} star",  # e.g. "Red Dwarf star"
        distance_desc_en,
        temp_desc_en,
        gravity_desc_en,
        tilt_desc_en,
        moon_desc_en,
        oxygen_desc_en,
        life_en
    )

    return child_summary, detail_th, prompt_all


##############################################
# 6) สร้าง UI (Gradio)
##############################################

css_code = """
body {
    background-color: #F9FBFF;
    font-family: "Kanit", sans-serif;
}
#title {
    color: #4A90E2;
    text-align: center;
    font-size: 2rem;
    margin-top: 20px;
    margin-bottom: 10px;
    font-weight: bold;
}
.game-desc {
    margin: 0 auto;
    width: 90%;
    background-color: #ECF6FF;
    border: 2px dashed #B3DAFF;
    border-radius: 10px;
    padding: 15px;
    color: #333;
    margin-bottom: 20px;
}
.btn-main {
    background-color: #FFE066;
    border: 2px solid #FFCA28;
    font-weight: bold;
    font-size: 1.1rem;
    padding: 10px 30px;
    border-radius: 10px;
    margin-right: 10px;
}
#child-summary, #detail-th, #prompt-en {
    background-color: #FFFDF5;
    border: 2px solid #FFE082;
    border-radius: 10px;
    padding: 10px;
    margin-bottom: 20px;
}
"""

def welcome_text():
    return "ยินดีต้อนรับสู่ Planetary Adventure+! ลองกรอกข้อมูลแล้วกด 'สร้างโลกแฟนตาซี' กันนะจ๊ะ"

copy_button_html = """
<button style="background-color: #F06292; border: 2px solid #E91E63; font-weight: bold; 
               font-size: 1.0rem; padding: 8px 20px; border-radius: 10px;"
        onclick="copyPromptText()">
  คัดลอก 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 แล้ว! นำไปใช้ใน AI สร้างภาพได้เลย~');
}
</script>
"""

import re

with gr.Blocks(css=css_code) as demo:
    gr.Markdown("<h1 id='title'>Planetary Adventure (Advanced) - Thai Input + Real Sci + No Direct Numbers</h1>")
    gr.Markdown("""
    <div class="game-desc">
        <p>สร้างดาวเคราะห์พร้อมหลักวิทยาศาสตร์เบื้องต้น (สูตร Stefan-Boltzmann) และใช้การแปล MarianMT</p>
        <ol>
            <li>กรอกชื่อดาวเคราะห์ (ภาษาไทย)</li>
            <li>เลือกชนิดดาวฤกษ์ (Red Dwarf / Sun-like / Blue Giant)</li>
            <li>กรอกระยะห่าง (หน่วย AU), ขนาดเท่าโลก, แกนเอียง, จำนวนดวงจันทร์, % ออกซิเจน</li>
            <li>ระบุประเภทดาวเคราะห์ (ไทย) เช่น "ดาวหิน", "ดาวก๊าซ", "ดาวน้ำแข็ง"</li>
            <li>พิมพ์สิ่งมีชีวิต (ไทย)</li>
            <li>กดปุ่ม <strong>"สร้างโลกแฟนตาซี"</strong></li>
        </ol>
        <p>ระบบจะ <em>คำนวณอุณหภูมิ</em>, <em>แรงโน้มถ่วง</em>, <em>ตีความเป็นคำบรรยาย</em>, และแปลไทย->อังกฤษ<br>
           จากนั้นสร้าง <strong>3 prompts</strong> ที่เลี่ยงตัวเลขตรงๆ ใช้คำเปรียบเทียบแทน</p>
    </div>
    """)

    planet_name_th = gr.Textbox(label="ชื่อดาวเคราะห์ (ภาษาไทย)", placeholder="ตัวอย่าง: ดาวซานาดา")

    star_type_select = gr.Dropdown(
        label="ประเภทดาวฤกษ์ (English)",
        choices=["Red Dwarf", "Sun-like", "Blue Giant"],
        value="Sun-like"
    )

    distance_au = gr.Textbox(label="ระยะห่างจากดาวฤกษ์ (AU)", placeholder="เช่น 1, 0.5, 2")
    diameter_factor = gr.Textbox(label="ขนาด (เท่าโลก)", placeholder="เช่น 1, 2, 0.5")
    tilt_deg = gr.Textbox(label="แกนเอียง (องศา)", placeholder="0-90")
    moon_number = gr.Textbox(label="จำนวนดวงจันทร์", placeholder="0, 1, 2, ...")

    oxygen_slider = gr.Slider(
        minimum=0, maximum=100, step=1, value=21,
        label="% ออกซิเจน (ในบรรยากาศ)"
    )

    planet_type_th = gr.Dropdown(
        label="ชนิดดาวเคราะห์ (ไทย)",
        choices=["ดาวหิน", "ดาวก๊าซ", "ดาวน้ำแข็ง"],
        value="ดาวหิน"
    )

    life_th = gr.Textbox(label="สิ่งมีชีวิต (ภาษาไทย)", placeholder="ตัวอย่าง: แมลงยักษ์เรืองแสง")

    create_btn = gr.Button("สร้างโลกแฟนตาซี", elem_classes="btn-main")

    child_summary_out = gr.Textbox(label="สรุปสำหรับเด็ก (ไทย)", interactive=False, elem_id="child-summary")
    detail_th_out = gr.Textbox(label="รายละเอียด (ไทย)", interactive=False, elem_id="detail-th")
    prompt_en_out = gr.Textbox(label="3 Prompts (English)", interactive=False, elem_id="prompt-en")

    gr.HTML(copy_button_html)

    create_btn.click(
        fn=generate_planet_info,
        inputs=[
            planet_name_th,     # text
            star_type_select,   # red dwarf / sun-like / blue giant
            distance_au,
            diameter_factor,
            tilt_deg,
            moon_number,
            oxygen_slider,
            planet_type_th,     # "ดาวหิน", "ดาวก๊าซ", "ดาวน้ำแข็ง"
            life_th
        ],
        outputs=[
            child_summary_out,
            detail_th_out,
            prompt_en_out
        ]
    )

    demo.load(fn=welcome_text, inputs=None, outputs=child_summary_out)

demo.launch()