#!/usr/bin/env python3
"""Generate a PIA WireGuard config and restart dependent containers."""

from __future__ import annotations

import json
import os
import random
import re
import socket
import subprocess
import sys
import tempfile
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from pathlib import Path
from typing import Any

# EX_TEMPFAIL — entrypoint retries after RATE_LIMIT_WAIT_SECONDS
EXIT_RATE_LIMITED = 75
SERVERLIST_URL = "https://serverlist.piaservers.net/vpninfo/servers/v6"

_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 require_env(name: str) -> str:
    value = os.environ.get(name, "")
    if not value:
        raise SystemExit(f"Missing required env var: {name}")
    return value


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 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 duration_go(seconds: int) -> str:
    """Format seconds as a Go duration string (pia-wg-config uses time.ParseDuration)."""
    if seconds % 3600 == 0:
        return f"{seconds // 3600}h"
    if seconds % 60 == 0:
        return f"{seconds // 60}m"
    return f"{seconds}s"


def parse_list(raw: str) -> list[str]:
    raw = raw.strip()
    if not raw:
        return []
    if raw.startswith("["):
        data = json.loads(raw)
        if not isinstance(data, list):
            raise SystemExit("Expected a JSON array")
        return [str(item).strip() for item in data if str(item).strip()]
    return [part.strip() for part in raw.split(",") if part.strip()]


def parse_regions() -> list[str]:
    regions = parse_list(require_env("PIA_REGIONS"))
    if not regions:
        raise SystemExit("PIA_REGIONS is empty after parsing")
    return regions


