"""
Reconciliation script: compare row counts between source schemas in
PROD_SHOPIFY_INTEGRATIONS and rolled-up tables in SHOPIFY_STORES_GLOBAL.QA_SHOPIFY_STORE_SCHEMA.

Key facts determined from prior exploration:
- Destination filter column: STORE_SCHEMA (lowercase schema names)
- All table names are quoted to guard against reserved words
- Destination tables are named identically to source entity names (e.g. ORDER, CUSTOMER)
"""
import warnings
warnings.filterwarnings("ignore")

import snowflake.connector
import os

USER = os.environ.get("SNOWFLAKE_USER", "bill.fleischer@sonymusic-pde.com")
ROLE = os.environ.get("SNOWFLAKE_ROLE", "FANSIFTER_ENGINEERING_PRIVACY")
WAREHOUSE = os.environ.get("SNOWFLAKE_WAREHOUSE", "EXPLORATION_WH")
ACCOUNT = "delphi.us-east-1"

SOURCE_DB = "PROD_SHOPIFY_INTEGRATIONS"
# Source schemas (uppercase for SHOW TABLES, but lowercase for STORE_SCHEMA filter)
SOURCE_SCHEMAS = [
    "STORE_82477_ORCHARD_ADVERTISING_MARKETING_SE_005",
    "STORE_82477_ORCHARD_ADVERTISING_MARKETING_SE_004",
    "STORE_34544_SME_AFRICA_PTY_LTD_002",
    "STORE_34588_SME_UK_LIMITED_001",
]
# These are the STORE_SCHEMA values in the destination tables (lowercase)
DEST_FILTER_VALUES = [s.lower() for s in SOURCE_SCHEMAS]

DEST_DB = "SHOPIFY_STORES_GLOBAL"
DEST_SCHEMA = "QA_SHOPIFY_STORE_SCHEMA"
FILTER_COL = "STORE_SCHEMA"

print(f"Connecting to Snowflake as {USER} (browser auth)...")
conn = snowflake.connector.connect(
    user=USER,
    account=ACCOUNT,
    authenticator="externalbrowser",
    role=ROLE,
    warehouse=WAREHOUSE,
)
cur = conn.cursor()
print("Connected.\n")


def run(sql):
    cur.execute(sql)
    return cur.fetchall()


def quote_table(name: str) -> str:
    return f'"{name}"'


# ── 1. List tables in each source schema ────────────────────────────────────
print("=" * 60)
print("STEP 1: Tables in source schemas")
print("=" * 60)

source_tables: dict[str, list[str]] = {}
for schema in SOURCE_SCHEMAS:
    rows = run(f"SHOW TABLES IN {SOURCE_DB}.{schema}")
    tables = [r[1].upper() for r in rows]
    source_tables[schema] = tables
    print(f"\n{schema} ({len(tables)} tables):")
    for t in sorted(tables):
        print(f"  {t}")

all_source_table_names = sorted(set(t for tables in source_tables.values() for t in tables))
print(f"\nTotal distinct source table names: {len(all_source_table_names)}")


# ── 2. List tables in destination schema ────────────────────────────────────
print("\n" + "=" * 60)
print("STEP 2: Tables in destination schema")
print("=" * 60)

dest_rows = run(f"SHOW TABLES IN {DEST_DB}.{DEST_SCHEMA}")
dest_tables = sorted([r[1].upper() for r in dest_rows])
print(f"\n{DEST_DB}.{DEST_SCHEMA} ({len(dest_tables)} tables):")
for t in dest_tables:
    print(f"  {t}")


# ── 3. Build source → destination table mapping ──────────────────────────────
# Destination tables are named identically to source entity names (ORDER, CUSTOMER, etc.)
dest_table_map: dict[str, str] = {dest: dest for dest in dest_tables}

print("\n" + "=" * 60)
print("STEP 3: Source entity → destination table mapping")
print("=" * 60)
matched = [(e, dest_table_map[e]) for e in all_source_table_names if e in dest_table_map]
unmatched = [e for e in all_source_table_names if e not in dest_table_map]
print(f"\nMatched ({len(matched)}): {[e for e, _ in matched]}")
print(f"\nNo destination table ({len(unmatched)}): {unmatched}")


# ── 4. Row count reconciliation ──────────────────────────────────────────────
print("\n" + "=" * 60)
print("STEP 4: Counting rows (this may take a few minutes)...")
print("=" * 60)

