import streamlit as st
import ollama
import json
from gpiozero import PWMOutputDevice
from time import sleep
import os

os.environ["OLLAMA_NUM_GPU"] = "1"

# =========================================================
# GPIO 初始化（只會執行一次）
# =========================================================
@st.cache_resource
def init_buzzer():
    return PWMOutputDevice(17)

buzzer = init_buzzer()

# =========================================================
# 音階表（白名單）
# =========================================================
NOTES = {
    "C": 262,
    "D": 294,
    "E": 330,
    "F": 349,
    "G": 392,
    "A": 440,
    "B": 494,
    "C5": 523
}

# =========================================================
# LLM 語意 → 旋律 JSON
# =========================================================
def analyze_music(text):
    prompt = f"""
你是一個「旋律生成器」。

請「只輸出 JSON」，不要解釋、不要多餘文字。
JSON 格式必須完全符合：

{{
  "melody": [
    {{"note":"C","duration":0.4}},
    {{"note":"E","duration":0.4}}
  ]
}}

規則：
- 只能使用音符：C D E F G A B C5
- duration 範圍：0.2 ~ 0.8
- 音符數量 8～16 個
- 不要輸出任何其他文字

使用者需求：
{text}
"""

    response = ollama.chat(
        model="llama3",
        messages=[
            {"role": "system", "content": "你只能輸出 JSON"},
            {"role": "user", "content": prompt}
        ]
    )

    raw = response["message"]["content"].strip()

    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        return None

# =========================================================
# 播放旋律（安全執行）
# =========================================================
def play_melody(melody):
    for item in melody:
        note = item.get("note")
        duration = float(item.get("duration", 0.4))

        if note not in NOTES:
            continue

        buzzer.frequency = NOTES[note]
        buzzer.value = 0.5
        sleep(duration)
        buzzer.off()

# =========================================================
# Streamlit UI
# =========================================================
st.title("🎼 LLM 自動作曲 → GPIO 彈奏")
st.write("輸入中文描述，LLM 生成旋律並由蜂鳴器播放")

cmd = st.text_input(
    "🗣️ 輸入需求（例如：快樂的旋律 / 悲傷慢歌）",
    "請生成一段快樂、輕快的旋律"
)

if st.button("🤖 LLM 作曲並播放"):
    if not cmd.strip():
        st.warning("請輸入描述")
        st.stop()

    result = analyze_music(cmd)

    if result is None:
        st.error("❌ LLM 回傳格式錯誤")
        st.stop()

    melody = result.get("melody", [])

    st.subheader("🎵 LLM 生成旋律（JSON）")
    st.json(result)

    play_melody(melody)

# =========================================================
# 安全停止
# =========================================================
st.divider()

if st.button("🔇 停止播放"):
    buzzer.off()