def read_state(state_path: Path) -> dict[str, Any]:
    if not state_path.is_file():
        return {}
    try:
        data = json.loads(state_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return {}
    return data if isinstance(data, dict) else {}


def read_previous_region(state_path: Path) -> str | None:
    region = read_state(state_path).get("region")
    return region if isinstance(region, str) and region else None


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


def fetch_serverlist() -> dict[str, Any]:
    """Fetch PIA server list, with on-disk cache shared with pia-wg-config."""
    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")
            cached_raw = None
            cache_age = None

    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:
            raw = resp.read()
        text = raw.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 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 region_wg_servers(serverlist: dict[str, Any], region_id: str) -> list[dict[str, str]]:
    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", {}).get("wg") or []
        return [
            {"ip": s["ip"], "cn": s.get("cn", "")}
            for s in servers
            if isinstance(s, dict) and s.get("ip")
        ]
    log(f"Region {region_id} not found in PIA server list; skipping")
    return []


def tcp_latency_ms(ip: str, port: int, timeout: float) -> float | None:
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(timeout)
    started = time.perf_counter()
    try:
        sock.connect((ip, port))
        return (time.perf_counter() - started) * 1000.0
    except OSError:
        return None
    finally:
        sock.close()


def average_tcp_latency_ms(ip: str, port: int, timeout: float, samples: int) -> float | None:
    readings: list[float] = []
    for _ in range(samples):
        value = tcp_latency_ms(ip, port, timeout)
        if value is None:
            return None
        readings.append(value)
    return sum(readings) / len(readings)


def pick_fastest_region(
    candidates: list[str],
    previous: str | None,
) -> tuple[str, list[dict[str, Any]]]:
    port = int(os.environ.get("LATENCY_PORT", "1337"))
    timeout = float(os.environ.get("LATENCY_TIMEOUT_SECONDS", "2"))
    samples = max(1, int(os.environ.get("LATENCY_SAMPLES", "2")))
    margin = float(os.environ.get("LATENCY_SWITCH_MARGIN_MS", "15"))

    serverlist = fetch_serverlist()
    probes: list[tuple[str, str]] = []
    server_counts: dict[str, int] = {}
    for region_id in candidates:
        servers = region_wg_servers(serverlist, region_id)
        if not servers:
            continue
        server_counts[region_id] = len(servers)
        for server in servers:
            probes.append((region_id, server["ip"]))

    if not probes:
        raise SystemExit("No reachable WireGuard servers found for configured regions")

    best_by_region: dict[str, dict[str, Any]] = {
        region_id: {
            "region": region_id,
            "latency_ms": None,
            "server_ip": None,
            "servers": server_counts[region_id],
            "failures": 0,
        }
        for region_id in server_counts
    }

    with ThreadPoolExecutor(max_workers=min(32, len(probes))) as pool:
        futures = {
            pool.submit(average_tcp_latency_ms, ip, port, timeout, samples): (region_id, ip)
            for region_id, ip in probes
        }
        for future in as_completed(futures):
            region_id, ip = futures[future]
            value = future.result()
            current = best_by_region[region_id]
            if value is None:
                current["failures"] += 1
                continue
            if current["latency_ms"] is None or value < current["latency_ms"]:
                current["latency_ms"] = round(value, 2)
                current["server_ip"] = ip

    results = sorted(
        best_by_region.values(),
        key=lambda item: (
            item["latency_ms"] is None,
            item["latency_ms"] if item["latency_ms"] is not None else float("inf"),
            item["region"],
        ),
    )

    for item in results:
        latency = "timeout" if item["latency_ms"] is None else f"{item['latency_ms']:.2f}ms"
        log(
            f"Latency {item['region']}: {latency} "
            f"(best_ip={item['server_ip']}, servers={item['servers']}, failures={item['failures']})"
        )

    reachable = [item for item in results if item["latency_ms"] is not None]
    if not reachable:
        raise SystemExit("All latency probes failed; cannot select fastest region")

    winner = reachable[0]
    if previous:
        previous_result = next((item for item in reachable if item["region"] == previous), None)
        if previous_result is not None:
            improvement = previous_result["latency_ms"] - winner["latency_ms"]
            if winner["region"] != previous and improvement < margin:
                log(
                    f"Keeping current region {previous} "
                    f"({previous_result['latency_ms']:.2f}ms); "
                    f"best {winner['region']} only {improvement:.2f}ms faster "
                    f"(margin {margin:g}ms)"
                )
                winner = previous_result
            elif winner["region"] == previous:
                log(f"Current region {previous} is still fastest")

    log(f"Fastest region: {winner['region']} ({winner['latency_ms']:.2f}ms via {winner['server_ip']})")
    return winner["region"], results


def pick_random_region(candidates: list[str], previous: str | None) -> str:
    pool = [region for region in candidates if region != previous] or list(candidates)
    if previous and previous not in pool and len(candidates) == 1:
        log(f"Only one region configured; reusing previous: {previous}")
    elif previous and previous not in pool:
        log(f"Excluding previous region: {previous}")
    return random.choice(pool)


def pick_region(state_path: Path) -> tuple[str, str, list[dict[str, Any]]]:
    regions = parse_regions()
    previous = read_previous_region(state_path)
    mode = os.environ.get("REGION_SELECT", "fastest").strip().lower() or "fastest"

    if mode == "random":
        return pick_random_region(regions, previous), mode, []
    if mode != "fastest":
        raise SystemExit(f"Invalid REGION_SELECT '{mode}' (expected fastest|random)")

    region, results = pick_fastest_region(regions, previous)
    return region, mode, results


def parse_restart_containers() -> list[str]:
    gluetun = os.environ.get("GLUETUN_CONTAINER", "m3u-filter-vpn").strip() or "m3u-filter-vpn"
    ordered = [gluetun]
    seen = {gluetun}
    for container in parse_list(os.environ.get("RESTART_CONTAINERS", "")):
        if container in seen:
            continue
        ordered.append(container)
        seen.add(container)
    return ordered


def restart_containers(containers: list[str]) -> None:
    if not containers:
        raise SystemExit("No containers configured to restart")
    for container in containers:
        log(f"Restarting container: {container}")
        result = subprocess.run(["docker", "restart", container], check=False)
        if result.returncode != 0:
            raise SystemExit(f"docker restart failed for container={container}")


def generate_wg_config(region: str, outfile: Path) -> None:
    user = require_env("PIA_USER")
    password = require_env("PIA_PASS")
    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)

    # Ensure cache exists so pia-wg-config can reuse it.
    fetch_serverlist()

    cmd = [
        "pia-wg-config",
        "-v",
        "-r",
        region,
        "-o",
        str(outfile),
        "--serverlist-cache",
        str(cache_path),
        "--serverlist-cache-ttl",
        duration_go(ttl),
        "--serverlist-cache-max-age",
        duration_go(max_age),
        user,
        password,
    ]
    result = subprocess.run(cmd, check=False, capture_output=True, text=True)
    if result.stdout:
        print(result.stdout, end="", file=sys.stderr)
    if result.stderr:
        print(result.stderr, end="", file=sys.stderr)

    output = f"{result.stdout}{result.stderr}"
    if result.returncode == 0:
        return
    if "too_many_attempts" in output or "status 429" in output:
        print(f"pia-wg-config rate-limited for region={region}", file=sys.stderr)
        raise SystemExit(EXIT_RATE_LIMITED)
    print(f"pia-wg-config failed for region={region}", file=sys.stderr)
    raise SystemExit(1)


