#!/usr/bin/env python3
"""Minimal PIA WireGuard client: serverlist, token, keygen, addKey."""

from __future__ import annotations

import base64
import json
import os
import re
import socket
import ssl
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import dataclass
from datetime import datetime
from http.client import HTTPSConnection
from pathlib import Path
from typing import Any

from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey

EXIT_RATE_LIMITED = 75
SERVERLIST_URL = "https://serverlist.piaservers.net/vpninfo/servers/v6"
TOKEN_URL = "https://www.privateinternetaccess.com/api/client/v2/token"
PIA_CA_URL = (
    "https://raw.githubusercontent.com/pia-foss/manual-connections/master/ca.rsa.4096.crt"
)

_DURATION_UNITS = {"s": 1, "m": 60, "h": 3600, "d": 86400}


def log(msg: str) -> None:
    print(f"[{datetime.now().astimezone().isoformat(timespec='seconds')}] {msg}", file=sys.stderr)


def parse_duration_seconds(raw: str, default_seconds: int) -> int:
    value = (raw or "").strip().lower()
    if not value:
        return default_seconds
    if value.isdigit():
        return int(value)
    match = re.fullmatch(r"(\d+)([smhd])", value)
    if not match:
        raise SystemExit(f"Invalid duration '{raw}' (use seconds or Ns/Nm/Nh/Nd)")
    return int(match.group(1)) * _DURATION_UNITS[match.group(2)]


def env_bool(name: str, default: bool = False) -> bool:
    raw = os.environ.get(name)
    if raw is None or raw.strip() == "":
        return default
    return raw.strip().lower() in {"1", "true", "yes", "on"}


def serverlist_cache_path() -> Path:
    return Path(os.environ.get("SERVERLIST_CACHE_PATH", "/config/cache/pia-serverlist.json"))


def token_cache_path() -> Path:
    return Path(os.environ.get("TOKEN_CACHE_PATH", "/config/cache/pia-token.json"))


def ca_cache_path() -> Path:
    return Path(os.environ.get("PIA_CA_PATH", "/config/cache/ca.rsa.4096.crt"))


@dataclass(frozen=True)
class WgServer:
    region: str
    ip: str
    cn: str


@dataclass(frozen=True)
class WgKeys:
    private_key: str
    public_key: str


def generate_wg_keys() -> WgKeys:
    private = X25519PrivateKey.generate()
    priv_bytes = private.private_bytes(
        encoding=serialization.Encoding.Raw,
        format=serialization.PrivateFormat.Raw,
        encryption_algorithm=serialization.NoEncryption(),
    )
    pub_bytes = private.public_key().public_bytes(
        encoding=serialization.Encoding.Raw,
        format=serialization.PublicFormat.Raw,
    )
    return WgKeys(
        private_key=base64.b64encode(priv_bytes).decode("ascii"),
        public_key=base64.b64encode(pub_bytes).decode("ascii"),
    )


def parse_serverlist_bytes(raw: bytes) -> dict[str, Any]:
    try:
        data, _ = json.JSONDecoder().raw_decode(raw.decode("utf-8", errors="replace"))
    except json.JSONDecodeError as exc:
        raise SystemExit(f"Failed to parse PIA server list: {exc}") from exc
    if not isinstance(data, dict) or "regions" not in data:
        raise SystemExit("PIA server list missing regions")
    return data


def fetch_serverlist() -> dict[str, Any]:
    cache_path = serverlist_cache_path()
    ttl = parse_duration_seconds(os.environ.get("SERVERLIST_CACHE_TTL", "24h"), 86400)
    max_age = parse_duration_seconds(os.environ.get("SERVERLIST_CACHE_MAX_AGE", "168h"), 604800)
    force = env_bool("SERVERLIST_FORCE_REFRESH", False)
    now = time.time()

    cached_raw: bytes | None = None
    cache_age: float | None = None
    if cache_path.is_file() and not force:
        try:
            cached_raw = cache_path.read_bytes()
            cache_age = now - cache_path.stat().st_mtime
        except OSError as exc:
            log(f"Server list cache unreadable ({exc}); fetching fresh copy")

    if cached_raw is not None and cache_age is not None and cache_age <= ttl:
        log(f"Using cached PIA server list (age {int(cache_age)}s <= TTL {ttl}s)")
        return parse_serverlist_bytes(cached_raw)

    try:
        with urllib.request.urlopen(SERVERLIST_URL, timeout=30) as resp:
            text = resp.read().decode("utf-8", errors="replace")
        data, end = json.JSONDecoder().raw_decode(text)
        stripped = text[:end].encode("utf-8")
        if not isinstance(data, dict) or "regions" not in data:
            raise SystemExit("PIA server list missing regions")
        cache_path.parent.mkdir(parents=True, exist_ok=True)
        tmp = cache_path.with_suffix(cache_path.suffix + ".tmp")
        tmp.write_bytes(stripped)
        os.replace(tmp, cache_path)
        log(f"Fetched and cached PIA server list -> {cache_path}")
        return data
    except (urllib.error.URLError, json.JSONDecodeError, OSError) as exc:
        if cached_raw is not None and cache_age is not None and cache_age <= max_age:
            log(f"Server list fetch failed ({exc}); using stale cache (age {int(cache_age)}s)")
            return parse_serverlist_bytes(cached_raw)
        raise SystemExit(f"Failed to fetch PIA server list: {exc}") from exc


