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

Produces a report of:
  - Warehouses present in both accounts with differing attributes or parameters
  - Warehouses only in Delphi

Warehouse-level parameters (MAX_CONCURRENCY_LEVEL, STATEMENT_TIMEOUT_IN_SECONDS,
STATEMENT_QUEUED_TIMEOUT_IN_SECONDS) are fetched via parallel SHOW PARAMETERS
calls and only compared when explicitly set at the warehouse level.

Usage:
  pip install snowflake-connector-python
  python3 compare-snowflake-warehouses.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: SYSADMIN)
  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 from SHOW WAREHOUSES to compare between accounts.
# Excluded: state, started_clusters, running, queued, actives, pendings,
#           failed, suspended (runtime metrics); is_default, is_current
#           (session-specific); available, provisioning, quiescing, other
#           (transient state); created_on, resumed_on, updated_on, uuid
#           (timestamps / internal identifiers).
# ---------------------------------------------------------------------------
COMPARE_ATTRS = [
    "type",
    "size",
    "min_cluster_count",
    "max_cluster_count",
    "auto_suspend",
    "auto_resume",
    # "owner",
    # "owner_role_type",
    "comment",
    "scaling_policy",
    "enable_query_acceleration",
    "query_acceleration_max_scale_factor",
    # "resource_monitor",
    "resource_constraint",
    "warehouse_credit_limit",
    # Warehouse-level parameters — injected after the bulk query (see fetch_warehouse_params).
    # Only included when explicitly set at the warehouse level (level = 'WAREHOUSE').
    "param:MAX_CONCURRENCY_LEVEL",
    "param:STATEMENT_TIMEOUT_IN_SECONDS",
    "param:STATEMENT_QUEUED_TIMEOUT_IN_SECONDS",
]


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 "SYSADMIN",
        "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_warehouse_params(
        conn_params: dict[str, Any],
        warehouse_names: list[str],
        label: str,
) -> dict[str, dict[str, str]]:
    """
    Fetch explicitly-set warehouse-level parameters for all warehouses in parallel.
    Returns {warehouse_name: {param_key: value}} — only params with level='WAREHOUSE'.
    """
    print(
        f"  Fetching parameters for {len(warehouse_names):,} warehouses in {label}…",
        flush=True,
    )

    def _fetch_one(name: str) -> tuple[str, dict[str, str]]:
        conn = snowflake.connector.connect(**conn_params)
        try:
            cur = conn.cursor(snowflake.connector.DictCursor)
            cur.execute(f"SHOW PARAMETERS IN WAREHOUSE {name}")
            rows = cur.fetchall()
        finally:
            conn.close()
        # Only include parameters explicitly set at the warehouse level
        return name, {
            r["key"]: r["value"]
            for r in rows
            if r.get("level") == "WAREHOUSE"
        }

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

    return results


def fetch_warehouses(account_name: str, env_prefix: str, label: str, default_warehouse: str) -> dict[str, dict]:
    """
    Connect to `account_name` and return {warehouse_name: {attr: value}} dict.
    Warehouse-level parameters are fetched via parallel SHOW PARAMETERS calls.
    """
    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 WAREHOUSES")
        rows = cur.fetchall()
    finally:
        conn.close()

    warehouses: dict[str, dict] = {}
    for row in rows:
        name = row["name"]
        warehouses[name] = {attr: row.get(attr) for attr in COMPARE_ATTRS if not attr.startswith("param:")}
        warehouses[name]["param:MAX_CONCURRENCY_LEVEL"] = None
        warehouses[name]["param:STATEMENT_TIMEOUT_IN_SECONDS"] = None
        warehouses[name]["param:STATEMENT_QUEUED_TIMEOUT_IN_SECONDS"] = None

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

    # Enrich with warehouse-level parameters
    wh_params = fetch_warehouse_params(params, list(warehouses.keys()), label)
    for name, wh_p in wh_params.items():
        for key, value in wh_p.items():
            warehouses[name][f"param:{key}"] = value

    return warehouses


def compare(
        orchard: dict[str, dict], delphi: dict[str, dict]
) -> tuple[list, list]:
    """
    Returns:
        common_diffs  - list of (name, list_of_(attr, orchard_val, delphi_val))
        delphi_only   - list of names
    """
    orchard_names = set(orchard)
    delphi_names = set(delphi)

    delphi_only = sorted(delphi_names - orchard_names)

    common_diffs = []
    for name in sorted(orchard_names & delphi_names):
        diffs = []
        for attr in COMPARE_ATTRS:
            o_val = orchard[name].get(attr)
            d_val = delphi[name].get(attr)

            # Ignore diffs where generation is "1" in one account and None in another
            if attr == 'resource_constraint' and o_val in ('STANDARD_GEN_1', None) and d_val in ('STANDARD_GEN_1', None):
                continue

            if str(o_val) != str(d_val):
                diffs.append((attr, o_val, d_val))
        if diffs:
            common_diffs.append((name, diffs))

    return common_diffs, delphi_only


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

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

    print(f"\nWAREHOUSES WITH ATTRIBUTE DIFFERENCES ({len(common_diffs)} warehouse(s)):\n")
    for name, diffs in common_diffs:
        print(f"  {name}")
        print(f"    {'ATTRIBUTE':<42} {'ORCHARD':<35} {'DELPHI'}")
        print(f"    {'-'*42} {'-'*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:<42} {o_str:<35} {d_str}")
        print()

    print(f"WAREHOUSES ONLY IN DELPHI ({len(delphi_only)} warehouse(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(
            ["warehouse", "category", "attribute", "orchard_value", "delphi_value"]
        )
        for name, diffs in common_diffs:
            for attr, o_val, d_val in diffs:
                writer.writerow([name, "DIFF", attr, o_val, d_val])
        for name in delphi_only:
            writer.writerow([name, "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_warehouses = fetch_warehouses(args.orchard_account, "ORCHARD_", "Orchard", "DEV_OWS_WAREHOUSE")
    delphi_warehouses = fetch_warehouses(args.delphi_account, "DELPHI_", "Delphi", "DEV_OWS_WH")

    common_diffs, delphi_only = compare(orchard_warehouses, delphi_warehouses)
    print_report(common_diffs, delphi_only)

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


if __name__ == "__main__":
    main()