#!/usr/bin/env python3
"""Mood VPN = Psiphon (free, no login, stealth protocols that get past school filters).

Runs as the user. psiphon-tunnel-core opens SOCKS 127.0.0.1:1081 + HTTP 127.0.0.1:8081 and, while it's up,
the desktop proxy (GSettings: Mood Browser, Firefox, GTK apps) and Spotify's own proxy setting point at it.
On top of that, a root sing-box TUN (mood-vpn-tun.service, via mood-admin vpn-tun) routes ALL device traffic
and DNS into Psiphon, so every app is tunnelled. If the TUN can't start, it stays proxy-only.
Turning it off puts the school proxy from Settings → School network back.

  mood-vpn on|off|toggle|status|restore|log
"""
import json
import os
import re
import signal
import subprocess
import sys
import time
from pathlib import Path

sys.path.insert(0, "/usr/lib/mood/runtime")
try:
    import moodactivity  # Dynamic Island pill while connecting
except Exception:  # pragma: no cover
    moodactivity = None


def _island(op, /, **kw):
    if not moodactivity:
        return
    try:
        if op == "clear":
            moodactivity.clear("vpn")
        elif op == "flash":
            moodactivity.flash("vpn", **kw)
        else:
            moodactivity.set("vpn", **kw)
    except Exception:
        pass

HOME = Path.home()
CORE = "/usr/lib/mood/psiphon/psiphon-tunnel-core"
BASE_CONFIG = "/usr/lib/mood/psiphon/psiphon.config"
STATE = HOME / ".local/state/mood-vpn"
PID, LOG, CFG, SAVED = STATE / "pid", STATE / "log.json", STATE / "config.json", STATE / "spotify-saved.json"
SETTINGS = HOME / ".config/mood/settings.json"
SOCKS, HTTP = 1081, 8081
SPOTIFY_PREFS = [HOME / ".config/spotify/prefs", HOME / ".var/app/com.spotify.Client/config/spotify/prefs"]
REGIONS = {"", "US", "GB", "CA", "DE", "NL", "FR", "JP", "SG", "AU", "IN"}


def _settings():
    try:
        return json.loads(SETTINGS.read_text())
    except Exception:
        return {}


def _save(patch):
    s = _settings()
    s.setdefault("vpn", {}).update(patch)
    SETTINGS.parent.mkdir(parents=True, exist_ok=True)
    tmp = SETTINGS.with_suffix(".vpn.tmp")
    tmp.write_text(json.dumps(s, indent=2))
    tmp.replace(SETTINGS)


def _pid():
    try:
        p = int(PID.read_text())
        os.kill(p, 0)
        if CORE in Path(f"/proc/{p}/cmdline").read_bytes().decode(errors="ignore"):
            return p
    except Exception:
        pass
    return None


def _notices():
    try:
        lines = LOG.read_text(errors="ignore").splitlines()[-400:]
    except Exception:
        return []
    out = []
    for ln in lines:
        try:
            out.append(json.loads(ln))
        except ValueError:
            pass
    return out


def _tunnel():
    """(connected, region) from Psiphon's notices: the last Tunnels count, and the region it reported."""
    up, region = False, ""
    for n in _notices():
        t, d = n.get("noticeType"), n.get("data") or {}
        if t == "Tunnels":
            up = d.get("count", 0) > 0
        elif t == "ActiveTunnel" and d.get("serverRegion"):
            region = d["serverRegion"]
        elif t == "ConnectedServerRegion" and d.get("serverRegion"):
            region = d["serverRegion"]
    return up, region


def _gs(*a):
    try:
        subprocess.run(["gsettings", *a], capture_output=True, timeout=10)
    except (OSError, subprocess.SubprocessError):
        pass


def _ignore(no):
    return "[" + ",".join(f"'{h.strip()}'" for h in no.split(",") if h.strip()) + "]"


def _school():
    return _settings().get("proxy") or {"mode": "none", "pac": "", "host": "", "port": 8080, "noProxy": "localhost,127.0.0.0/8,::1"}