def list_regions(serverlist: dict[str, Any] | None = None) -> list[dict[str, Any]]:
    data = serverlist or fetch_serverlist()
    regions = []
    for region in data.get("regions", []):
        regions.append(
            {
                "id": region.get("id"),
                "name": region.get("name"),
                "country": region.get("country"),
                "port_forward": bool(region.get("port_forward")),
                "offline": bool(region.get("offline")),
                "wg_servers": len((region.get("servers") or {}).get("wg") or []),
            }
        )
    regions.sort(key=lambda item: (item.get("country") or "", item.get("id") or ""))
    return regions


def region_wg_servers(serverlist: dict[str, Any], region_id: str) -> list[WgServer]:
    for region in serverlist.get("regions", []):
        if region.get("id") != region_id:
            continue
        if region.get("offline"):
            log(f"Region {region_id} is marked offline; skipping")
            return []
        servers = (region.get("servers") or {}).get("wg") or []
        out: list[WgServer] = []
        for server in servers:
            if not isinstance(server, dict) or not server.get("ip") or not server.get("cn"):
                continue
            out.append(WgServer(region=region_id, ip=server["ip"], cn=server["cn"]))
        return out
    log(f"Region {region_id} not found in PIA server list; skipping")
    return []


def ensure_pia_ca() -> Path:
    path = ca_cache_path()
    ttl = parse_duration_seconds(os.environ.get("PIA_CA_CACHE_TTL", "30d"), 30 * 86400)
    if path.is_file() and (time.time() - path.stat().st_mtime) <= ttl:
        return path
    try:
        with urllib.request.urlopen(PIA_CA_URL, timeout=30) as resp:
            data = resp.read()
        if b"BEGIN CERTIFICATE" not in data:
            raise SystemExit("Downloaded PIA CA does not look like a certificate")
        path.parent.mkdir(parents=True, exist_ok=True)
        tmp = path.with_suffix(path.suffix + ".tmp")
        tmp.write_bytes(data)
        os.replace(tmp, path)
        log(f"Cached PIA CA certificate -> {path}")
    except (urllib.error.URLError, OSError) as exc:
        if path.is_file():
            log(f"PIA CA download failed ({exc}); using existing {path}")
            return path
        raise SystemExit(f"Failed to download PIA CA certificate: {exc}") from exc
    return path


def _rate_limited(body: str, status_code: int) -> bool:
    return status_code == 429 or "too_many_attempts" in body


def get_token(username: str, password: str) -> str:
    cache_path = token_cache_path()
    ttl = parse_duration_seconds(os.environ.get("TOKEN_CACHE_TTL", "20h"), 20 * 3600)
    force = env_bool("FORCE_TOKEN_REFRESH", False)

    if cache_path.is_file() and not force:
        try:
            cached = json.loads(cache_path.read_text(encoding="utf-8"))
            token = cached.get("token")
            obtained = float(cached.get("obtained_at", 0))
            if isinstance(token, str) and token and (time.time() - obtained) <= ttl:
                log(f"Using cached PIA token (age {int(time.time() - obtained)}s <= TTL {ttl}s)")
                return token
        except (OSError, json.JSONDecodeError, TypeError, ValueError):
            pass

    form = urllib.parse.urlencode({"username": username, "password": password}).encode()
    req = urllib.request.Request(
        TOKEN_URL,
        data=form,
        method="POST",
        headers={"Content-Type": "application/x-www-form-urlencoded"},
    )
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            body = resp.read().decode("utf-8", errors="replace")
            status = resp.status
    except urllib.error.HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")
        if _rate_limited(body, exc.code):
            raise SystemExit(EXIT_RATE_LIMITED) from exc
        raise SystemExit(f"PIA token request failed with status {exc.code}: {body}") from exc
    except urllib.error.URLError as exc:
        raise SystemExit(f"PIA token request failed: {exc}") from exc

    if _rate_limited(body, status):
        raise SystemExit(EXIT_RATE_LIMITED)

    try:
        payload = json.loads(body)
    except json.JSONDecodeError as exc:
        raise SystemExit(f"Invalid PIA token response: {body}") from exc

    token = payload.get("token")
    if not isinstance(token, str) or not token:
        raise SystemExit(f"PIA token response missing token: {body}")

    cache_path.parent.mkdir(parents=True, exist_ok=True)
    cache_path.write_text(
        json.dumps({"token": token, "obtained_at": time.time()}, indent=2) + "\n",
        encoding="utf-8",
    )
    cache_path.chmod(0o600)
    log("Fetched and cached new PIA token")
    return token


