#!/usr/bin/env python3
"""Rotate the initialized PostgreSQL bootstrap role and verify before returning success."""

from __future__ import annotations

import os
import sys
from pathlib import Path

import psycopg2
from psycopg2 import sql


def read_secret(path: str) -> str:
    value = Path(path).read_text()
    if not value or "\x00" in value or any(character.isspace() for character in value):
        raise ValueError("secret must be non-empty and contain no whitespace or NUL bytes")
    return value


def connect(password: str):
    return psycopg2.connect(
        host=os.environ.get("THT_VECTOR_HOST", "vector-db"),
        port=int(os.environ.get("THT_VECTOR_PORT", "5432")),
        dbname=os.environ.get("THT_VECTOR_DATABASE", "thoth"),
        user=os.environ.get("THT_VECTOR_BOOTSTRAP_USER", "postgres"),
        password=password,
        connect_timeout=5,
    )


def alter_current_role(connection, password: str) -> None:
    with connection.cursor() as cursor:
        cursor.execute("SELECT current_user")
        current_user = cursor.fetchone()[0]
        expected = os.environ.get("THT_VECTOR_BOOTSTRAP_USER", "postgres")
        if current_user != expected:
            raise RuntimeError("authenticated role does not match THT_VECTOR_BOOTSTRAP_USER")
        cursor.execute(
            sql.SQL("ALTER ROLE {} PASSWORD {}").format(
                sql.Identifier(current_user), sql.Literal(password)
            )
        )
    connection.commit()


def main() -> int:
    if len(sys.argv) != 3:
        print("usage: rotate-bootstrap-password.py OLD_SECRET NEW_SECRET", file=sys.stderr)
        return 2
    try:
        old_password = read_secret(sys.argv[1])
        new_password = read_secret(sys.argv[2])
        if old_password == new_password:
            raise ValueError("old and new bootstrap passwords must differ")
        old_connection = connect(old_password)
    except Exception as exc:
        print(f"bootstrap rotation refused before change: {type(exc).__name__}", file=sys.stderr)
        return 1

    try:
        alter_current_role(old_connection, new_password)
        try:
            verification = connect(new_password)
            verification.close()
        except Exception as verify_exc:
            try:
                alter_current_role(old_connection, old_password)
            except Exception as restore_exc:
                print(
                    "bootstrap rotation verification failed and password restore failed: "
                    f"{type(verify_exc).__name__}/{type(restore_exc).__name__}",
                    file=sys.stderr,
                )
                return 3
            print(
                f"bootstrap rotation verification failed; old password restored: "
                f"{type(verify_exc).__name__}",
                file=sys.stderr,
            )
            return 1
    except Exception as exc:
        print(f"bootstrap rotation failed: {type(exc).__name__}", file=sys.stderr)
        return 1
    finally:
        old_connection.close()

    print("bootstrap database password rotated and new login verified")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())