#!/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"
GTOKEN_URL = "https://www.privateinternetaccess.com/gtoken/generateToken"
PIA_CA_URL = (
    "https://raw.githubusercontent.com/pia-foss/manual-connections/master/ca.rsa.4096.crt"
)

# Cloudflare error 1010 blocks Python's default User-Agent.
HTTP_HEADERS = {
    "User-Agent": (
        "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
        "(KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36"
    ),
    "Accept": "application/json, text/plain, */*",
    "Accept-Language": "en-US,en;q=0.9",
}

_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]:
    return _region_servers(serverlist, region_id, "wg")


def region_meta_servers(serverlist: dict[str, Any], region_id: str) -> list[WgServer]:
    return _region_servers(serverlist, region_id, "meta")


def _region_servers(serverlist: dict[str, Any], region_id: str, kind: 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(kind) 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 pia_ssl_context() -> ssl.SSLContext:
    """SSL context trusting PIA's CA.

    OpenSSL 3.2+ / Python 3.13+ enable X509_STRICT by default, which rejects
    PIA's ca.rsa.4096.crt because basicConstraints is not marked critical.
    """
    context = ssl.create_default_context(cafile=str(ensure_pia_ca()))
    if hasattr(ssl, "VERIFY_X509_STRICT"):
        context.verify_flags &= ~ssl.VERIFY_X509_STRICT
    return context


def rate_limit_path() -> Path:
    return Path(os.environ.get("RATE_LIMIT_PATH", "/config/cache/pia-rate-limit.json"))


def rate_limit_wait_seconds() -> int:
    return int(os.environ.get("RATE_LIMIT_WAIT_SECONDS", "3600"))


def rate_limit_remaining_seconds() -> int:
    path = rate_limit_path()
    if not path.is_file():
        return 0
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
        until = float(data.get("until", 0))
    except (OSError, json.JSONDecodeError, TypeError, ValueError):
        return 0
    return max(0, int(until - time.time()))


def mark_rate_limited(wait_s: int | None = None, reason: str = "") -> int:
    wait = rate_limit_wait_seconds() if wait_s is None else wait_s
    until = time.time() + wait
    path = rate_limit_path()
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps(
            {
                "until": until,
                "wait_seconds": wait,
                "reason": reason,
                "marked_at": datetime.now().astimezone().isoformat(timespec="seconds"),
            },
            indent=2,
        )
        + "\n",
        encoding="utf-8",
    )
    log(f"Marked PIA rate-limit cooldown for {wait}s ({reason or 'rate limited'})")
    return wait


def clear_rate_limit() -> None:
    path = rate_limit_path()
    try:
        path.unlink(missing_ok=True)
    except OSError:
        pass


def raise_if_rate_limit_cooldown() -> None:
    remaining = rate_limit_remaining_seconds()
    if remaining > 0:
        log(f"PIA rate-limit cooldown active; {remaining}s remaining (no API calls)")
        raise SystemExit(EXIT_RATE_LIMITED)


class RateLimited(RuntimeError):
    """One auth endpoint reported rate limiting; other methods may still work."""


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


def _cloudflare_blocked(status_code: int, body: str) -> bool:
    return status_code == 403 and ("error code: 1010" in body or "Attention Required" in body)


def _basic_auth_header(username: str, password: str) -> str:
    token = base64.b64encode(f"{username}:{password}".encode("utf-8")).decode("ascii")
    return f"Basic {token}"


def _parse_token_body(body: str, source: str) -> str:
    try:
        payload = json.loads(body)
    except json.JSONDecodeError as exc:
        raise RuntimeError(f"{source}: invalid JSON: {body[:200]}") from exc
    token = payload.get("token")
    if not isinstance(token, str) or not token:
        raise RuntimeError(f"{source}: missing token in response: {body[:200]}")
    return token


def _store_token(token: str) -> str:
    cache_path = token_cache_path()
    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)
    return token