class _HTTPSConnectionToIP(HTTPSConnection):
    """HTTPS connection to a fixed IP while presenting server_hostname for SNI/verify."""

    def __init__(self, ip: str, port: int, server_hostname: str, context: ssl.SSLContext):
        super().__init__(server_hostname, port=port, context=context, timeout=30)
        self._connect_ip = ip
        self._server_hostname = server_hostname

    def connect(self) -> None:
        sock = socket.create_connection((self._connect_ip, self.port), self.timeout)
        self.sock = self._context.wrap_socket(sock, server_hostname=self._server_hostname)


def add_key(server: WgServer, token: str, public_key: str) -> dict[str, Any]:
    ca_path = ensure_pia_ca()
    context = ssl.create_default_context(cafile=str(ca_path))
    query = urllib.parse.urlencode({"pt": token, "pubkey": public_key})
    path = f"/addKey?{query}"

    conn = _HTTPSConnectionToIP(server.ip, 1337, server.cn, context)
    try:
        conn.request("GET", path, headers={"Content-Type": "application/json"})
        resp = conn.getresponse()
        body = resp.read().decode("utf-8", errors="replace")
        status = resp.status
    finally:
        conn.close()

    if _rate_limited(body, status):
        raise SystemExit(EXIT_RATE_LIMITED)
    if status != 200:
        raise SystemExit(f"addKey failed for {server.cn}/{server.ip}: status {status}: {body}")

    try:
        payload = json.loads(body)
    except json.JSONDecodeError as exc:
        raise SystemExit(f"Invalid addKey response: {body}") from exc

    if payload.get("status") not in (None, "OK"):
        raise SystemExit(f"addKey rejected: {body}")

    required = ("server_key", "server_ip", "peer_ip", "dns_servers")
    missing = [key for key in required if key not in payload]
    if missing:
        raise SystemExit(f"addKey response missing {missing}: {body}")
    if not payload["dns_servers"]:
        raise SystemExit(f"addKey response missing dns_servers: {body}")

    return payload


def render_wg_config(keys: WgKeys, addkey: dict[str, Any], server: WgServer) -> str:
    endpoint_ip = addkey.get("server_ip") or server.ip
    endpoint_port = addkey.get("server_port") or 1337
    dns = addkey["dns_servers"][0]
    return (
        "[Interface]\n"
        f"PrivateKey = {keys.private_key}\n"
        f"Address = {addkey['peer_ip']}\n"
        f"DNS = {dns}\n"
        "[Peer]\n"
        f"PublicKey = {addkey['server_key']}\n"
        "AllowedIPs = 0.0.0.0/0\n"
        f"Endpoint = {endpoint_ip}:{endpoint_port}\n"
        "PersistentKeepalive = 25\n"
    )


def generate_wg_config(
    username: str,
    password: str,
    server: WgServer,
    outfile: Path,
) -> dict[str, Any]:
    log(f"Requesting PIA token (cached when possible)")
    token = get_token(username, password)
    log("Generating local WireGuard keypair")
    keys = generate_wg_keys()
    log(f"Registering pubkey via addKey on {server.cn} ({server.ip})")
    addkey = add_key(server, token, keys.public_key)
    config = render_wg_config(keys, addkey, server)
    outfile.write_text(config, encoding="utf-8")
    outfile.chmod(0o600)
    return {
        "region": server.region,
        "server_ip": server.ip,
        "server_cn": server.cn,
        "endpoint_ip": addkey.get("server_ip") or server.ip,
        "peer_ip": addkey.get("peer_ip"),
    }