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

###################################
# 1) โหลดโมเดล MarianMT (Thai->En)
###################################
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)
###################################
def approximate_temp_with_star(star_type, distance_au, albedo=0.3):
    """
    - star_type: เพิ่ม White Dwarf, Supergiant
    - distance_au: ระยะห่าง (AU)
    - albedo: สัดส่วนสะท้อน (0.3)
    - กลไก: Stefan-Boltzmann + greenhouse + lum_ratio
    """
    if star_type == "Red Dwarf":
        lum_ratio = 0.02
    elif star_type == "White Dwarf":
        lum_ratio = 0.001  # สมมุติ (white dwarf มักมี luminosity ต่ำ)
    elif star_type == "Sun-like":
        lum_ratio = 1.0
    elif star_type == "Blue Giant":
        lum_ratio = 10.0
    elif star_type == "Supergiant":
        lum_ratio = 100.0
    else:
        lum_ratio = 1.0  # default

    luminosity = 3.828e26 * lum_ratio
    dist_m = distance_au * 1.496e11

    sigma = 5.67e-8
    T_k = ((1 - albedo) * luminosity / (16 * math.pi * sigma * dist_m**2))**0.25
    T_c = T_k - 273.15
    T_c += 15  # สมมุติ Greenhouse

    return round(T_c)

###################################
# 3) สูตรแรงโน้มถ่วงตามกฎนิวตัน (สมมุติ)
###################################
def approximate_gravity_nw(diameter_factor):
    """
    ถ้าคิดว่าดาวเคราะห์มีความหนาแน่นเท่าโลก:
    - มวล (M) ~ (diameter_factor^3) (เพราะ R^3 => Volume)
    - รัศมี (R) ~ diameter_factor
    - g ~ GM/R^2 ~ R^3 / R^2 = R (ถ้า density คงที่)
    ดังนั้น g ~ diameter_factor (เหมือนเดิม)
    แต่เราจะใส่ “สูตรจริง” ในปุ่มโชว์สูตร
    """
    return round(diameter_factor, 2)

###################################
# 4) ตีความเชิงบรรยาย (ภาษาไทย)
###################################
def describe_distance(distance_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):
    if temp_c < -30:
        return "อากาศหนาวแข็ง"
    elif temp_c < 10:
        return "เย็นสบาย"
    elif temp_c < 35:
        return "พอเหมาะกำลังดี"
    else:
        return "ร้อนระอุ"

def describe_gravity(g):
    if g < 0.5:
        return "โน้มถ่วงเบา (ลอยง่าย)"
    elif g < 1.2:
        return "คล้ายโลก"
    elif g < 2.0:
        return "หนักกว่าปกติ"
    else:
        return "หนักมาก"

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

def describe_moons(n):
    if n <= 0:
        return "ไม่มีดวงจันทร์"
    elif n == 1:
        return "มีดวงจันทร์หนึ่งดวง"
    else:
        # ถ้ามีเยอะ => tides effect
        if n > 2:
            return f"มีดวงจันทร์ {n} ดวง ซึ่งทำให้ปรากฏการณ์น้ำขึ้นน้ำลงรุนแรง"
        else:
            return f"มีดวงจันทร์ {n} ดวง"

###################################
# 5) สร้าง Prompt 3 แบบ (ไม่มีตัวเลข)
###################################
def build_prompts_en(planet_name_en, star_type_en, dist_desc_en, temp_desc_en, grav_desc_en, tilt_desc_en, moon_desc_en, oxygen_desc_en, life_en):
    prompt1 = (
        f"A vibrant space painting of planet '{planet_name_en}' orbiting a {star_type_en}. "
        f"It is {dist_desc_en}, with {temp_desc_en} conditions and {grav_desc_en} gravity. "
        f"{tilt_desc_en}, {moon_desc_en}, atmosphere has {oxygen_desc_en}. Cinematic details."
    )

    prompt2 = (
        f"On planet '{planet_name_en}', we discover {life_en} thriving in {temp_desc_en} weather, "
        f"{grav_desc_en} pull, and {oxygen_desc_en} in the air. Surreal alien ecosystem, rich concept art."
    )

    prompt3 = (
        f"Exploring the surface of '{planet_name_en}': {temp_desc_en} climate, {grav_desc_en}, "
        f"{tilt_desc_en} tilt, and {moon_desc_en}. Epic environment design, atmospheric perspective."
    )

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

