#!/usr/bin/env python3
"""Schedule PIA WireGuard rotations via ROTATE_CRON (croniter)."""

from __future__ import annotations

import os
import re
import subprocess
import sys
import time
from datetime import datetime
from zoneinfo import ZoneInfo

from croniter import croniter

# Matches rotate.sh EXIT_RATE_LIMITED (EX_TEMPFAIL)
EXIT_RATE_LIMITED = 75

CRON_MACROS = {
    "@yearly": "0 0 1 1 *",
    "@annually": "0 0 1 1 *",
    "@monthly": "0 0 1 * *",
    "@weekly": "0 0 * * 0",
    "@daily": "0 0 * * *",
    "@midnight": "0 0 * * *",
    "@hourly": "0 * * * *",
}


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


def zone() -> ZoneInfo:
    name = os.environ.get("TZ") or "UTC"
    try:
        return ZoneInfo(name)
    except Exception as exc:
        raise SystemExit(f"Invalid TZ '{name}': {exc}") from exc


def hhmm_to_cron(value: str) -> str:
    match = re.fullmatch(r"([0-9]{1,2}):([0-9]{2})", value.strip())
    if not match:
        raise ValueError(f"Invalid ROTATE_AT '{value}' (expected HH:MM)")
    hour = int(match.group(1))
    minute = int(match.group(2))
    if not (0 <= hour <= 23 and 0 <= minute <= 59):
        raise ValueError(f"Invalid ROTATE_AT '{value}' (expected HH:MM)")
    return f"{minute} {hour} * * *"


def resolve_cron_expr() -> str:
    cron = os.environ.get("ROTATE_CRON", "").strip()
    at = os.environ.get("ROTATE_AT", "").strip()
    if cron:
        expr = cron
    elif at:
        expr = hhmm_to_cron(at)
    else:
        expr = "0 3 * * *"

    expr = CRON_MACROS.get(expr.lower(), expr)
    if not croniter.is_valid(expr):
        raise SystemExit(f"Invalid ROTATE_CRON '{expr}' (expected a 5-field cron expression)")
    return expr


def next_run(expr: str, after: datetime) -> datetime:
    return croniter(expr, after).get_next(datetime)


def sleep_until_next_rotate(expr: str) -> None:
    tz = zone()
    now = datetime.now(tz)
    nxt = next_run(expr, now)
    wait_s = max(0.0, (nxt - now).total_seconds())
    log(f"Next rotation at {nxt.isoformat(timespec='seconds')} (cron '{expr}', TZ={tz.key}) in {int(wait_s)}s")
    time.sleep(wait_s)


def run_rotation(reason: str) -> None:
    wait_s = int(os.environ.get("RATE_LIMIT_WAIT_SECONDS", "3600"))
    while True:
        log(reason)
        result = subprocess.run(["/usr/local/bin/rotate.py"], check=False)
        if result.returncode == 0:
            return
        if result.returncode == EXIT_RATE_LIMITED:
            log(f"PIA rate-limited (too many attempts); waiting {wait_s}s before retry")
            time.sleep(wait_s)
            reason = "Retrying rotation after rate-limit wait"
            continue
        raise SystemExit(f"Rotation failed with exit code {result.returncode}")


def main() -> None:
    expr = resolve_cron_expr()
    # Fail fast on bad cron / TZ before rotating.
    next_run(expr, datetime.now(zone()))
    log(f"Starting gluetun PIA WireGuard rotator (TZ={zone().key}, ROTATE_CRON='{expr}')")

    run_rotation("Running rotation on startup")
    while True:
        sleep_until_next_rotate(expr)
        run_rotation("Running scheduled rotation")


if __name__ == "__main__":
    main()