def _proxy_vpn():
    no = _school().get("noProxy") or "localhost,127.0.0.0/8,::1"
    _gs("set", "org.gnome.system.proxy", "mode", "manual")
    _gs("set", "org.gnome.system.proxy", "autoconfig-url", "")
    _gs("set", "org.gnome.system.proxy", "ignore-hosts", _ignore(no))
    for scheme in ("http", "https", "ftp"):
        _gs("set", f"org.gnome.system.proxy.{scheme}", "host", "127.0.0.1")
        _gs("set", f"org.gnome.system.proxy.{scheme}", "port", str(HTTP))
    _gs("set", "org.gnome.system.proxy.socks", "host", "127.0.0.1")
    _gs("set", "org.gnome.system.proxy.socks", "port", str(SOCKS))


def _proxy_school():
    """Same GSettings writes as Settings → School network (plugins/connect.py proxy_set)."""
    p = _school()
    mode = p.get("mode") if p.get("mode") in ("none", "auto", "manual") else "none"
    _gs("set", "org.gnome.system.proxy", "mode", mode)
    _gs("set", "org.gnome.system.proxy", "autoconfig-url", p.get("pac", "") if mode == "auto" else "")
    _gs("set", "org.gnome.system.proxy", "ignore-hosts", _ignore(p.get("noProxy") or "localhost,127.0.0.0/8,::1"))
    for scheme in ("http", "https", "ftp"):
        _gs("set", f"org.gnome.system.proxy.{scheme}", "host", p.get("host", "") if mode == "manual" else "")
        _gs("set", f"org.gnome.system.proxy.{scheme}", "port", str(p.get("port") or 8080))
    _gs("set", "org.gnome.system.proxy.socks", "host", "")
    _gs("set", "org.gnome.system.proxy.socks", "port", "0")


SP_RE = re.compile(r"^network\.proxy\.(mode|addr)=")


def _spotify_on():
    saved = {}
    for f in SPOTIFY_PREFS:
        if not f.is_file():
            continue
        lines = f.read_text(errors="ignore").splitlines()
        saved[str(f)] = [ln for ln in lines if SP_RE.match(ln)]
        keep = [ln for ln in lines if not SP_RE.match(ln)]
        f.write_text("\n".join(keep + [f'network.proxy.addr="127.0.0.1:{HTTP}@http"', "network.proxy.mode=2"]) + "\n")
    if saved and not SAVED.exists():
        SAVED.write_text(json.dumps(saved))


def _spotify_off():
    try:
        saved = json.loads(SAVED.read_text())
    except Exception:
        saved = {}
    for f in SPOTIFY_PREFS:
        if not f.is_file():
            continue
        lines = f.read_text(errors="ignore").splitlines()
        if not any(f"127.0.0.1:{HTTP}@http" in ln for ln in lines):
            continue
        keep = [ln for ln in lines if not SP_RE.match(ln)]
        f.write_text("\n".join(keep + saved.get(str(f), [])) + "\n")
    SAVED.unlink(missing_ok=True)


def _spotify_running():
    return subprocess.run(["pgrep", "-u", str(os.getuid()), "-x", "spotify"], capture_output=True).returncode == 0


def _tun(cmd):
    try:
        r = subprocess.run(["sudo", "-n", "/usr/lib/mood/runtime/mood-admin", "vpn-tun", cmd],
                           stdin=subprocess.DEVNULL, capture_output=True, text=True, timeout=30)
        return r.returncode == 0 and bool(json.loads(r.stdout or "{}").get("tun"))
    except Exception:
        return False


