#!/usr/bin/env python3
"""Pick a PIA region/server and write a WireGuard config for Gluetun."""

from __future__ import annotations

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

import pia


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 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 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 probe_servers(servers: list[pia.WgServer]) -> list[dict[str, Any]]:
    """Probe all servers and return per-IP results sorted by latency."""
    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")))

    results: list[dict[str, Any]] = []
    with ThreadPoolExecutor(max_workers=min(32, max(1, len(servers)))) as pool:
        futures = {
            pool.submit(average_tcp_latency_ms, server.ip, port, timeout, samples): server
            for server in servers
        }
        for future in as_completed(futures):
            server = futures[future]
            value = future.result()
            results.append(
                {
                    "region": server.region,
                    "server_ip": server.ip,
                    "server_cn": server.cn,
                    "latency_ms": None if value is None else round(value, 2),
                }
            )

    results.sort(
        key=lambda item: (
            item["latency_ms"] is None,
            item["latency_ms"] if item["latency_ms"] is not None else float("inf"),
            item["region"],
            item["server_ip"],
        )
    )
    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} ({item['server_cn']}/{item['server_ip']})")
    return results


def pick_fastest(
    candidates: list[str],
    previous_region: str | None,
    previous_server_ip: str | None,
    *,
    exclude_ips: set[str] | None = None,
    apply_margin: bool = True,
) -> tuple[pia.WgServer, list[dict[str, Any]]]:
    exclude_ips = exclude_ips or set()
    margin = float(os.environ.get("LATENCY_SWITCH_MARGIN_MS", "15"))

    serverlist = pia.fetch_serverlist()
    servers: list[pia.WgServer] = []
    for region_id in candidates:
        for server in pia.region_wg_servers(serverlist, region_id):
            if server.ip in exclude_ips:
                continue
            servers.append(server)
    if not servers:
        raise SystemExit("No reachable WireGuard servers left after excludes")

    results = probe_servers(servers)
    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 server")

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

    server = pia.WgServer(
        region=winner["region"],
        ip=winner["server_ip"],
        cn=winner["server_cn"],
    )
    log(f"Selected {server.region} via {server.cn} ({server.ip}) at {winner['latency_ms']:.2f}ms")
    return server, results


def pick_random(
    candidates: list[str],
    previous_region: str | None,
    *,
    exclude_ips: set[str] | None = None,
) -> pia.WgServer:
    exclude_ips = exclude_ips or set()
    serverlist = pia.fetch_serverlist()
    pool = [region for region in candidates if region != previous_region] or list(candidates)
    if previous_region and previous_region not in pool and len(candidates) == 1:
        log(f"Only one region configured; reusing previous: {previous_region}")
    elif previous_region and previous_region not in pool:
        log(f"Excluding previous region: {previous_region}")

    random.shuffle(pool)
    for region_id in pool:
        servers = [
            server
            for server in pia.region_wg_servers(serverlist, region_id)
            if server.ip not in exclude_ips
        ]
        if servers:
            server = random.choice(servers)
            log(f"Randomly selected {server.region} via {server.cn} ({server.ip})")
            return server
    raise SystemExit("No WireGuard servers left after excludes")


def pick_server(
    state_path: Path,
    *,
    exclude_ips: set[str] | None = None,
    apply_margin: bool = True,
) -> tuple[pia.WgServer, str, list[dict[str, Any]]]:
    regions = parse_regions()
    state = read_state(state_path)
    previous_region = state.get("region") if isinstance(state.get("region"), str) else None
    previous_server_ip = state.get("server_ip") if isinstance(state.get("server_ip"), str) else None
    mode = os.environ.get("REGION_SELECT", "fastest").strip().lower() or "fastest"

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

    server, results = pick_fastest(
        regions,
        previous_region,
        previous_server_ip,
        exclude_ips=exclude_ips,
        apply_margin=apply_margin,
    )
    return server, 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 write_state(
    state_path: Path,
    server: pia.WgServer,
    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") == server.region and item.get("server_ip") == server.ip
        ),
        next((item for item in latency_results if item.get("server_ip") == server.ip), None),
    )
    now = datetime.now().astimezone().isoformat(timespec="seconds")
    payload: dict[str, Any] = {
        "region": server.region,
        "server_ip": server.ip,
        "server_cn": server.cn,
        "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"]
    elif skipped and previous.get("latency_ms") is not None:
        payload["latency_ms"] = previous["latency_ms"]
    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 config_looks_valid(wg_path: Path) -> bool:
    if not wg_path.is_file():
        return False
    text = wg_path.read_text(encoding="utf-8", errors="replace")
    return bool(re.search(r"^\[Interface\]", text, flags=re.MULTILINE))


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Rotate PIA WireGuard config for Gluetun")
    parser.add_argument("--list-regions", action="store_true")
    parser.add_argument(
        "--skip-same-server",
        action="store_true",
        help="Skip token/addKey/restarts when the selected server is unchanged",
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Always generate a new keypair and restart containers",
    )
    parser.add_argument(
        "--exclude-server",
        action="append",
        default=[],
        help="Exclude this server IP from selection (repeatable)",
    )
    return parser.parse_args(argv)


def rotate_once(args: argparse.Namespace) -> 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"))
    state = read_state(state_path)
    exclude_ips = {ip for ip in args.exclude_server if ip}
    # Unhealthy/force path wants true #1/#2 ranking without stickiness margin.
    apply_margin = not args.force and not exclude_ips

    server, mode, latency_results = pick_server(
        state_path,
        exclude_ips=exclude_ips,
        apply_margin=apply_margin,
    )
    log(f"Selected endpoint: {server.region} / {server.cn} / {server.ip} (mode={mode})")

    previous_ip = state.get("server_ip") if isinstance(state.get("server_ip"), str) else None
    if args.skip_same_server and not args.force and previous_ip == server.ip and config_looks_valid(wg_path):
        log(
            f"Skipping PIA token/addKey; endpoint unchanged ({server.region}/{server.ip})"
        )
        write_state(
            state_path,
            server,
            mode,
            [],
            latency_results,
            skipped=True,
        )
        log(f"Rotation skipped for {server.region}/{server.ip}")
        return

    # Write temp file on the same filesystem as the destination so os.replace works
    # across Docker bind mounts (/tmp is often a different device than /config).
    wg_path.parent.mkdir(parents=True, exist_ok=True)
    tmp_conf = wg_path.with_name(wg_path.name + ".tmp")
    try:
        log("Generating WireGuard config via native PIA client")
        pia.generate_wg_config(
            require_env("PIA_USER"),
            require_env("PIA_PASS"),
            server,
            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")

        os.replace(tmp_conf, wg_path)
        wg_path.chmod(0o600)
        log(f"Wrote {wg_path}")
    finally:
        tmp_conf.unlink(missing_ok=True)

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


def list_regions_main() -> None:
    for region in pia.list_regions():
        print(
            f"{region['id']:<35} {region['name']:<35} "
            f"country={region['country']:<4} port-forward={region['port_forward']!s:<5} "
            f"offline={region['offline']} wg={region['wg_servers']}"
        )


if __name__ == "__main__":
    # Compat: `rotate.py list-regions`
    argv = sys.argv[1:]
    if argv and argv[0] == "list-regions":
        argv = ["--list-regions", *argv[1:]]
    args = parse_args(argv)
    if args.list_regions:
        list_regions_main()
    else:
        rotate_once(args)