schema_list = ", ".join(f"'{s}'" for s in DEST_FILTER_VALUES)
results = []

for src_entity in sorted(all_source_table_names):
    dest_table = dest_table_map.get(src_entity)
    quoted_src = quote_table(src_entity)

    # Sum source counts across all schemas that have this table
    source_total = 0
    source_by_schema = {}
    for schema in SOURCE_SCHEMAS:
        if src_entity in source_tables.get(schema, []):
            try:
                count_row = run(f"SELECT COUNT(*) FROM {SOURCE_DB}.{schema}.{quoted_src}")
                cnt = count_row[0][0]
                source_by_schema[schema] = cnt
                source_total += cnt
            except Exception as e:
                print(f"  ERROR counting {schema}.{src_entity}: {e}")
                source_by_schema[schema] = "ERR"
        else:
            source_by_schema[schema] = None

    # Destination count filtered to these 4 store schemas
    dest_total = None
    if dest_table:
        try:
            dest_count_row = run(
                f"SELECT COUNT(*) FROM {DEST_DB}.{DEST_SCHEMA}.{dest_table} "
                f"WHERE {FILTER_COL} IN ({schema_list})"
            )
            dest_total = dest_count_row[0][0]
        except Exception as e:
            print(f"  ERROR counting dest {dest_table}: {e}")
            dest_total = "ERR"

    if isinstance(source_total, int) and isinstance(dest_total, int):
        match_status = "MATCH" if source_total == dest_total else "MISMATCH"
    elif dest_total is None:
        match_status = "NO DEST"
    else:
        match_status = "ERROR"

    results.append({
        "entity": src_entity,
        "dest_table": dest_table or "—",
        "source_total": source_total,
        "dest_total": dest_total,
        "status": match_status,
        "source_by_schema": source_by_schema,
    })
    status_icon = {"MATCH": "OK", "MISMATCH": "!!!!", "NO DEST": "---", "ERROR": "ERR"}.get(match_status, "?")
    print(f"  [{status_icon:4s}] {src_entity:<45} src={source_total:>8}  dest={str(dest_total):>8}")


# ── 5. Print reconciliation table ───────────────────────────────────────────
print("\n\n" + "=" * 95)
print("RECONCILIATION SUMMARY")
print("=" * 95)
print(f"{'Entity':<40} {'Dest Table':<35} {'Source':>8} {'Dest':>8} {'Status':>10}")
print("-" * 95)

for r in results:
    dest_str = str(r["dest_total"]) if r["dest_total"] is not None else "N/A"
    src_str = f"{r['source_total']:,}" if isinstance(r['source_total'], int) else str(r['source_total'])
    print(f"{r['entity']:<40} {r['dest_table']:<35} {src_str:>8} {dest_str:>8} {r['status']:>10}")

# Summary counts
matched_count = sum(1 for r in results if r["status"] == "MATCH")
mismatch_count = sum(1 for r in results if r["status"] == "MISMATCH")
no_dest_count = sum(1 for r in results if r["status"] == "NO DEST")
print(f"\nTotals: {len(results)} source entities | {matched_count} MATCH | {mismatch_count} MISMATCH | {no_dest_count} NO DEST")

if mismatch_count:
    print("\nMISMATCHED ENTITIES:")
    for r in results:
        if r["status"] == "MISMATCH":
            diff = r["source_total"] - (r["dest_total"] if isinstance(r["dest_total"], int) else 0)
            print(f"  {r['entity']}: source={r['source_total']:,}  dest={r['dest_total']:,}  diff={diff:+,}")

# Per-schema breakdown
print("\n\nPer-schema source row counts:")
abbrevs = ["SE_005", "SE_004", "AFR_002", "UK_001"]
print(f"{'Entity':<40}", end="")
for a in abbrevs:
    print(f" {a:>10}", end="")
print()
print("-" * (40 + 11 * len(SOURCE_SCHEMAS)))
for r in results:
    print(f"{r['entity']:<40}", end="")
    for schema in SOURCE_SCHEMAS:
        v = r["source_by_schema"].get(schema)
        s = f"{v:,}" if isinstance(v, int) else ("—" if v is None else str(v))
        print(f" {s:>10}", end="")
    print()

cur.close()
conn.close()
print("\nDone.")