#!/usr/bin/env python3
"""
Compare account-level role grants for every role or user between Delphi and Orchard.

For each entity (role or user) in the Delphi account, reports which account-level
roles are granted to it in Delphi but absent in Orchard — i.e., what you need to
add to Orchard to make it match Delphi.

Database roles and application roles are excluded; only account-level (ROLE type)
grants are considered.

Usage:
  pip install snowflake-connector-python
  python3 compare-snowflake-role-grants.py --roles [--entity ROLE_NAME] [--csv out.csv]
  python3 compare-snowflake-role-grants.py --users [--entity USER_NAME] [--csv out.csv]

Authentication:
  Uses key-pair auth if SNOWFLAKE_PRIVATE_KEY_PATH is set, otherwise falls back
  to externalbrowser SSO. Per-account overrides via ORCHARD_* / DELPHI_* prefixes.

Required env vars:
  SNOWFLAKE_USER   - your Snowflake username

Optional env vars:
  SNOWFLAKE_PRIVATE_KEY_PATH         - path to PEM private key file
  SNOWFLAKE_PRIVATE_KEY_PASSPHRASE   - passphrase if key is encrypted
  SNOWFLAKE_ROLE                     - role to use (default: SECURITYADMIN)
  ORCHARD_SNOWFLAKE_USER / DELPHI_SNOWFLAKE_USER   - per-account user override
  ORCHARD_SNOWFLAKE_ROLE / DELPHI_SNOWFLAKE_ROLE   - per-account role override
  ORCHARD_ACCOUNT / DELPHI_ACCOUNT                 - account identifier overrides
  ORCHARD_SNOWFLAKE_WAREHOUSE / DELPHI_SNOWFLAKE_WAREHOUSE - warehouse overrides
"""

import argparse
import csv
import os
import sys
from typing import Any

try:
    import snowflake.connector
except ImportError:
    sys.exit(
        "snowflake-connector-python is required.\n"
        "Install it with: pip install snowflake-connector-python"
    )

# Account-level role grants to roles (GRANTS_TO_ROLES, joined with ROLES to
# filter to ROLE type only — excludes DATABASE_ROLE and APPLICATION_ROLE).
ROLE_GRANTS_QUERY = """
SELECT
    g.GRANTEE_NAME AS entity,
    g.NAME         AS granted_role
FROM SNOWFLAKE.ACCOUNT_USAGE.GRANTS_TO_ROLES g
JOIN SNOWFLAKE.ACCOUNT_USAGE.ROLES r
     ON r.NAME = g.NAME AND r.DELETED_ON IS NULL
WHERE g.GRANTED_ON  = 'ROLE'
  AND g.PRIVILEGE   = 'USAGE'
  AND g.DELETED_ON  IS NULL
  AND g.GRANTED_TO  = 'ROLE'
  AND r.ROLE_TYPE   = 'ROLE'
ORDER BY g.GRANTEE_NAME, g.NAME
"""

# Account-level role grants to users (GRANTS_TO_USERS, joined with ROLES to
# filter to ROLE type only).
USER_GRANTS_QUERY = """
SELECT
    g.GRANTEE_NAME AS entity,
    g.ROLE         AS granted_role
FROM SNOWFLAKE.ACCOUNT_USAGE.GRANTS_TO_USERS g
JOIN SNOWFLAKE.ACCOUNT_USAGE.ROLES r
     ON r.NAME = g.ROLE AND r.DELETED_ON IS NULL
WHERE g.DELETED_ON IS NULL
  AND r.ROLE_TYPE   = 'ROLE'
ORDER BY g.GRANTEE_NAME, g.ROLE;
"""

ALL_ROLES_QUERY = """
SELECT NAME
FROM SNOWFLAKE.ACCOUNT_USAGE.ROLES
WHERE DELETED_ON  IS NULL
  AND ROLE_TYPE   = 'ROLE'
ORDER BY NAME
"""

ALL_USERS_QUERY = """
SELECT NAME
FROM SNOWFLAKE.ACCOUNT_USAGE.USERS
WHERE DELETED_ON IS NULL
ORDER BY NAME
"""


def _private_key_bytes() -> bytes | None:
    path = os.environ.get("SNOWFLAKE_PRIVATE_KEY_PATH")
    if not path:
        return None
    from cryptography.hazmat.backends import default_backend
    from cryptography.hazmat.primitives import serialization

    passphrase_str = os.environ.get("SNOWFLAKE_PRIVATE_KEY_PASSPHRASE")
    passphrase = passphrase_str.encode() if passphrase_str else None
    with open(path, "rb") as fh:
        private_key = serialization.load_pem_private_key(
            fh.read(), password=passphrase, backend=default_backend()
        )
    return private_key.private_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    )


def _conn_params(account_name: str, env_prefix: str, default_warehouse: str) -> dict[str, Any]:
    def env(key: str) -> str | None:
        return os.environ.get(f"{env_prefix}{key}") or os.environ.get(key)

    params: dict[str, Any] = {
        "account": account_name,
        "user": env("SNOWFLAKE_USER"),
        "role": env("SNOWFLAKE_ROLE") or "SECURITYADMIN",
        "warehouse": env("SNOWFLAKE_WAREHOUSE") or default_warehouse,
    }

    pk_bytes = _private_key_bytes()
    if pk_bytes:
        params["private_key"] = pk_bytes
        params["authenticator"] = "snowflake_jwt"
    else:
        params["authenticator"] = "externalbrowser"

    return params


