import subprocess
import time
import dns.resolver
import re
import os
import psutil
import configparser
import threading
import tkinter as tk
from tkinter import ttk, scrolledtext
from PIL import Image, ImageDraw, ImageFont
import pystray

# ==========================
# 配置读取
# ==========================
config = configparser.ConfigParser()
config.read("config.ini", encoding="utf-8")

WIREGUARD_EXE = config["wireguard"]["wireguard_exe"]
CONF_FILE = config["wireguard"]["conf_file"]
CHECK_IP = config["wireguard"]["check_ip"]

TXT_DOMAIN = config["dns"]["txt_domain"]
DNS_SERVERS = [x.strip() for x in config["dns"]["dns_server"].split(",") if x.strip()]
DNS_LIFETIME = int(config["dns"]["lifetime"])

CHECK_INTERVAL = int(config["check"]["interval"])
PING_TIMEOUT = config["check"]["ping_timeout_ms"]

ENDPOINT_PATTERN = r"(\d{1,3}(?:\.\d{1,3}){3}:\d+)"

resolver = dns.resolver.Resolver()
resolver.nameservers = DNS_SERVERS
resolver.lifetime = DNS_LIFETIME

last_endpoint = None
loop_running = False
loop_thread = None
tray_icon = None

# ==========================
# GUI 日志输出
# ==========================
def log(msg):
    timestamp = time.strftime("%H:%M:%S")
    text_area.configure(state="normal")
    text_area.insert(tk.END, f"[{timestamp}] {msg}\n")
    text_area.see(tk.END)
    text_area.configure(state="disabled")
    print(msg)

# ==========================
# WireGuard 功能
# ==========================
def get_endpoint_from_txt():
    try:
        answers = resolver.resolve(TXT_DOMAIN, "TXT", tcp=True)
        for r in answers:
            if hasattr(r, "strings"):
                txt = "".join([s.decode(errors="ignore") for s in r.strings])
            else:
                txt = r.to_text().strip('"')
            match = re.search(ENDPOINT_PATTERN, txt)
            if match:
                return match.group(1)
    except Exception as e:
        log(f"TXT query failed: {e}")
    return None

def update_conf_endpoint(new_endpoint):
    lines = []
    with open(CONF_FILE, "r", encoding="utf-8") as f:
        for line in f:
            if line.strip().lower().startswith("endpoint"):
                lines.append(f"Endpoint = {new_endpoint}\n")
            else:
                lines.append(line)
    with open(CONF_FILE, "w", encoding="utf-8") as f:
        f.writelines(lines)

def run_wg(args):
    result = subprocess.run(
        [WIREGUARD_EXE] + args,
        capture_output=True,
        creationflags=subprocess.CREATE_NO_WINDOW  # ✅ 防止弹出窗口
    )
    if result.returncode != 0:
        log(result.stderr.decode(errors="ignore"))
    return result

def uninstall_tunnel():
    killed = 0
    for proc in psutil.process_iter(["name", "pid"]):
        try:
            if proc.info["name"] and proc.info["name"].lower() == "wireguard.exe":
                proc.kill()
                killed += 1
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            continue
    log(f"{killed} WireGuard process(es) terminated.")

def install_tunnel():
    log("Installing tunnelservice...")
    run_wg(["/installtunnelservice", CONF_FILE])

def is_tunnel_up():
    try:
        result = subprocess.run(
            ["ping", "-n", "1", "-w", str(PING_TIMEOUT), CHECK_IP],
            capture_output=True,
            creationflags=subprocess.CREATE_NO_WINDOW  # ✅ 防止闪现窗口
        )
        return b"TTL=" in result.stdout
    except Exception as e:
        log(f"Ping check failed: {e}")
        return False

def ensure_tunnel(new_endpoint):
    log(f"Updating tunnel -> {new_endpoint}")
    uninstall_tunnel()
    update_conf_endpoint(new_endpoint)
    install_tunnel()
    log("Tunnel updated")

# ==========================
# 主循环
# ==========================
def wg_loop():
    global last_endpoint, loop_running
    status_label.config(text="已连接")
    while loop_running:
        endpoint = get_endpoint_from_txt()
        if endpoint and endpoint != last_endpoint:
            log(f"New endpoint detected: {endpoint}")
            ensure_tunnel(endpoint)
            last_endpoint = endpoint
        elif endpoint and not is_tunnel_up():
            log("Tunnel down, restarting")
            ensure_tunnel(endpoint)
        time.sleep(CHECK_INTERVAL)
    uninstall_tunnel()
    status_label.config(text="未连接")
    log("Tunnel stopped")
    toggle_button.config(state="normal", text="启动")

# ==========================
# GUI 控制
# ==========================
def start_stop_loop():
    global loop_running, loop_thread
    if not loop_running:
        loop_running = True
        loop_thread = threading.Thread(target=wg_loop, daemon=True)
        loop_thread.start()
        toggle_button.config(text="停止")
    else:
        toggle_button.config(state="disabled")
        loop_running = False

# ==========================
# 托盘支持
# ==========================
def make_icon():
    size = 64
    image = Image.new("RGBA", (size, size), (0, 0, 0, 0))  # 透明背景
    draw = ImageDraw.Draw(image)

    # 画圆形背景
    draw.ellipse((0, 0, size-1, size-1), fill=(0, 120, 215, 255))  # Windows 蓝

    # 画字母 "WG"
    try:
        font = ImageFont.truetype("seguiemj.ttf", 28)  # 系统字体
    except:
        font = ImageFont.load_default()

    text = "WG"
    bbox = draw.textbbox((0, 0), text, font=font)  # 获取文字边界
    w = bbox[2] - bbox[0]
    h = bbox[3] - bbox[1]
    draw.text(((size - w) / 2, (size - h) / 2), text, font=font, fill="white")

    return image

def create_tray_icon_once():
    global tray_icon
    if tray_icon is None:
    
        def on_show(icon, item):
            icon.stop()
            root.after(0, root.deiconify)

        def on_quit(icon, item):
            icon.stop()
            global loop_running
            loop_running = False
            uninstall_tunnel()
            root.after(0, root.destroy)

        def on_toggle(icon, item):
            root.after(0, start_stop_loop)

        menu = pystray.Menu(
            pystray.MenuItem("显示窗口", on_show),
            pystray.MenuItem("启动/停止 WireGuard", on_toggle),
            pystray.MenuItem("退出", on_quit)
        )
        tray_icon = pystray.Icon("wg_updater", make_icon(), "WireGuard Updater", menu)
        threading.Thread(target=tray_icon.run, daemon=True).start()

def minimize_to_tray():
    root.withdraw()  # 隐藏窗口，任务栏不显示
    create_tray_icon_once()

# ==========================
# GUI 布局
# ==========================
root = tk.Tk()
root.title("WireGuard under STUN Peer Using DNS TXT Updater")
root.geometry("800x600")

style = ttk.Style()
style.configure("TButton", font=("Segoe UI", 16, "bold"), padding=10)
style.configure("TLabel", font=("Segoe UI", 14, "bold"))

toggle_button = ttk.Button(root, text="启动", command=start_stop_loop)
toggle_button.pack(pady=15)

status_label = ttk.Label(root, text="未连接")
status_label.pack(pady=10)

text_area = scrolledtext.ScrolledText(root, width=100, height=25, state="disabled", font=("Consolas", 11))
text_area.pack(padx=10, pady=10, fill="both", expand=True)

# 自动最小化到托盘
def on_minimize(event):
    minimize_to_tray()
root.bind("<Unmap>", on_minimize)

root.mainloop()