#!/usr/bin/env python3
"""
Rotate SECRET_KEY in .env with seamless fallback using PREVIOUS_SECRET_KEY.

Usage:
  python scripts/rotate-secret-key.py [--env-path .env]

Behavior:
- Reads the .env file (default .env)
- Sets PREVIOUS_SECRET_KEY to current SECRET_KEY
- Generates a new SECRET_KEY
- Preserves other variables
- Writes back atomically and sets file mode 600
"""
from __future__ import annotations

import argparse
import os
import secrets
import tempfile
from pathlib import Path


def generate_secret_key(length: int = 32) -> str:
    return secrets.token_urlsafe(length)


def parse_env(contents: str) -> dict[str, str]:
    env: dict[str, str] = {}
    for line in contents.splitlines():
        if not line or line.strip().startswith("#"):
            continue
        if "=" not in line:
            continue
        key, value = line.split("=", 1)
        env[key.strip()] = value.strip()
    return env


def render_env(original: str, updates: dict[str, str]) -> str:
    lines = original.splitlines()
    seen_keys: set[str] = set()
    out_lines: list[str] = []
    for line in lines:
        if not line or line.strip().startswith("#") or "=" not in line:
            out_lines.append(line)
            continue
        key, _ = line.split("=", 1)
        k = key.strip()
        if k in updates:
            out_lines.append(f"{k}={updates[k]}")
            seen_keys.add(k)
        else:
            out_lines.append(line)
            seen_keys.add(k)
    # Append any new keys not present originally
    for k, v in updates.items():
        if k not in seen_keys:
            out_lines.append(f"{k}={v}")
    return "\n".join(out_lines) + "\n"


def main() -> None:
    parser = argparse.ArgumentParser(description="Rotate SECRET_KEY in .env")
    parser.add_argument("--env-path", default=".env", help="Path to .env file")
    args = parser.parse_args()

    env_path = Path(args.env_path)
    if not env_path.exists():
        raise SystemExit(f".env file not found at {env_path}")

    original = env_path.read_text()
    env = parse_env(original)

    current_secret = env.get("SECRET_KEY")
    if not current_secret:
        raise SystemExit("SECRET_KEY not found in .env")

    new_secret = generate_secret_key(32)
    updates = {
        "PREVIOUS_SECRET_KEY": current_secret,
        "SECRET_KEY": new_secret,
    }

    rendered = render_env(original, updates)

    # Atomic write
    with tempfile.NamedTemporaryFile("w", delete=False, dir=str(env_path.parent)) as tmp:
        tmp.write(rendered)
        temp_name = tmp.name
    os.replace(temp_name, env_path)
    os.chmod(env_path, 0o600)

    print("✅ SECRET_KEY rotated successfully.")
    print("   PREVIOUS_SECRET_KEY updated for seamless token validation.")
    print("   Restart the application to apply the new key.")


if __name__ == "__main__":
    main()