#!/usr/bin/env python3
"""
Sub2SOCKS bootstrap - fetch subscription -> pick VLESS node -> sing-box SOCKS proxy.

Only required env var: SUB_URL  (subscription group URL with VLESS nodes)

Optional env vars:
  SOCKS_PORT    (default 1080)         SOCKS5 listen port
  SOCKS_USER / SOCKS_PASS              auth credentials  (empty -> no auth)
  DNS_SERVER    (default tls://8.8.8.8) upstream resolver e.g. tls://<adguard-ip>
  PICK_STRATEGY first | random         (default "first")
  LOG_LEVEL     info | debug | warn | error  (default "info")
"""

import base64
import json
import os
import random
import sys
import textwrap
import urllib.error
import urllib.parse as up
import urllib.request

SUB_URL       = os.environ["SUB_URL"]
SOCKS_PORT    = int(os.environ.get("SOCKS_PORT",  "1080"))
SOCKS_USER    = os.environ.get("SOCKS_USER",     "")
SOCKS_PASS    = os.environ.get("SOCKS_PASS",     "")
DNS_SERVER    = os.environ.get("DNS_SERVER",     "tls://8.8.8.8")
PICK_STRATEGY = os.environ.get("PICK_STRATEGY",  "first")
LOG_LEVEL     = os.environ.get("LOG_LEVEL",      "info")

CONFIG_FILE   = "/tmp/.sub2socks.json"


def _fetch(url):
    """Download sub content, auto decode base64 if needed."""
    req = urllib.request.Request(url, headers={"User-Agent": "clash.meta"})
    with urllib.request.urlopen(req, timeout=30) as r:
        raw = r.read()
    try:
        return raw.decode("utf-8")
    except UnicodeDecodeError:
        return base64.b64decode(raw).decode("utf-8")


def _vless_lines(text):
    """Return VLESS URI lines (handles base64-encoded lines too)."""
    out = []
    for line in text.splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        if line.lower().startswith("vless://"):
            out.append(line)
            continue
        try:
            d = base64.b64decode(line).decode("utf-8")
            if d.lower().startswith("vless://"):
                out.append(d)
        except Exception:
            pass
    return out


def _parse(url_str):
    """Parse a VLESS URL into component dict."""
    p = up.urlparse(url_str.strip())
    q = up.parse_qs(p.query)
    def g(key, default=None):
        if key in q:
            return q[key][0]
        return default
    return {
        "uuid":        p.username,
        "server":      p.hostname,
        "port":        p.port or (443 if p.scheme == "vless+tls" else 80),
        "security":    g("security", "tls"),
        "type":        g("type", "tcp"),
        "sni":         g("sni") or "",
        "alpn":        g("alpn")   or "",
        "pbk":         g("public_key") or g("pbk") or "",
        "sid":         g("sid")    or "",
        "spX":         g("spX") or "",
        "path":        g("path")  or "/",
        "serviceName": g("serviceName") or g("service-name") or "",
        "host":        g("host")  or "",
        "flow":        g("flow")  or "",
    }


def _make_dns_server(addr_str):
    """Parse common DNS address strings into native sing-box v1.13+ server dict.

    Supports: udp://host:port, tls://host(:port), https://host(:port), bare IP/hostname.
    """
    addr = addr_str.strip()

    # Strip scheme to find protocol type
    parts = up.urlparse(addr)
    scheme = parts.scheme.lower() if parts.scheme else ""

    host   = parts.hostname or addr
    port   = parts.port or 53

    # Default TLS SNI (DNS-over-TLS requires server_name)
    sni_map = {
        "8.8.8.8":      "dns.google",
        "8.8.4.4":      "dns.google",
        "1.1.1.1":      "cloudflare-dns.com",
        "1.0.0.1":      "cloudflare-dns.com",
    }

    obj = {"type": scheme or "udp", "server": host, "server_port": port}

    if scheme == "tls":
        obj["tls"] = {"server_name": sni_map.get(host, host)}
    elif scheme == "https":
        port = port or 443
        obj["server_port"] = port
        obj["tls"] = {"server_name": host}

    return obj