def fetch_grants(
    account_name: str,
    env_prefix: str,
    label: str,
    default_warehouse: str,
    mode: str,
) -> dict[str, set[str]]:
    """Return {entity_name: set(granted_account_role_names)} for all entities."""
    print(f"Connecting to {label} ({account_name})…", flush=True)
    params = _conn_params(account_name, env_prefix, default_warehouse)
    conn = snowflake.connector.connect(**params)
    try:
        cur = conn.cursor(snowflake.connector.DictCursor)

        # Baseline — ensures entities with zero grants still appear in the dict
        baseline_query = ALL_ROLES_QUERY if mode == "roles" else ALL_USERS_QUERY
        cur.execute(baseline_query)
        entities: dict[str, set[str]] = {row["NAME"]: set() for row in cur.fetchall()}

        grants_query = ROLE_GRANTS_QUERY if mode == "roles" else USER_GRANTS_QUERY
        cur.execute(grants_query)
        for row in cur.fetchall():
            entity = row["ENTITY"]
            if entity in entities:
                entities[entity].add(row["GRANTED_ROLE"])

    finally:
        conn.close()

    grant_count = sum(len(v) for v in entities.values())
    print(
        f"  → {len(entities):,} {mode} in {label}, "
        f"{grant_count:,} account-role grants total",
        flush=True,
    )
    return entities


def compute_diff(
    delphi: dict[str, set[str]],
    orchard: dict[str, set[str]],
) -> list[tuple[str, list[str]]]:
    """Return [(entity, [missing_roles])] for entities that have gaps."""
    results = []
    for entity in sorted(delphi):
        missing = sorted(delphi[entity] - orchard.get(entity, set()))
        if missing:
            results.append((entity, missing))
    return results


def print_report(diff: list[tuple[str, list[str]]], mode: str) -> None:
    sep = "=" * 72
    grantee_kw = "ROLE" if mode == "roles" else "USER"

    print(f"\n{sep}")
    print(f"ACCOUNT-ROLE GRANT DIFF: DELPHI → ORCHARD  (mode: --{mode})")
    print("Grants present in Delphi but MISSING in Orchard")
    print(sep)

    if not diff:
        print(f"\nAll Delphi {mode} are fully covered in Orchard. Nothing to add.\n")
        print(f"{sep}\n")
        return

    print(f"\n{mode.upper()} WITH MISSING GRANTS ({len(diff)}):\n")
    for entity, missing in diff:
        print(f"  {entity}  — {len(missing)} grant(s) to add:")
        for role in missing:
            print(f"    GRANT ROLE {role} TO {grantee_kw} {entity};")
        print()

    print(f"{sep}\n")


def write_csv(path: str, diff: list[tuple[str, list[str]]], mode: str) -> None:
    grantee_kw = "ROLE" if mode == "roles" else "USER"
    with open(path, "w", newline="") as fh:
        writer = csv.writer(fh)
        writer.writerow(["entity", "grantee_type", "missing_role", "sql_statement"])
        for entity, missing in diff:
            for role in missing:
                writer.writerow(
                    [entity, grantee_kw, role, f"GRANT ROLE {role} TO {grantee_kw} {entity};"]
                )
    print(f"CSV written to: {path}")


def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )

    mode_group = parser.add_mutually_exclusive_group(required=True)
    mode_group.add_argument("--roles", action="store_true", help="Analyse role-to-role grants")
    mode_group.add_argument("--users", action="store_true", help="Analyse role-to-user grants")

    parser.add_argument(
        "--entity",
        metavar="NAME",
        help="Restrict analysis to a single role or user (case-sensitive)",
    )
    parser.add_argument("--csv", metavar="FILE", help="Also write missing grants to a CSV file")
    parser.add_argument(
        "--orchard-account",
        default=os.environ.get("ORCHARD_ACCOUNT", "sme-orchard"),
        help="Orchard account identifier (default: sme-orchard)",
    )
    parser.add_argument(
        "--delphi-account",
        default=os.environ.get("DELPHI_ACCOUNT", "sme-delphi"),
        help="Delphi account identifier (default: sme-delphi)",
    )
    args = parser.parse_args()

    mode = "roles" if args.roles else "users"

    delphi_grants = fetch_grants(args.delphi_account, "DELPHI_", "Delphi", "DEV_OWS_WH", mode)
    orchard_grants = fetch_grants(args.orchard_account, "ORCHARD_", "Orchard", "DEV_OWS_WAREHOUSE", mode)

    if args.entity:
        name = args.entity
        if name not in delphi_grants:
            sys.exit(f"Entity '{name}' not found among Delphi {mode}.")
        delphi_grants = {name: delphi_grants[name]}

    diff = compute_diff(delphi_grants, orchard_grants)
    print_report(diff, mode)

    if args.csv:
        write_csv(args.csv, diff, mode)


if __name__ == "__main__":
    main()