#!/usr/bin/env python3
from __future__ import annotations

import argparse
from pathlib import Path

import yaml


MAX_CONSUMERS = 3
MIN_RESERVED_CONNECTIONS = 40
REVIEWED_CELLS = ("apps-pg", "apps-pg-2")


def verify(paths: list[Path]) -> dict[str, int]:
    cells: dict[str, dict] = {}
    role_owner: dict[str, str] = {}
    prefixes: dict[str, str] = {}
    for path in paths:
        document = yaml.safe_load(path.read_text(encoding="utf-8"))
        if document.get("kind") != "Cluster":
            raise ValueError(f"{path}: expected a CNPG Cluster")
        name = document["metadata"]["name"]
        if name not in REVIEWED_CELLS:
            raise ValueError(f"{path}: {name!r} is not a reviewed apps-pg cell")
        cells[name] = document
    missing = set(REVIEWED_CELLS) - cells.keys()
    if missing:
        raise ValueError(f"missing reviewed overflow desired state: {sorted(missing)}")

    counts: dict[str, int] = {}
    for name, document in cells.items():
        spec = document["spec"]
        roles = spec.get("managed", {}).get("roles", [])
        if len(roles) > MAX_CONSUMERS:
            raise ValueError(
                f"{name}: {len(roles)} roles exceeds the {MAX_CONSUMERS}-consumer ceiling"
            )
        max_connections = int(spec["postgresql"]["parameters"]["max_connections"])
        allocated = 0
        for role in roles:
            role_name = role["name"]
            if role_name in role_owner:
                raise ValueError(
                    f"role {role_name!r} appears in both {role_owner[role_name]} and {name}"
                )
            role_owner[role_name] = name
            limit = role.get("connectionLimit")
            if not isinstance(limit, int) or isinstance(limit, bool) or limit < 1:
                raise ValueError(f"{name}/{role_name}: positive connectionLimit is required")
            allocated += limit
        if allocated + MIN_RESERVED_CONNECTIONS > max_connections:
            raise ValueError(
                f"{name}: {allocated} consumer connections leave fewer than "
                f"{MIN_RESERVED_CONNECTIONS} reserved connections"
            )
        resources = spec.get("resources", {})
        if not resources.get("requests") or not resources.get("limits"):
            raise ValueError(f"{name}: pod requests and limits are required")
        prefix = spec["backup"]["barmanObjectStore"]["destinationPath"]
        if prefix in prefixes:
            raise ValueError(f"{name}: backup prefix is shared with {prefixes[prefix]}")
        prefixes[prefix] = name
        counts[name] = len(roles)
    return counts


def main() -> int:
    parser = argparse.ArgumentParser(description="Verify apps-pg cell admission controls")
    parser.add_argument("manifests", nargs="+", type=Path)
    args = parser.parse_args()
    try:
        counts = verify(args.manifests)
    except (KeyError, TypeError, ValueError, yaml.YAMLError) as error:
        parser.error(str(error))
    print(
        "apps-pg capacity verified: "
        + ", ".join(f"{name}={counts[name]}/{MAX_CONSUMERS}" for name in REVIEWED_CELLS)
    )
    return 0


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