import socket, json, threading, webbrowser, sys, os, tkinter as tk
import pystray                         # pip install pystray pillow
from pystray import MenuItem as item
from PIL import Image, ImageDraw, ImageTk

UDP_PORT     = 9876
opened_once  = False
WIN_W, WIN_H = 400, 200
ICON_FILE    = "logo.ico"              # ← 你的图标文件名（ico/png 都行）

# ───── 通用资源路径（兼容 PyInstaller） ───────────────────────────
def resource_path(rel_path: str):
    """Return absolute path to resource, works for dev & PyInstaller"""
    base = getattr(sys, "_MEIPASS", os.path.abspath("."))
    return os.path.join(base, rel_path)

# ───── 收到广播就跳浏览器 ───────────────────────────────────────────
def open_browser(url):
    try:
        webbrowser.open(url)
    except Exception as e:
        print(f"Failed to open browser: {e}")

def listen_udp():
    global opened_once
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind(("", UDP_PORT))
    print(f"Listening for UDP broadcasts on port {UDP_PORT}...")

    while True:
        try:
            data, _ = sock.recvfrom(1024)
            msg = json.loads(data.decode("utf-8"))
            if msg.get("type") != "ai_server_announce":
                continue

            ip, port, name = msg.get("ip"), msg.get("port"), msg.get("name", "Unknown")
            url = f"http://{ip}:80"
            print(f"Discovered {name} at {url}")

            if not opened_once:
                opened_once = True
                open_browser(url)
        except Exception as e:
            print(f"UDP error: {e}")

# ───── 托盘辅助 ───────────────────────────────────────────────────
def load_icon_img():
    """Load ICO/PNG; fall back to blue square if文件缺失/损坏"""
    try:
        return Image.open(resource_path(ICON_FILE))
    except Exception as e:
        print(f"[warn] load icon failed: {e}, fallback to default")
        return create_fallback_icon()

def create_fallback_icon(size=(64, 64), bg="#004080", fg="#ffffff"):
    img = Image.new("RGB", size, bg)
    ImageDraw.Draw(img).rectangle([16,16,48,48], fill=fg)
    return img

def show_window(icon=None, _=None):
    root.deiconify(); root.focus_force()
    if icon: icon.visible = False

def hide_to_tray():
    root.withdraw(); tray_icon.visible = True

def on_exit(icon=None, _=None):
    if icon: icon.stop()
    root.quit(); sys.exit()

tray_icon = pystray.Icon(
    "ai_discovery",
    icon  = load_icon_img(),           # ⭐ 自定义托盘图标
    title = "PlugAI Launcher",
    menu  = pystray.Menu(
        item("显示窗口", show_window),
        item("退出程序", on_exit)
    )
)

# ───── 主窗口 ─────────────────────────────────────────────────────
root = tk.Tk()
root.title("PlugAI Launcher")

# ⭐ 窗口标题栏图标（ICO 必须；PNG 需转 PhotoImage）
ico_path = resource_path(ICON_FILE)
if ico_path.lower().endswith(".ico"):
    try:
        root.iconbitmap(ico_path)
    except Exception as e:
        print(f"[warn] set window icon failed: {e}")
else:  # PNG fallback
    try:
        img = ImageTk.PhotoImage(file=ico_path)
        root.iconphoto(False, img)
    except Exception as e:
        print(f"[warn] set window icon failed: {e}")

def center_window():
    sw, sh = root.winfo_screenwidth(), root.winfo_screenheight()
    x, y = (sw - WIN_W) // 2, (sh - WIN_H) // 2
    root.geometry(f"{WIN_W}x{WIN_H}+{x}+{y}")

center_window()
tk.Label(root, text="正在等待 PlugAI 准备好中...\n端口: 9876",
         font=("Arial", 12)).pack(pady=60)

root.protocol("WM_DELETE_WINDOW", on_exit)

def on_minimize(event):
    if root.state() == "iconic":
        hide_to_tray()

root.bind("<Unmap>", on_minimize)

# ───── 启动后台线程 ───────────────────────────────────────────────
threading.Thread(target=listen_udp, daemon=True).start()
threading.Thread(target=tray_icon.run, daemon=True).start()

root.mainloop()