import os, subprocess
import qrcode
from io import BytesIO
import base64

WG_CONF_DIR = "/etc/wireguard"
WG_SERVER_CONF = f"{WG_CONF_DIR}/wg0.conf"
WG_CLIENTS_DIR = os.path.join(WG_CONF_DIR, "clients")
WG_KEYS_DIR = os.path.join(WG_CONF_DIR, "keys")

os.makedirs(WG_CLIENTS_DIR, exist_ok=True)
os.makedirs(WG_KEYS_DIR, exist_ok=True)

def get_public_ip():
    try:
        ip = subprocess.check_output("curl -s ifconfig.me", shell=True, text=True, timeout=3)
        return ip.strip()
    except Exception:
        return ""

def get_or_create_server_keys():
    priv_path = os.path.join(WG_KEYS_DIR, "server.key")
    pub_path = os.path.join(WG_KEYS_DIR, "server.pub")
    if not os.path.exists(priv_path):
        subprocess.run(f"wg genkey | tee {priv_path} | wg pubkey > {pub_path}", shell=True)
    with open(priv_path) as f:
        priv = f.read().strip()
    with open(pub_path) as f:
        pub = f.read().strip()
    return priv, pub

def load_server_conf():
    if not os.path.exists(WG_SERVER_CONF):
        # 기본값 리턴 (최초 실행시)
        return {
            "interface": {
                "Address": "10.8.0.1/24",
                "ListenPort": "51820",
                "PrivateKey": get_or_create_server_keys()[0]
            },
            "peers": []
        }
    with open(WG_SERVER_CONF) as f:
        lines = f.readlines()
    conf = {"interface": {}, "peers": []}
    cur = None
    for l in lines:
        l = l.strip()
        if l.startswith("[Interface]"):
            cur = "interface"
        elif l.startswith("[Peer]"):
            cur = "peer"
            conf["peers"].append({})
        elif "=" in l:
            k, v = [i.strip() for i in l.split("=", 1)]
            if cur == "interface":
                conf["interface"][k] = v
            elif cur == "peer":
                conf["peers"][-1][k] = v
    return conf

def save_server_conf(interface, peers):
    with open(WG_SERVER_CONF, "w") as f:
        f.write("[Interface]\n")
        for k, v in interface.items():
            f.write(f"{k} = {v}\n")
        for peer in peers:
            f.write("\n[Peer]\n")
            for k, v in peer.items():
                f.write(f"{k} = {v}\n")

def restart_wg():
    subprocess.run("systemctl restart wg-quick@wg0", shell=True)

def list_clients():
    clients = []
    for fn in os.listdir(WG_CLIENTS_DIR):
        if fn.endswith(".conf"):
            clients.append(fn[:-5])
    return clients

def gen_client_keys(client_name):
    priv_path = os.path.join(WG_KEYS_DIR, f"{client_name}.key")
    pub_path = os.path.join(WG_KEYS_DIR, f"{client_name}.pub")
    if not os.path.exists(priv_path):
        subprocess.run(f"wg genkey | tee {priv_path} | wg pubkey > {pub_path}", shell=True)
    with open(priv_path) as f:
        priv = f.read().strip()
    with open(pub_path) as f:
        pub = f.read().strip()
    return priv, pub

def create_client_conf(client_name, server_ip, server_pub, vpn_subnet="10.8.0.0/24", port="51820"):
    priv, pub = gen_client_keys(client_name)
    # IP할당: 10.8.0.100~ 로 자동 부여 (이미 할당된 IP 피함)
    used = [p.split('.')[-1].split('/')[0] for p in [c.get('AllowedIPs', '') for c in load_server_conf()["peers"]]]
    idx = 100
    while str(idx) in used:
        idx += 1
    client_addr = f"10.8.0.{idx}/32"
    conf = f"""[Interface]
PrivateKey = {priv}
Address = {client_addr}
DNS = 8.8.8.8

[Peer]
PublicKey = {server_pub}
Endpoint = {server_ip}:{port}
AllowedIPs = 0.0.0.0/0
PersistentKeepalive = 25
"""
    conf_path = os.path.join(WG_CLIENTS_DIR, f"{client_name}.conf")
    with open(conf_path, "w") as f:
        f.write(conf)
    # 서버 conf에 peer 자동 추가(중복 방지)
    srv = load_server_conf()
    peer_obj = {
        "PublicKey": pub,
        "AllowedIPs": client_addr
    }
    exists = False
    for p in srv["peers"]:
        if p.get("PublicKey") == pub:
            exists = True
            break
    if not exists:
        srv["peers"].append(peer_obj)
        save_server_conf(srv["interface"], srv["peers"])
        restart_wg()
    return conf_path

def get_client_conf(client_name):
    path = os.path.join(WG_CLIENTS_DIR, f"{client_name}.conf")
    if os.path.exists(path):
        with open(path) as f:
            return f.read()
    return ""

def delete_client(client_name):
    pub_path = os.path.join(WG_KEYS_DIR, f"{client_name}.pub")
    pubkey = ""
    if os.path.exists(pub_path):
        with open(pub_path) as f:
            pubkey = f.read().strip()
    srv = load_server_conf()
    if pubkey:
        srv["peers"] = [p for p in srv["peers"] if p.get("PublicKey") != pubkey]
        save_server_conf(srv["interface"], srv["peers"])
        restart_wg()
    # 파일 삭제
    for ext in ["key", "pub"]:
        fp = os.path.join(WG_KEYS_DIR, f"{client_name}.{ext}")
        if os.path.exists(fp):
            os.remove(fp)
    confp = os.path.join(WG_CLIENTS_DIR, f"{client_name}.conf")
    if os.path.exists(confp):
        os.remove(confp)

def gen_qrcode(conf_str):
    qr = qrcode.make(conf_str)
    buf = BytesIO()
    qr.save(buf, format="PNG")
    data = base64.b64encode(buf.getvalue()).decode()
    return data
# ... (이전 코드 그대로)
def gen_qrcode_base64(conf_str):
    qr = qrcode.make(conf_str)
    buf = BytesIO()
    qr.save(buf, format="PNG")
    data = base64.b64encode(buf.getvalue()).decode()
    return data

def gen_qr_image(conf_str):
    qr = qrcode.make(conf_str)
    return qr