def config_age_seconds(wg_path: Path) -> float | None:
    if not wg_path.is_file():
        return None
    try:
        return time.time() - wg_path.stat().st_mtime
    except OSError:
        return None


def should_skip_regeneration(region: str, previous: str | None, wg_path: Path) -> bool:
    if env_bool("FORCE_ROTATE", False):
        log("FORCE_ROTATE=true; regenerating WireGuard config")
        return False

    if previous != region:
        return False

    age = config_age_seconds(wg_path)
    if age is None:
        return False

    max_age = parse_duration_seconds(os.environ.get("WG_CONFIG_MAX_AGE", "7d"), 604800)
    if age > max_age:
        log(f"Existing WireGuard config is stale (age {int(age)}s > max {max_age}s); regenerating")
        return False

    if not re.search(
        r"^\[Interface\]",
        wg_path.read_text(encoding="utf-8", errors="replace"),
        flags=re.MULTILINE,
    ):
        return False

    log(
        f"Skipping PIA token/config request; region unchanged ({region}) "
        f"and config age {int(age)}s <= {max_age}s"
    )
    return True


def write_state(
    state_path: Path,
    region: str,
    mode: str,
    restarted: list[str],
    latency_results: list[dict[str, Any]],
    *,
    skipped: bool = False,
) -> None:
    previous = read_state(state_path)
    winner = next((item for item in latency_results if item.get("region") == region), None)
    now = datetime.now().astimezone().isoformat(timespec="seconds")
    payload: dict[str, Any] = {
        "region": region,
        "selection": mode,
        "wg_config": os.environ.get("WG_CONFIG_PATH", "/config/wireguard/wg0.conf"),
        "restarted_containers": restarted,
        "last_checked_at": now,
        "skipped_regeneration": skipped,
    }
    if skipped and isinstance(previous.get("rotated_at"), str):
        payload["rotated_at"] = previous["rotated_at"]
    else:
        payload["rotated_at"] = now
    if winner and winner.get("latency_ms") is not None:
        payload["latency_ms"] = winner["latency_ms"]
        payload["server_ip"] = winner.get("server_ip")
    if latency_results:
        payload["latency_results"] = latency_results

    state_path.parent.mkdir(parents=True, exist_ok=True)
    state_path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
    state_path.chmod(0o644)


def rotate_once() -> None:
    require_env("PIA_USER")
    require_env("PIA_PASS")
    require_env("PIA_REGIONS")

    state_path = Path(os.environ.get("ROTATOR_STATE_PATH", "/config/rotator-state.json"))
    wg_path = Path(os.environ.get("WG_CONFIG_PATH", "/config/wireguard/wg0.conf"))
    previous = read_previous_region(state_path)

    region, mode, latency_results = pick_region(state_path)
    log(f"Selected region: {region} (mode={mode})")

    if should_skip_regeneration(region, previous, wg_path):
        write_state(state_path, region, mode, [], latency_results, skipped=True)
        log(f"Rotation skipped for region={region}")
        return

    with tempfile.TemporaryDirectory(prefix="pia-rotate-") as tmp:
        tmp_conf = Path(tmp) / "wg0.conf"
        log("Generating WireGuard config with pia-wg-config")
        generate_wg_config(region, tmp_conf)
        text = tmp_conf.read_text(encoding="utf-8", errors="replace")
        if not re.search(r"^\[Interface\]", text, flags=re.MULTILINE):
            raise SystemExit("Generated config missing [Interface] section")

        wg_path.parent.mkdir(parents=True, exist_ok=True)
        os.replace(tmp_conf, wg_path)
        wg_path.chmod(0o600)
        log(f"Wrote {wg_path}")

    restarted = parse_restart_containers()
    restart_containers(restarted)
    write_state(state_path, region, mode, restarted, latency_results, skipped=False)
    log(f"Rotation complete for region={region}")


if __name__ == "__main__":
    rotate_once()