def _notify(body):
    subprocess.Popen(["notify-send", "-a", "Mood VPN", "-i", "network-vpn", "Mood VPN", body],
                     stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def status():
    p = _pid()
    up, region = _tunnel() if p else (False, "")
    return {"active": bool(p and up), "busy": bool(p and not up), "running": bool(p), "region": region,
            "want": bool((_settings().get("vpn") or {}).get("want")), "egress": (_settings().get("vpn") or {}).get("region", ""),
            "socks": SOCKS, "http": HTTP, "full": bool(p and up and _tun("status"))}


def _stop():
    _tun("off")
    p = _pid()
    if p:
        os.kill(p, signal.SIGTERM)
        for _ in range(30):
            if not _pid():
                break
            time.sleep(0.1)
        else:
            os.kill(p, signal.SIGKILL)
    PID.unlink(missing_ok=True)


def off(remember=True):
    was = bool(_pid())
    _stop()
    if was:
        _island("flash", title="VPN off", sub="Back on your normal connection", icon="shield", color="#64748b", seconds=2.5)
    _proxy_school()
    _spotify_off()
    if remember:
        _save({"want": False})
    return status()


def on(timeout=90):
    if _pid() and _tunnel()[0]:
        return status()
    _stop()
    STATE.mkdir(parents=True, exist_ok=True)
    (STATE / "data").mkdir(exist_ok=True)
    cfg = json.loads(Path(BASE_CONFIG).read_text())
    cfg["DataRootDirectory"] = str(STATE / "data")
    region = (_settings().get("vpn") or {}).get("region", "")
    if region in REGIONS and region:
        cfg["EgressRegion"] = region
    school = _school()
    if school.get("mode") == "manual" and school.get("host"):
        cfg["UpstreamProxyUrl"] = f"http://{school['host']}:{school.get('port') or 8080}"
    CFG.write_text(json.dumps(cfg))
    _island("set", kind="live", title="Connecting VPN", sub=f"Finding a server{' in ' + region if region else ''}…", icon="shield",
            color="#8b5cf6", spinner=True, priority=60, expires=int((time.time() + timeout + 5) * 1000), app="settings")
    with open(LOG, "w") as log:
        proc = subprocess.Popen([CORE, "-config", str(CFG)], stdout=log, stderr=subprocess.STDOUT,
                                stdin=subprocess.DEVNULL, start_new_session=True, cwd=str(STATE))
    PID.write_text(str(proc.pid))
    _save({"want": True})
    end = time.time() + timeout
    while time.time() < end:
        if proc.poll() is not None:
            break
        if _tunnel()[0]:
            _proxy_vpn()
            _spotify_on()
            _tun("on")
            st = status()
            _island("flash", title="VPN on", sub=(st.get("region") or "Connected") + (" · everything tunnelled" if st.get("full") else " · browsers tunnelled"),
                    icon="shield-check", color="#22c55e", seconds=4)
            return st
        time.sleep(0.5)
    _stop()
    _save({"want": False})
    _island("flash", title="VPN couldn't connect", sub="Try another country or a hotspot", icon="shield-alert", color="#ef4444", seconds=6)
    err = next((str((n.get("data") or {}).get("message", "")) for n in reversed(_notices()) if n.get("noticeType") in ("Alert", "Error")), "")
    raise RuntimeError("Couldn't reach a Psiphon server from this network" + (f" ({err[:120]})" if err else "") +
                       ". Try again, pick another country, or use a phone hotspot.")


def restore():
    """At login: fix a proxy left pointing at a dead tunnel, and reconnect if the VPN was on last time."""
    if _pid() and _tunnel()[0]:
        _tun("on")
        return status()
    _stop()
    _proxy_school()
    _spotify_off()
    if (_settings().get("vpn") or {}).get("want"):
        try:
            return on()
        except Exception:
            return status()
    return status()


def log_text():
    out = []
    for n in _notices()[-150:]:
        t, d = n.get("noticeType", ""), n.get("data") or {}
        if t in ("Info", "Alert", "Error", "Tunnels", "ActiveTunnel", "ConnectingServer", "ConnectedServer", "CandidateServers",
                 "AvailableEgressRegions", "Exiting", "UpstreamProxyError", "ServerAlert"):
            out.append(f"{n.get('timestamp', '')[11:19]} {t}: {json.dumps(d)[:300]}")
    return "\n".join(out)


def main():
    cmd = sys.argv[1] if len(sys.argv) > 1 else "toggle"
    try:
        if cmd == "toggle":
            cmd = "off" if _pid() else "on"
        if cmd == "on":
            st = on()
            msg = f"On ({st['region'] or 'connected'}). " + ("Everything is tunnelled." if st.get("full") else "Browsers are tunnelled.")
            if _spotify_running():
                msg += " Restart Spotify so it uses the VPN."
            if not os.environ.get("MOOD_VPN_QUIET"):
                _notify(msg)
            print(json.dumps(st))
        elif cmd == "off":
            print(json.dumps(off()))
            if not os.environ.get("MOOD_VPN_QUIET"):
                _notify("VPN off")
        elif cmd == "status":
            print(json.dumps(status()))
        elif cmd == "restore":
            print(json.dumps(restore()))
        elif cmd == "log":
            print(log_text())
        else:
            print(__doc__.strip())
            return 2
    except Exception as e:
        print(str(e), file=sys.stderr)
        if not os.environ.get("MOOD_VPN_QUIET"):
            _notify(str(e)[:200])
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())