###################################
# 6) ฟังก์ชันหลัก
###################################
def generate_planet_info(
    planet_name_th,
    star_type_en,       # Red Dwarf / White Dwarf / Sun-like / Blue Giant / Supergiant
    distance_str,
    diameter_str,
    tilt_value,
    moon_value,
    oxygen_percent,
    planet_type_th,     # (ดาวหิน, ดาวก๊าซ, ดาวน้ำแข็ง)
    life_th
):
    # parse
    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_value)
    except:
        tilt_deg = 23.5

    try:
        moon_count = int(moon_value)
    except:
        moon_count = 1

    # คำนวณอุณหภูมิ
    temp_c = approximate_temp_with_star(star_type_en, distance_au)
    # คำนวณแรงโน้มถ่วงแบบนิวตันสมมุติ (R^3 => M, g => M/R^2 => R)
    g_approx = approximate_gravity_nw(diameter_factor)

    # สรุปสำหรับเด็ก (ภาษาไทย)
    child_summary = (
        f"ดาว {planet_name_th} เป็น{planet_type_th} โคจรรอบดาวฤกษ์ {star_type_en}\n"
        f"ระยะ ~{distance_au} AU => อุณหภูมิราว {temp_c}°C,\n"
        f"แกนเอียง {tilt_deg}°, ดวงจันทร์ {moon_count}, "
        f"แรงโน้มถ่วง ~{g_approx}g, O2 ~{oxygen_percent}%\n"
        f"มีสิ่งมีชีวิต: {life_th}\n"
        f"น่าตื่นเต้นมาก!"
    )

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

    # แปลชื่อดาว, สิ่งมีชีวิต
    planet_name_en = translate_th_to_en(planet_name_th)
    life_en = translate_th_to_en(life_th)

    # ตีความ distance, temp, grav, tilt, moon, oxygen
    dist_desc_th = describe_distance(distance_au)
    dist_desc_en = translate_th_to_en(dist_desc_th)

    temp_desc_th = describe_temp(temp_c)
    temp_desc_en = translate_th_to_en(temp_desc_th)

    grav_desc_th = describe_gravity(g_approx)
    grav_desc_en = translate_th_to_en(grav_desc_th)

    tilt_desc_th = describe_tilt(tilt_deg)
    tilt_desc_en = translate_th_to_en(tilt_desc_th)

    moon_desc_th = describe_moons(moon_count)
    moon_desc_en = translate_th_to_en(moon_desc_th)

    # ตีความ O2
    if oxygen_percent < 1:
        o2_desc_th = "ไม่มีออกซิเจน"
    elif oxygen_percent < 10:
        o2_desc_th = "ออกซิเจนน้อย"
    elif oxygen_percent < 25:
        o2_desc_th = "ออกซิเจนพอเหมาะ"
    else:
        o2_desc_th = "ออกซิเจนสูง"
    oxygen_desc_en = translate_th_to_en(o2_desc_th)

    # สร้าง 3 Prompts
    prompts_en = build_prompts_en(
        planet_name_en,
        star_type_en + " star",
        dist_desc_en,
        temp_desc_en,
        grav_desc_en,
        tilt_desc_en,
        moon_desc_en,
        oxygen_desc_en,
        life_en
    )

    return child_summary, detail_th, prompts_en

###################################
# สูตร (Markdown) เพิ่มเนื้อหา
###################################
formula_text = r"""
**สูตรอุณหภูมิ (Stefan-Boltzmann)**

\\[
T = \left(\frac{(1 - A) \times L}{16 \pi \sigma \, d^2}\right)^{\frac{1}{4}} - 273.15 + 15^\circ\text{C (Greenhouse)}
\\]

- \\(A\\) = Albedo
- \\(L\\) = ความสว่างของดาว (W)
- \\(\sigma\\) = 5.67\\times10^{-8} (ค่าคงที่ Stefan-Boltzmann)
- \\(d\\) = ระยะทาง (m)

**สูตรแรงโน้มถ่วงนิวตัน**: \\(g = \frac{GM}{R^2}\\)

ถ้าสมมติความหนาแน่นเท่าโลก => \\(M \propto R^3\\) => \\(g \propto \frac{R^3}{R^2} = R\\)

*(ในโค้ดเราประยุกต์ใช้ค่าสัดส่วนง่าย ๆ)*
"""