def _config(proxy):
    """Build sing-box JSON config."""
    ob = {
        "type":            "vless",
        "tag":             "proxy",
        "server":         proxy["server"],
        "server_port":    proxy["port"],
        "uuid":           proxy["uuid"],
        "packet_encoding": "xudp",
    }

    sec = proxy.get("security") or "tls"

    if sec == "reality":
        ob["tls"] = {
            "enabled":       True,
            "server_name":   proxy.get("sni"),
            "utls":          {"enabled": True, "fingerprint": "chrome"},
            "reality": {
                "enabled":    True,
                "public_key": proxy["pbk"],
                "short_id":   proxy.get("sid") or proxy.get("spX"),
            },
        }
    elif sec == "tls":
        ob["tls"] = {
            "enabled":       True,
            "server_name":  proxy.get("sni"),
            "utls":          {"enabled": True, "fingerprint": "chrome"},
        }
        alpn_raw = proxy.get("alpn")
        if alpn_raw:
            ob["tls"]["alpn"] = [x.strip() for x in alpn_raw.split(",") if x.strip()]

    ttype = proxy.get("type") or "tcp"

    # Plain TCP is the default in sing-box; only add transport block for non-tcp types
    if ttype == "ws":
        tp = {"type": "ws", "path": proxy.get("path") or "/"}
        hv = proxy.get("host")
        if hv:
            tp["headers"] = {"Host": hv}
        ob["transport"] = tp
    elif ttype == "http":
        tp = {"type": "http"}
        hv = proxy.get("host")
        if hv:
            tp["host"] = [x.strip() for x in hv.split(",") if x.strip()]
        tp["path"]   = proxy.get("path") or "/"
        ob["transport"] = tp
    elif ttype == "grpc":
        tp = {"type": "grpc"}
        svc = proxy.get("serviceName") or proxy.get("path")
        if svc:
            tp["service_name"] = svc
        ob["transport"] = tp
    elif ttype == "httpupgrade":
        tp = {"type": "httpupgrade", "path": proxy.get("path") or "/"}
        hv = proxy.get("host")
        if hv:
            tp["headers"] = {"Host": hv}
        ob["transport"] = tp

    # ── Build native v1.13+ DNS server object ──────────────────────
    dns_server_obj = _make_dns_server(DNS_SERVER)

    cfg = {
        "log": {"level": LOG_LEVEL},
        "dns": {
            "servers": [
                {**dns_server_obj, "tag": "remote"},
                {"type": "local",  "tag": "local"},
            ],
            "final": "remote",
        },
        "inbounds": [{
            "type":         "socks",
            "listen":       "::",
            "listen_port":  SOCKS_PORT,
        }],
        "outbounds": [
            ob,                                    # index 0 -> tag: proxy
            {"type": "direct", "tag": "direct"},   # bypass
        ],
        "route": {
            "default_domain_resolver": "remote",
            "rules": [{
                "clash_mode": "direct",
                "outbound":   "direct",
            }],
            "final":     "proxy",                   # default -> outbound[0]
        },
    }

    if SOCKS_USER:
        cfg["inbounds"][0]["users"] = [
            {"username": SOCKS_USER,  "password": SOCKS_PASS},
        ]
    return cfg


def main():
    print(textwrap.dedent(f"""
        +-------------------------------------------------------+
        |   Sub2SOCKS -- bootstrapping                          |
        +-------------------------------------------------------+
        |  SUB_URL       : {SUB_URL}
        |  SOCKS port    : {SOCKS_PORT:>6d}
        |  DNS upstream  : {DNS_SERVER}
        |  PICK_STRATEGY : {PICK_STRATEGY}
        |  LOG_LEVEL     : {LOG_LEVEL}
        +-------------------------------------------------------+
    """).strip())

    # 1. fetch and filter
    try:
        data = _fetch(SUB_URL)
    except urllib.error.URLError as e:
        print(f"ERROR: Cannot reach {SUB_URL}: {e}", flush=True)
        sys.exit(1)
    except Exception as e:
        print(f"ERROR: Subscription error: {e}", flush=True)
        sys.exit(1)

    lines = _vless_lines(data)
    if not lines:
        print("ERROR: No VLESS nodes found in subscription!")
        print(data[:500])
        sys.exit(1)

    print(f"  Found {len(lines)} VLESS node(s) - picking one.", flush=True)

    chosen = random.choice(lines) if PICK_STRATEGY == "random" else lines[0]

    proxy  = _parse(chosen)
    print(f"  Node: {proxy['server']}:{proxy['port']}", flush=True)

    conf = _config(proxy)
    with open(CONFIG_FILE, "w") as fp:
        json.dump(conf, fp, indent=2)

    print(f"  Config written -> {CONFIG_FILE}", flush=True)

    # exec sing-box in-place (native v1.13+ DNS, no shim required)
    os.execve(
        "/usr/local/bin/sing-box",
        ["sing-box", "run", "-c", CONFIG_FILE],
        dict(os.environ),
    )


if __name__ == "__main__":
    main()