import streamlit as st
import ollama
import json
from gpiozero import LED
import os

# =========================================================
# 環境設定
# =========================================================
os.environ["OLLAMA_NUM_GPU"] = "1"

# =========================================================
# GPIO 初始化（只會執行一次）
# =========================================================
@st.cache_resource
def init_leds():
    red = LED(18)
    green = LED(23)
    return red, green

red_led, green_led = init_leds()

# =========================================================
# LLM 預熱（第一次呼叫較慢）
# =========================================================
@st.cache_resource
def warmup_llm():
    try:
        ollama.chat(
            model="llama3.2:3b",
            messages=[{"role": "user", "content": "hi"}],
            options={"num_predict": 1}
        )
    except:
        pass

warmup_llm()

# =========================================================
# LLM 語意分析（快速 JSON）
# =========================================================
def analyze_command(text):
    response = ollama.chat(
        model="llama3.2:3b",
        messages=[
            {
                "role": "system",
                "content": (
                    "只輸出 JSON，必須是以下之一："
                    '{"color":"red","action":"on"},'
                    '{"color":"red","action":"off"},'
                    '{"color":"green","action":"on"},'
                    '{"color":"green","action":"off"},'
                    '{"color":"all","action":"on"},'
                    '{"color":"all","action":"off"}'
                )
            },
            {"role": "user", "content": text}
        ],
        options={
            "temperature": 0,
            "num_predict": 20
        }
    )

    raw = response["message"]["content"].strip()
    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        return None

# =========================================================
# Streamlit UI - LLM 控制
# =========================================================
st.title("💡 LLM 智慧燈號控制系統")
st.write("請輸入中文指令，例如：**請亮綠燈 / 把紅燈關掉 / 全部開 / 全部關**")

cmd = st.text_input("🗣️ 輸入指令")

if st.button("送出指令"):
    if not cmd.strip():
        st.warning("請先輸入指令")
        st.stop()

    result = analyze_command(cmd)

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

    color = result.get("color")
    action = result.get("action")

    st.write("🤖 LLM JSON 回傳：", result)

    # =====================================================
    # 燈號控制（增加全部燈選項）
    # =====================================================
    targets = []
    if color == "red":
        targets.append(red_led)
    elif color == "green":
        targets.append(green_led)
    elif color == "all":
        targets = [red_led, green_led]

    if not targets:
        st.error("❌ 無法識別的控制指令")
        st.stop()

    for led in targets:
        if action == "on":
            led.on()
        elif action == "off":
            led.off()

    if color == "all":
        st.success(f"💡 全部燈已 {'開啟' if action=='on' else '關閉'}")
    else:
        st.success(f"💡 {color}燈已 {'開啟' if action=='on' else '關閉'}")

# =========================================================
# 手動測試區（除錯用）
# =========================================================
st.divider()
st.subheader("🧪 手動控制（測試用）")

c1, c2 = st.columns(2)

with c1:
    if st.button("🔴 紅燈 ON"):
        red_led.on()
        st.success("紅燈已開啟")
    if st.button("🔴 紅燈 OFF"):
        red_led.off()
        st.success("紅燈已關閉")

with c2:
    if st.button("🟢 綠燈 ON"):
        green_led.on()
        st.success("綠燈已開啟")
    if st.button("🟢 綠燈 OFF"):
        green_led.off()
        st.success("綠燈已關閉")

st.write("💡 也可以在 LLM 指令中輸入 **全部開 / 全部關** 控制所有燈")