def _token_via_central(username: str, password: str) -> str:
    form = urllib.parse.urlencode({"username": username, "password": password}).encode()
    req = urllib.request.Request(
        TOKEN_URL,
        data=form,
        method="POST",
        headers={
            **HTTP_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 RateLimited(f"central token API: {body[:200]}") from exc
        if _cloudflare_blocked(exc.code, body):
            raise RuntimeError("central token API blocked by Cloudflare (1010)") from exc
        raise RuntimeError(f"central token API status {exc.code}: {body[:200]}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"central token API network error: {exc}") from exc

    if _rate_limited(body, status):
        raise RateLimited(f"central token API: {body[:200]}")
    return _parse_token_body(body, "central token API")


def _token_via_gtoken(username: str, password: str) -> str:
    req = urllib.request.Request(
        GTOKEN_URL,
        method="GET",
        headers={
            **HTTP_HEADERS,
            "Authorization": _basic_auth_header(username, password),
        },
    )
    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 RateLimited(f"gtoken API: {body[:200]}") from exc
        if _cloudflare_blocked(exc.code, body):
            raise RuntimeError("gtoken API blocked by Cloudflare (1010)") from exc
        raise RuntimeError(f"gtoken API status {exc.code}: {body[:200]}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"gtoken API network error: {exc}") from exc

    if _rate_limited(body, status):
        raise RateLimited(f"gtoken API: {body[:200]}")
    return _parse_token_body(body, "gtoken API")


def _token_via_meta(username: str, password: str, meta: WgServer) -> str:
    context = pia_ssl_context()
    conn = _HTTPSConnectionToIP(meta.ip, 443, meta.cn, context, timeout=10)
    try:
        conn.request(
            "GET",
            "/authv3/generateToken",
            headers={
                **HTTP_HEADERS,
                "Authorization": _basic_auth_header(username, password),
            },
        )
        resp = conn.getresponse()
        body = resp.read().decode("utf-8", errors="replace")
        status = resp.status
    finally:
        conn.close()

    if _rate_limited(body, status):
        raise RateLimited(f"meta {meta.cn}: {body[:200]}")
    if status != 200:
        raise RuntimeError(f"meta {meta.cn}/{meta.ip} status {status}: {body[:200]}")
    return _parse_token_body(body, f"meta {meta.cn}")


def get_token(username: str, password: str, preferred_region: str | None = None) -> str:
    raise_if_rate_limit_cooldown()

    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

    errors: list[str] = []
    saw_rate_limit = False

    for label, getter in (
        ("central v2 token API", lambda: _token_via_central(username, password)),
        ("gtoken API", lambda: _token_via_gtoken(username, password)),
    ):
        try:
            log(f"Requesting PIA token via {label}")
            token = getter()
            clear_rate_limit()
            log(f"Fetched and cached new PIA token via {label}")
            return _store_token(token)
        except RateLimited as exc:
            saw_rate_limit = True
            errors.append(f"{label}: {exc}")
            log(f"Token via {label} rate-limited: {exc}")
        except Exception as exc:  # noqa: BLE001 - try next auth method
            errors.append(f"{label}: {exc}")
            log(f"Token via {label} failed: {exc}")

    # Meta servers bypass Cloudflare; prefer the selected region, then any region.
    try:
        serverlist = fetch_serverlist()
    except SystemExit as exc:
        errors.append(f"serverlist for meta auth: {exc}")
        serverlist = None

    meta_candidates: list[WgServer] = []
    if serverlist is not None:
        if preferred_region:
            meta_candidates.extend(region_meta_servers(serverlist, preferred_region))
        if not meta_candidates:
            for region in serverlist.get("regions", []):
                rid = region.get("id")
                if not isinstance(rid, str):
                    continue
                meta_candidates.extend(region_meta_servers(serverlist, rid))
                if len(meta_candidates) >= 8:
                    break
        # Prefer legacy CN hostnames; new Server-* meta hosts often hang on /authv3.
        meta_candidates.sort(key=lambda item: item.cn.startswith("Server-"))

    for meta in meta_candidates[:5]:
        label = f"meta {meta.cn} ({meta.ip})"
        try:
            log(f"Requesting PIA token via {label}")
            token = _token_via_meta(username, password, meta)
            clear_rate_limit()
            log(f"Fetched and cached new PIA token via {label}")
            return _store_token(token)
        except RateLimited as exc:
            saw_rate_limit = True
            errors.append(f"{label}: {exc}")
            log(f"Token via {label} rate-limited: {exc}")
        except Exception as exc:  # noqa: BLE001
            errors.append(f"{label}: {exc}")
            log(f"Token via {label} failed: {exc}")

    if saw_rate_limit:
        mark_rate_limited(reason="token endpoints rate-limited")
        raise SystemExit(EXIT_RATE_LIMITED)
    raise SystemExit("PIA token request failed:\n- " + "\n- ".join(errors))


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,
        timeout: float = 30,
    ):
        super().__init__(server_hostname, port=port, context=context, timeout=timeout)
        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]:
    raise_if_rate_limit_cooldown()

    context = pia_ssl_context()
    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={**HTTP_HEADERS, "Content-Type": "application/json"})
        resp = conn.getresponse()
        body = resp.read().decode("utf-8", errors="replace")
        status = resp.status
    except ssl.SSLError as exc:
        raise SystemExit(f"addKey TLS failed for {server.cn}/{server.ip}: {exc}") from exc
    finally:
        conn.close()

    if _rate_limited(body, status):
        mark_rate_limited(reason=f"addKey {server.cn}")
        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("Requesting PIA token (cached when possible)")
    token = get_token(username, password, preferred_region=server.region)
    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"),
    }