#!/usr/bin/env python3
from __future__ import annotations

import argparse
import csv
import gzip
import json
from pathlib import Path
from typing import Any

import pymysql


INSERT_COLUMNS = [
    "Id",
    "Name",
    "Label",
    "IsCustom",
    "ProfileId",
    "IsOwnedByProfile",
    "CreatedDate",
    "LastModifiedDate",
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Supplement missing Salesforce PermissionSet rows into TiDB without truncating local/system rows.")
    parser.add_argument("--snapshot-dir", required=True)
    parser.add_argument("--db-host", required=True)
    parser.add_argument("--db-port", type=int, default=4000)
    parser.add_argument("--db-user", required=True)
    parser.add_argument("--db-password", required=True)
    parser.add_argument("--db-database", required=True)
    return parser.parse_args()


def normalize_bool(raw: str) -> int:
    return 1 if str(raw).strip().lower() == "true" else 0


def normalize_datetime(raw: str) -> str | None:
    value = str(raw or "").strip()
    if not value:
        return None
    return value.replace("T", " ").replace(".000+0000", "").replace(".000Z", "")


def load_source_rows(snapshot_dir: Path) -> dict[str, dict[str, Any]]:
    rows: dict[str, dict[str, Any]] = {}
    with gzip.open(snapshot_dir / "PermissionSet.csv.gz", "rt", encoding="utf-8-sig", newline="") as src:
        for row in csv.DictReader(src):
            row_id = (row.get("Id") or "").strip()
            if not row_id:
                continue
            rows[row_id] = {
                "Id": row_id,
                "Name": (row.get("Name") or "").strip(),
                "Label": (row.get("Label") or "").strip(),
                "IsCustom": normalize_bool(row.get("IsCustom") or ""),
                "ProfileId": (row.get("ProfileId") or "").strip() or None,
                "IsOwnedByProfile": normalize_bool(row.get("IsOwnedByProfile") or ""),
                "CreatedDate": normalize_datetime(row.get("CreatedDate") or ""),
                "LastModifiedDate": normalize_datetime(row.get("LastModifiedDate") or ""),
            }
    return rows


def main() -> int:
    args = parse_args()
    snapshot_dir = Path(args.snapshot_dir).expanduser().resolve()
    source_rows = load_source_rows(snapshot_dir)

    conn = pymysql.connect(
        host=args.db_host,
        port=args.db_port,
        user=args.db_user,
        password=args.db_password,
        database=args.db_database,
        ssl={"ssl": {}},
        charset="utf8mb4",
        autocommit=False,
    )

    report: dict[str, Any] = {"updated_existing": 0, "inserted": 0, "skipped_existing": 0}
    try:
        with conn.cursor() as cur:
            source_ids = sorted(source_rows.keys())

            # Normalize already-loaded public permission sets to use NULL ProfileId,
            # matching Salesforce semantics instead of a collapsing empty-string key.
            cur.execute(
                "UPDATE `PermissionSet` SET `ProfileId` = NULL "
                "WHERE `Id` IN (%s) AND `IsOwnedByProfile` = 0 AND `ProfileId` = ''"
                % ",".join(["%s"] * len(source_ids)),
                source_ids,
            )
            report["updated_existing"] = cur.rowcount

            cur.execute("SELECT Id FROM `PermissionSet`")
            existing_ids = {row[0] for row in cur.fetchall()}

            placeholders = ", ".join(["%s"] * len(INSERT_COLUMNS))
            insert_sql = (
                f"INSERT INTO `PermissionSet` ({', '.join(f'`{column}`' for column in INSERT_COLUMNS)}) "
                f"VALUES ({placeholders})"
            )

            for row_id, row in source_rows.items():
                if row_id in existing_ids:
                    report["skipped_existing"] += 1
                    continue
                values = [row[column] for column in INSERT_COLUMNS]
                cur.execute(insert_sql, values)
                existing_ids.add(row_id)
                report["inserted"] += 1

        conn.commit()
    finally:
        conn.close()

    print(json.dumps(report, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())