"""Generate attach (PC-level) and detach (CB-level) backfill folders from cb-roles.csv.

Snowflake query
SELECT 
  ir.identity_uuid,
  ir.tenant_uuid,
  ir.tenant_type,
  ir.roles
FROM PERMISSIONS_PLATFORM.PROD.PP_IDENTITY ir
WHERE ir.tenant_type = 'company_brand'
AND (
  ARRAY_CONTAINS(OBJECT_CONSTRUCT('role', 'contract_viewer')::VARIANT, ir.roles)
  OR  ARRAY_CONTAINS(OBJECT_CONSTRUCT('role', 'account_admin')::VARIANT, ir.roles)
  OR ARRAY_CONTAINS(OBJECT_CONSTRUCT('role', 'contract_admin')::VARIANT, ir.roles)
  OR ARRAY_CONTAINS(OBJECT_CONSTRUCT('role', 'account_edit')::VARIANT, ir.roles)
);

Reads the Snowflake export cb-roles.csv where each row has:
    IDENTITY_UUID, TENANT_UUID, TENANT_TYPE, ROLES (JSON array)

Creates two folders, each with a CSV and manifest.json:
    - <name>-attach/  (PC-level attach rows)
    - <name>-detach/  (CB-level detach rows)

Usage:
    python generate_cb_to_pc_csv.py <folder_name>
    # e.g. python generate_cb_to_pc_csv.py cb_to_pc
    # produces: cb_to_pc-attach/ and cb_to_pc-detach/
"""

import argparse
import csv
import json
import os

ORCHARD_PARENT_COMPANY_UUID = "955a1bbd-b623-4ea1-ab5f-8d6620c442fb"
SME_PARENT_COMPANY_UUID = "f1594122-7f99-4916-b103-08b0444c7b46"

HEADER = ["identity_uuid", "tenant_uuid", "tenant_type", "role", "operation"]


def parse_roles(roles_str: str) -> list[str]:
    """Parse the JSON roles column from Snowflake export."""
    data = json.loads(roles_str)
    return [entry["role"] for entry in data]


def write_folder(folder: str, csv_name: str, rows: list[list[str]], bucket: str) -> None:
    """Write a CSV and manifest.json into the given folder."""
    os.makedirs(folder, exist_ok=True)

    csv_path = os.path.join(folder, csv_name)
    with open(csv_path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(HEADER)
        writer.writerows(rows)

    manifest = {
        "bucket": bucket,
        "jobs": [
            {
                "job_type": "attach_and_detach",
                "keys": [f"{folder}/{csv_name}"],
            }
        ],
    }
    manifest_path = os.path.join(folder, "manifest.json")
    with open(manifest_path, "w") as f:
        json.dump(manifest, f, indent=4)
        f.write("\n")


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate PC attach / CB detach backfill CSVs")
    parser.add_argument("name", help="Base folder name (creates <name>-attach and <name>-detach)")
    parser.add_argument("--bucket", default="prod-pdp-backfill", help="S3 bucket name (default: prod-pdp-backfill)")
    parser.add_argument("--input", default="cb-roles.csv", help="Input CSV from Snowflake (default: cb-roles.csv)")
    args = parser.parse_args()

    attach_folder = f"{args.name}-attach"
    detach_folder = f"{args.name}-detach"

    # Track unique identity+role combos for attach (avoid duplicates across CBs)
    attach_seen: set[tuple[str, str]] = set()
    attach_rows: list[list[str]] = []
    detach_rows: list[list[str]] = []

    with open(args.input, newline="") as f:
        reader = csv.DictReader(f)
        for row in reader:
            identity_uuid = row["IDENTITY_UUID"].strip()
            cb_uuid = row["TENANT_UUID"].strip()
            roles = parse_roles(row["ROLES"])

            for role in roles:
                # Detach: one row per CB assignment
                detach_rows.append([identity_uuid, cb_uuid, "company_brand", role, "detach"])

                # Attach: one row per PC, but only once per identity+role
                key = (identity_uuid, role)
                if key not in attach_seen:
                    attach_seen.add(key)
                    attach_rows.append(
                        [identity_uuid, ORCHARD_PARENT_COMPANY_UUID, "parent_company", role, "attach"]
                    )
                    attach_rows.append(
                        [identity_uuid, SME_PARENT_COMPANY_UUID, "parent_company", role, "attach"]
                    )

    # Sort for deterministic output
    attach_rows.sort()
    detach_rows.sort()

    write_folder(attach_folder, "attach.csv", attach_rows, args.bucket)
    write_folder(detach_folder, "detach.csv", detach_rows, args.bucket)

    unique_identities = {r[0] for r in attach_rows}
    print(f"Unique identities: {len(unique_identities)}")
    print(f"Attach rows: {len(attach_rows)} -> {attach_folder}/")
    print(f"Detach rows: {len(detach_rows)} -> {detach_folder}/")


if __name__ == "__main__":
    main()