#!/usr/bin/env python3
"""
Compare all roles across the Orchard and Delphi Snowflake accounts.

Produces a report of:
  - Roles present in both accounts with differing attributes
  - Roles present in both accounts with differing parent-role grants (inheritance)
  - Roles only in Delphi

Uses real-time SHOW ROLES (and SHOW GRANTS TO ROLE for parent-role inheritance)
rather than SNOWFLAKE.ACCOUNT_USAGE.ROLES, which has a latency of up to ~2 hours.
Only account-level roles are compared (database / application roles are excluded).

Usage:
  pip install snowflake-connector-python
  python3 compare-snowflake-roles.py [--csv output.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 concurrent.futures import ThreadPoolExecutor, as_completed
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"
    )

# ---------------------------------------------------------------------------
# Attributes fetched from SHOW ROLES and compared between accounts.
# SHOW ROLES gives real-time results; ACCOUNT_USAGE.ROLES has a ~2h delay.
# Database / application roles are excluded (SHOW ROLES only returns account
# roles), so ROLE_TYPE is constant ("ROLE") and ROLE_DATABASE_NAME is null.
# ---------------------------------------------------------------------------
COMPARE_ATTRS = [
    "ROLE_TYPE",
    "ROLE_DATABASE_NAME",
    "COMMENT",
    # "OWNER",
    # "OWNER_ROLE_TYPE",
    "IS_FROM_ORGANIZATION_USER_GROUP",
    # Derived from SHOW GRANTS TO ROLE — sorted comma-separated list of parent roles
    # "PARENT_ROLES",
]

# Roles that exist in every Snowflake account by default — ignored.
SYSTEM_ROLES = {"AUTO_FULFILLMENT_EXECUTOR", "MLADMIN", "MLCONSUMER"}


def _private_key_bytes() -> bytes | None:
    """Load a private key from disk if configured."""
    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]:
    """Build snowflake.connector connection params for one account."""

    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 _ci(row: dict, key: str) -> Any:
    """Case-insensitive lookup — SHOW commands return lowercase keys, queries return upper."""
    return row.get(key) if key in row else row.get(key.lower())


def _show_grants_parent_roles(conn: "snowflake.connector.SnowflakeConnection", role_name: str) -> set[str]:
    """Return parent roles (USAGE on ROLE granted to this role) for a single role."""
    cur = conn.cursor(snowflake.connector.DictCursor)
    cur.execute(f'SHOW GRANTS TO ROLE "{role_name}"')
    parents: set[str] = set()
    for row in cur.fetchall():
        if (
            _ci(row, "PRIVILEGE") == "USAGE"
            and _ci(row, "GRANTED_ON") == "ROLE"
            and _ci(row, "GRANTED_TO") == "ROLE"
        ):
            parents.add(_ci(row, "NAME"))
    return parents


def _fetch_parent_roles(
    conn_params: dict[str, Any], names: list[str], label: str
) -> dict[str, set[str]]:
    """Parallel SHOW GRANTS TO ROLE for every role; returns {name: {parent_names}}."""
    if not names:
        return {}

    print(f"  Fetching parent-role grants for {len(names):,} roles in {label}…", flush=True)
    results: dict[str, set[str]] = {}

    def _fetch_one(name: str) -> tuple[str, set[str]]:
        c = snowflake.connector.connect(**conn_params)
        try:
            return name, _show_grants_parent_roles(c, name)
        finally:
            c.close()

    with ThreadPoolExecutor(max_workers=20) as pool:
        futures = {pool.submit(_fetch_one, n): n for n in names}
        done = 0
        for future in as_completed(futures):
            name, parents = future.result()
            results[name] = parents
            done += 1
            if done % 100 == 0:
                print(f"    …{done}/{len(names)}", flush=True)

    return results


def fetch_roles(account_name: str, env_prefix: str, label: str, default_warehouse: str) -> dict[str, dict]:
    """Connect to `account_name` and return {role_key: {attr: value}} dict.

    Uses real-time SHOW ROLES (and SHOW GRANTS TO ROLE per role if PARENT_ROLES is
    enabled) instead of the SNOWFLAKE.ACCOUNT_USAGE views, which lag by up to ~2h.
    """
    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)
        cur.execute("SHOW ROLES")
        rows = cur.fetchall()
    finally:
        conn.close()

    roles: dict[str, dict] = {}
    for row in rows:
        name = _ci(row, "NAME")
        if name in SYSTEM_ROLES:
            continue
        # SHOW ROLES only returns account-level roles, so ROLE_TYPE/DATABASE_NAME
        # are constant — kept for back-compat with the previous CSV/report shape.
        roles[name] = {
            "ROLE_TYPE": "ROLE",
            "ROLE_DATABASE_NAME": None,
            "COMMENT": _ci(row, "COMMENT"),
            "IS_FROM_ORGANIZATION_USER_GROUP": _ci(row, "IS_FROM_ORGANIZATION_USER_GROUP"),
        }

    print(f"  → {len(roles):,} roles found in {label}", flush=True)

    if "PARENT_ROLES" in COMPARE_ATTRS:
        parents = _fetch_parent_roles(params, list(roles), label)
        for name, parent_set in parents.items():
            roles[name]["PARENT_ROLES"] = (
                ", ".join(sorted(parent_set)) if parent_set else None
            )

    return roles


def compare(
        orchard: dict[str, dict], delphi: dict[str, dict]
) -> tuple[list, list]:
    """
    Returns:
        common_diffs  - list of (key, list_of_(attr, orchard_val, delphi_val))
        delphi_only   - list of keys
    """
    orchard_keys = set(orchard)
    delphi_keys = set(delphi)

    delphi_only = sorted(delphi_keys - orchard_keys)

    common_diffs = []
    for key in sorted(orchard_keys & delphi_keys):
        diffs = []
        for attr in COMPARE_ATTRS:
            o_val = orchard[key].get(attr)
            d_val = delphi[key].get(attr)
            if str(o_val) != str(d_val):
                diffs.append((attr, o_val, d_val))
        if diffs:
            common_diffs.append((key, diffs))

    return common_diffs, delphi_only


def print_report(common_diffs, delphi_only):
    sep = "=" * 72

    print(f"\n{sep}")
    print("SNOWFLAKE ROLE COMPARISON: ORCHARD vs DELPHI")
    print(sep)

    # Split attribute diffs from pure parent-role diffs for readability
    attr_diffs = [(k, d) for k, d in common_diffs if any(a != "PARENT_ROLES" for a, _, _ in d)]
    parent_only_diffs = [(k, d) for k, d in common_diffs if all(a == "PARENT_ROLES" for a, _, _ in d)]

    print(f"\nROLES WITH ATTRIBUTE DIFFERENCES ({len(attr_diffs)} role(s)):\n")
    for key, diffs in attr_diffs:
        print(f"  {key}")
        print(f"    {'ATTRIBUTE':<40} {'ORCHARD':<35} {'DELPHI'}")
        print(f"    {'-'*40} {'-'*35} {'-'*35}")
        for attr, o_val, d_val in diffs:
            o_str = _truncate(str(o_val), 34)
            d_str = _truncate(str(d_val), 34)
            print(f"    {attr:<40} {o_str:<35} {d_str}")
        print()

    print(f"ROLES WITH DIFFERENT PARENT-ROLE GRANTS ({len(parent_only_diffs)} role(s)):\n")
    for key, diffs in parent_only_diffs:
        for _, o_val, d_val in diffs:
            o_parents = set(o_val.split(", ")) if o_val else set()
            d_parents = set(d_val.split(", ")) if d_val else set()
            only_orchard = sorted(o_parents - d_parents)
            only_delphi = sorted(d_parents - o_parents)
            print(f"  {key}")
            if only_orchard:
                print(f"    only in Orchard : {', '.join(only_orchard)}")
            if only_delphi:
                print(f"    only in Delphi  : {', '.join(only_delphi)}")
        print()

    print(f"ROLES ONLY IN DELPHI ({len(delphi_only)} role(s)):")
    if delphi_only:
        for chunk in _chunks(delphi_only, 4):
            print("  " + "  ".join(chunk))
    else:
        print("  (none)")

    print(f"\n{sep}\n")


def write_csv(path: str, common_diffs, delphi_only):
    with open(path, "w", newline="") as fh:
        writer = csv.writer(fh)
        writer.writerow(
            ["role_key", "category", "attribute", "orchard_value", "delphi_value"]
        )
        for key, diffs in common_diffs:
            for attr, o_val, d_val in diffs:
                writer.writerow([key, "DIFF", attr, o_val, d_val])
        for key in delphi_only:
            writer.writerow([key, "DELPHI_ONLY", "", "", ""])
    print(f"CSV written to: {path}")


def _truncate(s: str, maxlen: int) -> str:
    return s if len(s) <= maxlen else s[: maxlen - 1] + "…"


def _chunks(lst, n):
    for i in range(0, len(lst), n):
        yield lst[i : i + n]


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--csv", metavar="FILE", help="also write results to CSV")
    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()

    orchard_roles = fetch_roles(args.orchard_account, "ORCHARD_", "Orchard", "DEV_OWS_WAREHOUSE")
    delphi_roles = fetch_roles(args.delphi_account, "DELPHI_", "Delphi", "DEV_OWS_WH")

    common_diffs, delphi_only = compare(orchard_roles, delphi_roles)
    print_report(common_diffs, delphi_only)

    if args.csv:
        write_csv(args.csv, common_diffs, delphi_only)


if __name__ == "__main__":
    main()