###################################
# 7) สร้าง 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, #formula-box {
    background-color: #FFFDF5;
    border: 2px solid #FFE082;
    border-radius: 10px;
    padding: 10px;
    margin-bottom: 20px;
}
"""

def show_formula(state):
    new_state = not state
    return new_state, gr.update(visible=new_state)

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

with gr.Blocks(css=css_code) as demo:
    gr.Markdown("<h1 id='title'>ZenityX Planetary Adventure</h1>")
    gr.Markdown("""
    <div class="game-desc">
        <p>เพิ่มประเภทดาวฤกษ์ (White Dwarf, Supergiant), ใส่เงื่อนไขน้ำขึ้นน้ำลง (tides) ถ้าดวงจันทร์เยอะ, 
           และใช้สูตรแรงโน้มถ่วงนิวตัน (สมมุติ).</p>
        <p>กรอกข้อมูลแล้วกด <strong>"สร้างโลกแฟนตาซี"</strong>!</p>
        <p>หากอยากดู <strong>สูตรคำนวณ</strong> ให้กดปุ่ม "โชว์สูตรคำนวณ" ด้านล่าง</p>
    </div>
    """)

    formula_state = gr.State(value=False)
    formula_md = gr.Markdown(formula_text, visible=False, elem_id="formula-box")

    show_formula_btn = gr.Button("โชว์สูตรคำนวณ")
    show_formula_btn.click(fn=show_formula, inputs=formula_state, outputs=[formula_state, formula_md])

    with gr.Row():
        with gr.Column():
            planet_name_th = gr.Textbox(label="ชื่อดาวเคราะห์ (ไทย)", placeholder="เช่น: ดาวซานาดา")
            star_type_en = gr.Dropdown(
                label="ประเภทดาวฤกษ์",
                choices=["Red Dwarf", "White Dwarf", "Sun-like", "Blue Giant", "Supergiant"],
                value="Sun-like"
            )
            distance_au = gr.Textbox(label="ระยะห่าง (AU)", placeholder="1, 0.5, 2...")
            diameter_factor = gr.Textbox(label="ขนาด (เท่าโลก)", placeholder="1, 2, 0.5...")

        with gr.Column():
            tilt_slider = gr.Slider(0, 90, step=1, value=23.5, label="แกนเอียง (องศา)")
            moon_slider = gr.Slider(0, 10, step=1, value=1, label="จำนวนดวงจันทร์")
            oxygen_slider = gr.Slider(0, 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="Prompts (English)", interactive=False, elem_id="prompt-en")

    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>
    """
    gr.HTML(copy_button_html)

    def generate_wrapper(
        p_name_th, s_type_en, dist_au, dia_fac, tilt_val, moon_val, oxy_val, p_type_th, l_th
    ):
        return generate_planet_info(
            planet_name_th=p_name_th,
            star_type_en=s_type_en,
            distance_str=dist_au,
            diameter_str=dia_fac,
            tilt_value=str(tilt_val),
            moon_value=str(moon_val),
            oxygen_percent=oxy_val,
            planet_type_th=p_type_th,
            life_th=l_th
        )

    create_btn.click(
        fn=generate_wrapper,
        inputs=[
            planet_name_th,
            star_type_en,
            distance_au,
            diameter_factor,
            tilt_slider,
            moon_slider,
            oxygen_slider,
            planet_type_th,
            life_th
        ],
        outputs=[
            child_summary_out,
            detail_th_out,
            prompt_en_out
        ]
    )

    def welcome_text():
        return "ยินดีต้อนรับสู่ Planetary Adventure++! ลองกรอกข้อมูลแล้วกด 'สร้างโลกแฟนตาซี' สิ!"
    
    demo.load(fn=welcome_text, inputs=None, outputs=child_summary_out)

demo.launch()