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

import argparse
import csv
import gzip
import json
import os
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import pymysql


DEFAULT_QUERY_BATCH_SIZE = 500
DEFAULT_APPLY_BATCH_SIZE = 100
DEFAULT_SUMMARY_NAME = "lead_field_repair_summary.json"

TEXT_FIELDS = (
    "HubSpot_Original_Source_Drill_Down_1__c",
    "CLDR__Active_Leads__c",
    "Lead_Status_Formula_Hidden__c",
)
BOOL_FIELD = "awsapn__Eligible_to_Sync_with_APN__c"


@dataclass
class DbConfig:
    host: str
    port: int
    user: str
    password: str
    database: str


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Repair selected Lead fields in TiDB from a frozen Lead snapshot using missing-only updates."
    )
    parser.add_argument("--snapshot-file", required=True, help="Path to Lead.csv.gz.")
    parser.add_argument("--query-batch-size", type=int, default=DEFAULT_QUERY_BATCH_SIZE)
    parser.add_argument("--apply-batch-size", type=int, default=DEFAULT_APPLY_BATCH_SIZE)
    parser.add_argument("--summary-json", help="Optional summary output path.")
    parser.add_argument("--apply", action="store_true", help="Apply updates. Default is dry run.")
    parser.add_argument("--db-host", default=os.environ.get("TIDB_PROD_HOST", os.environ.get("TIDB_HOST", "")))
    parser.add_argument("--db-port", type=int, default=int(os.environ.get("TIDB_PROD_PORT", os.environ.get("TIDB_PORT", "4000"))))
    parser.add_argument("--db-user", default=os.environ.get("TIDB_PROD_USER", os.environ.get("TIDB_USER", "")))
    parser.add_argument("--db-password", default=os.environ.get("TIDB_PROD_PASSWORD", os.environ.get("TIDB_PASSWORD", "")))
    parser.add_argument("--db-database", default=os.environ.get("TIDB_PROD_DATABASE", os.environ.get("TIDB_DATABASE", "")))
    return parser.parse_args()


def load_snapshot_rows(path: Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    seen: set[str] = set()
    with gzip.open(path, "rt", encoding="utf-8-sig", newline="") as src:
        reader = csv.DictReader(src)
        for row in reader:
            row_id = str(row.get("Id") or "").strip()
            if not row_id or row_id in seen:
                continue
            seen.add(row_id)
            rows.append(
                {
                    "Id": row_id,
                    TEXT_FIELDS[0]: normalize_text(row.get(TEXT_FIELDS[0])),
                    TEXT_FIELDS[1]: normalize_text(row.get(TEXT_FIELDS[1])),
                    TEXT_FIELDS[2]: normalize_text(row.get(TEXT_FIELDS[2])),
                    BOOL_FIELD: normalize_bool(row.get(BOOL_FIELD)),
                }
            )
    return rows


def normalize_text(value: Any) -> str | None:
    if value is None:
        return None
    text = str(value)
    return text if text != "" else None


def normalize_bool(value: Any) -> int:
    if value is None:
        return 0
    text = str(value).strip().lower()
    return 1 if text in {"1", "true", "t", "yes"} else 0


def batched[T](values: list[T], size: int) -> list[list[T]]:
    return [values[index : index + size] for index in range(0, len(values), size)]


def connect_db(db: DbConfig) -> pymysql.connections.Connection:
    return pymysql.connect(
        host=db.host,
        port=db.port,
        user=db.user,
        password=db.password,
        database=db.database,
        ssl={},
        autocommit=False,
        cursorclass=pymysql.cursors.DictCursor,
    )


def log(message: str) -> None:
    print(f"[lead-repair] {message}", file=sys.stderr, flush=True)


def fetch_target_rows(conn: pymysql.connections.Connection, ids: list[str]) -> dict[str, dict[str, Any]]:
    placeholders = ",".join(["%s"] * len(ids))
    sql = (
        "SELECT `Id`, "
        f"`{TEXT_FIELDS[0]}`, `{TEXT_FIELDS[1]}`, `{TEXT_FIELDS[2]}`, `{BOOL_FIELD}` "
        f"FROM `lead` WHERE `Id` IN ({placeholders})"
    )
    with conn.cursor() as cur:
        cur.execute(sql, ids)
        rows = cur.fetchall()
    return {str(row["Id"]): row for row in rows}


def build_patch(source_row: dict[str, Any], target_row: dict[str, Any] | None) -> dict[str, Any] | None:
    if target_row is None:
        return None
    updates: dict[str, Any] = {}
    for field in TEXT_FIELDS:
        source_value = source_row[field]
        target_value = normalize_text(target_row.get(field))
        if source_value is not None and target_value is None:
            updates[field] = source_value
    source_bool = int(source_row[BOOL_FIELD] or 0)
    target_bool = 1 if int(target_row.get(BOOL_FIELD) or 0) == 1 else 0
    if source_bool == 1 and target_bool != 1:
        updates[BOOL_FIELD] = 1
    if not updates:
        return None
    return {"Id": source_row["Id"], "updates": updates}


def apply_patch_batch(conn: pymysql.connections.Connection, patches: list[dict[str, Any]]) -> int:
    if not patches:
        return 0
    ids = [str(patch["Id"]) for patch in patches]
    params: list[Any] = []
    assignments: list[str] = []
    for field in (*TEXT_FIELDS, BOOL_FIELD):
        field_patches = [patch for patch in patches if field in patch["updates"]]
        if not field_patches:
            continue
        clauses: list[str] = []
        for patch in field_patches:
            clauses.append("WHEN %s THEN %s")
            params.extend([patch["Id"], patch["updates"][field]])
        assignments.append(
            f"`{field}` = CASE `Id` {' '.join(clauses)} ELSE `{field}` END"
        )
    if not assignments:
        return 0
    where_placeholders = ",".join(["%s"] * len(ids))
    params.extend(ids)
    sql = f"UPDATE `lead` SET {', '.join(assignments)} WHERE `Id` IN ({where_placeholders})"
    with conn.cursor() as cur:
        cur.execute(sql, params)
        total = cur.rowcount
    conn.commit()
    return total


def make_summary(
    *,
    snapshot_file: Path,
    snapshot_count: int,
    scanned_target_count: int,
    missing_target_count: int,
    planned_patches: list[dict[str, Any]],
    applied_rows: int,
    apply: bool,
    query_batch_size: int,
    apply_batch_size: int,
) -> dict[str, Any]:
    field_counts = {
        TEXT_FIELDS[0]: 0,
        TEXT_FIELDS[1]: 0,
        TEXT_FIELDS[2]: 0,
        BOOL_FIELD: 0,
    }
    for patch in planned_patches:
        for field in patch["updates"]:
            field_counts[field] += 1
    return {
        "snapshot_file": str(snapshot_file),
        "snapshot_count": snapshot_count,
        "scanned_target_count": scanned_target_count,
        "missing_target_count": missing_target_count,
        "planned_record_count": len(planned_patches),
        "planned_field_counts": field_counts,
        "applied_rows": applied_rows,
        "apply": apply,
        "query_batch_size": query_batch_size,
        "apply_batch_size": apply_batch_size,
        "sample_patches": planned_patches[:20],
    }


def write_summary(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, ensure_ascii=True), encoding="utf-8")


