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

# =========================================================
# 基本設定
# =========================================================
st.set_page_config(
    page_title="LLM 智慧燈號 + RPi 5 監控",
    layout="wide"
)

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

# =========================================================
# 主畫面 - 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()

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

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

    targets = []
    if color == "red":
        targets = [red_led]
    elif color == "green":
        targets = [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()

    msg = "開啟" if action == "on" else "關閉"
    st.success(f"💡 {color} 燈已 {msg}")

# =========================================================
# 手動測試區
# =========================================================
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("綠燈已關閉")

# =========================================================
# Sidebar - Raspberry Pi 5 系統監控
# =========================================================
def run(cmd):
    return subprocess.check_output(cmd, shell=True, text=True).strip()

def cpu_usage():
    idle = run(
        "top -bn1 | grep 'Cpu(s)' | "
        "sed 's/,/ /g' | "
        "awk '{for(i=1;i<=NF;i++) if ($i==\"id\") print $(i-1)}'"
    )
    return 100 - float(idle)

def ram_usage():
    used, total, pct = run(
        "free -m | awk 'NR==2{print $3,$2,$3/$2*100}'"
    ).split()
    return int(used), int(total), float(pct)

def swap_usage():
    used, total, pct = run(
        "free -m | awk 'NR==3{print $3,$2,($3/$2)*100}'"
    ).split()
    return int(used), int(total), float(pct)

def disk_usage():
    used, total, pct = run(
        "df -h / | awk 'NR==2{print $3,$2,$5}'"
    ).split()
    return used, total, pct

def cpu_temp():
    return float(run("vcgencmd measure_temp | cut -d= -f2 | tr -d \"'C\""))

def net_usage():
    iface = run(
        "ip -o link show | awk -F': ' '{print $2}' | "
        "grep -E '^(eth|wlan|usb)' | head -n1"
    )
    rx, tx = run(
        f"cat /proc/net/dev | grep {iface} | awk '{{print $2,$10}}'"
    ).split()
    return iface, int(tx)//1024**2, int(rx)//1024**2

with st.sidebar:
    st.subheader("🍓 Raspberry Pi 5 監控")

    c1, c2 = st.columns(2)
    c3, c4 = st.columns(2)
    c5, c6 = st.columns(2)

    c1.metric("🧠 CPU", f"{cpu_usage():.0f}%")

    ram_used, ram_total, ram_pct = ram_usage()
    c2.metric("💾 RAM", f"{ram_pct:.0f}%")

    swap_used, swap_total, swap_pct = swap_usage()
    c3.metric("🌀 Swap", f"{swap_pct:.0f}%")

    disk_used, disk_total, disk_pct = disk_usage()
    c4.metric("📀 Disk", disk_pct)

    temp = cpu_temp()
    if temp >= 80:
        c5.error(f"🌡 {temp:.1f} °C\n⚠️ 過熱")
    else:
        c5.metric("🌡 Temp", f"{temp:.1f} °C")

    iface, tx, rx = net_usage()
    c6.metric("🌐 Net", iface, f"↑{tx}MB ↓{rx}MB")

    time.sleep(2)
    st.rerun()