def main() -> int:
    args = parse_args()
    if args.query_batch_size < 1:
        raise SystemExit("--query-batch-size must be >= 1")
    if args.apply_batch_size < 1:
        raise SystemExit("--apply-batch-size must be >= 1")

    db = DbConfig(
        host=args.db_host,
        port=args.db_port,
        user=args.db_user,
        password=args.db_password,
        database=args.db_database,
    )
    if not all([db.host, db.user, db.password, db.database]):
        raise SystemExit("db host/user/password/database are required")

    snapshot_file = Path(args.snapshot_file).expanduser().resolve()
    snapshot_rows = load_snapshot_rows(snapshot_file)
    log(f"loaded snapshot rows={len(snapshot_rows)} file={snapshot_file}")

    planned_patches: list[dict[str, Any]] = []
    scanned_target_count = 0
    missing_target_count = 0

    conn = connect_db(db)
    try:
        query_chunks = batched(snapshot_rows, args.query_batch_size)
        for index, chunk in enumerate(query_chunks, start=1):
            target_by_id = fetch_target_rows(conn, [row["Id"] for row in chunk])
            scanned_target_count += len(target_by_id)
            missing_target_count += len(chunk) - len(target_by_id)
            for source_row in chunk:
                patch = build_patch(source_row, target_by_id.get(source_row["Id"]))
                if patch is not None:
                    planned_patches.append(patch)
            if index == 1 or index % 10 == 0 or index == len(query_chunks):
                log(
                    f"scan chunk={index}/{len(query_chunks)} scanned_target={scanned_target_count} "
                    f"planned_records={len(planned_patches)}"
                )

        applied_rows = 0
        if args.apply and planned_patches:
            apply_chunks = batched(planned_patches, args.apply_batch_size)
            for index, patch_chunk in enumerate(apply_chunks, start=1):
                applied_rows += apply_patch_batch(conn, patch_chunk)
                if index == 1 or index % 20 == 0 or index == len(apply_chunks):
                    log(f"apply chunk={index}/{len(apply_chunks)} applied_rows={applied_rows}")

        summary = make_summary(
            snapshot_file=snapshot_file,
            snapshot_count=len(snapshot_rows),
            scanned_target_count=scanned_target_count,
            missing_target_count=missing_target_count,
            planned_patches=planned_patches,
            applied_rows=applied_rows,
            apply=args.apply,
            query_batch_size=args.query_batch_size,
            apply_batch_size=args.apply_batch_size,
        )
    finally:
        conn.close()

    if args.summary_json:
        write_summary(Path(args.summary_json).expanduser().resolve(), summary)
    else:
        default_summary = snapshot_file.parent / DEFAULT_SUMMARY_NAME
        write_summary(default_summary, summary)
    print(json.dumps(summary, indent=2, ensure_ascii=True))
    return 0


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