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

import argparse
import csv
import json
import os
import shutil
import subprocess
import sys
import urllib.parse
import urllib.request
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from export_salesforce_snapshot import SalesforceClient, load_export_config, load_kv_config
from load_salesforce_snapshot_to_tidb import build_export_object_map


DEFAULT_EXPORT_CONFIG = "configs/daily_full_export.json"
DEFAULT_MONITOR_CONFIG = "configs/prod_reconcile_monitor.json"
DEFAULT_SF_CONFIG = "~/.salesforce/readonly_config"
MYSQL_CLIENT_CANDIDATES = [
    shutil.which("mysql"),
    "/usr/bin/mysql",
    "/usr/local/bin/mysql",
    "/opt/homebrew/opt/mysql-client@8.4/bin/mysql",
]


def first_existing(items: list[str | None]) -> str | None:
    for item in items:
        if item and Path(item).exists():
            return item
    return None


MYSQL_CLIENT = first_existing(MYSQL_CLIENT_CANDIDATES)
MYSQL_SSL_MODE_SUPPORT_CACHE: dict[str, bool] = {}


def log(message: str) -> None:
    timestamp = datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
    print(f"[reconcile] {timestamp} {message}", file=sys.stderr, flush=True)


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


@dataclass
class SalesforceRuntimeConfig:
    config_path: Path
    api_version: str


def parse_args() -> argparse.Namespace:
    script_dir = Path(__file__).resolve().parent
    parser = argparse.ArgumentParser(
        description="Compare live Salesforce data against TiDB production per monitored object."
    )
    parser.add_argument(
        "--monitor-config",
        default=str(script_dir / DEFAULT_MONITOR_CONFIG),
        help="JSON config listing monitored objects and allowed prod-only extras.",
    )
    parser.add_argument(
        "--export-config",
        default=str(script_dir / DEFAULT_EXPORT_CONFIG),
        help="Export config used to resolve object -> target table mapping.",
    )
    parser.add_argument("--object", action="append", help="Optional object allowlist. Repeatable.")
    parser.add_argument("--sample-limit", type=int, default=20, help="Max mismatch IDs included per object.")
    parser.add_argument("--output-json", help="Optional path to write reconciliation JSON.")
    parser.add_argument(
        "--sf-config",
        default=os.environ.get("SF_READONLY_CONFIG_PATH", DEFAULT_SF_CONFIG),
        help="Path to the live Salesforce readonly config.",
    )
    parser.add_argument(
        "--sf-api-version",
        default=os.environ.get("SF_API_VERSION", "v63.0"),
        help="Salesforce API version used for live queries.",
    )
    parser.add_argument("--mysql-client", default=MYSQL_CLIENT, help="Path to mysql client.")
    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", "")))
    parser.add_argument("--ssl-mode", default=os.environ.get("TIDB_SSL_MODE", "REQUIRED"))
    return parser.parse_args()


def load_json(path: Path) -> dict[str, Any]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(payload, dict):
        raise ValueError(f"{path} must be a JSON object")
    return payload


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


def sanitize_identifier(name: str) -> str:
    if not name or any(ch not in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789_" for ch in name):
        raise ValueError(f"unsafe identifier: {name!r}")
    return name


def mysql_env(password: str) -> dict[str, str]:
    env = os.environ.copy()
    env["MYSQL_PWD"] = password
    return env


def mysql_supports_ssl_mode(mysql_client: str) -> bool:
    cached = MYSQL_SSL_MODE_SUPPORT_CACHE.get(mysql_client)
    if cached is not None:
        return cached
    result = subprocess.run(
        [mysql_client, "--help"],
        capture_output=True,
        text=True,
        check=False,
        env=mysql_env(""),
    )
    combined = f"{result.stdout}\n{result.stderr}"
    supported = "--ssl-mode" in combined
    MYSQL_SSL_MODE_SUPPORT_CACHE[mysql_client] = supported
    return supported


def mysql_ssl_args(mysql_client: str, ssl_mode: str) -> list[str]:
    normalized = str(ssl_mode or "").strip().upper()
    if not normalized:
        return []
    if mysql_supports_ssl_mode(mysql_client):
        return ["--ssl-mode", normalized]
    if normalized == "DISABLED":
        return []
    return ["--ssl"]


def run_mysql_query(mysql_client: str, db: DbConfig, sql: str) -> list[str]:
    cmd = [
        mysql_client,
        "--host",
        db.host,
        "--port",
        str(db.port),
        "--user",
        db.user,
        "--database",
        db.database,
        "--batch",
        "--raw",
        "--skip-column-names",
        "-e",
        sql,
    ]
    cmd[10:10] = mysql_ssl_args(mysql_client, db.ssl_mode)
    result = subprocess.run(cmd, check=True, capture_output=True, text=True, env=mysql_env(db.password))
    return [line.rstrip("\n") for line in result.stdout.splitlines()]


def require_db_config(args: argparse.Namespace) -> DbConfig:
    db = DbConfig(
        host=args.db_host,
        port=args.db_port,
        user=args.db_user,
        password=args.db_password,
        database=args.db_database,
        ssl_mode=args.ssl_mode,
    )
    if not args.mysql_client:
        raise RuntimeError("mysql client not found")
    if not all([db.host, db.user, db.database]):
        raise ValueError("db host/user/database are required")
    return db


def require_salesforce_runtime_config(args: argparse.Namespace) -> SalesforceRuntimeConfig:
    config_path = Path(args.sf_config).expanduser().resolve()
    if not config_path.exists():
        raise FileNotFoundError(f"salesforce config not found: {config_path}")
    return SalesforceRuntimeConfig(
        config_path=config_path,
        api_version=str(args.sf_api_version).strip() or "v63.0",
    )


def fetch_table_columns(mysql_client: str, db: DbConfig, table: str) -> list[str]:
    safe_table = sanitize_identifier(table)
    sql = (
        "SELECT column_name FROM information_schema.columns "
        f"WHERE table_schema = '{db.database}' AND table_name = '{safe_table}' "
        "ORDER BY ordinal_position"
    )
    return [line.strip() for line in run_mysql_query(mysql_client, db, sql) if line.strip()]


def fetch_target_ids(mysql_client: str, db: DbConfig, table: str) -> set[str]:
    safe_table = sanitize_identifier(table)
    sql = f"SELECT `Id` FROM `{safe_table}`"
    return {line.strip() for line in run_mysql_query(mysql_client, db, sql) if line.strip()}


def load_source_ids_from_csv_bytes(content: bytes) -> set[str]:
    result: set[str] = set()
    decoded = content.decode("utf-8-sig")
    reader = csv.DictReader(decoded.splitlines())
    header = list(reader.fieldnames or [])
    if "Id" not in header:
        raise ValueError("bulk query results missing Id column")
    for row in reader:
        row_id = str(row.get("Id") or "").strip()
        if row_id:
            result.add(row_id)
    return result


def fetch_live_ids_via_rest(client: SalesforceClient, object_name: str) -> set[str]:
    soql = f"SELECT Id FROM {object_name}"
    next_url = client._api(f"/query/?q={urllib.parse.quote(soql, safe='')}")
    result: set[str] = set()
    while next_url:
        req = urllib.request.Request(next_url, headers=client._headers())
        with urllib.request.urlopen(req, timeout=300) as resp:
            payload = json.loads(resp.read().decode("utf-8"))
        for row in payload.get("records", []):
            row_id = str((row or {}).get("Id") or "").strip()
            if row_id:
                result.add(row_id)
        relative_next = str(payload.get("nextRecordsUrl") or "").strip()
        next_url = f"{client.auth.instance_url}{relative_next}" if relative_next else ""
    return result


def fetch_live_source_ids(client: SalesforceClient, object_name: str) -> tuple[set[str], str]:
    try:
        job_id = client.create_bulk_query_job(f"SELECT Id FROM {object_name}")
        client.wait_for_bulk_query_job(job_id, poll_seconds=2, timeout_seconds=1800)
        content = client.download_bulk_query_results(job_id)
        return load_source_ids_from_csv_bytes(content), "bulk"
    except Exception:
        return fetch_live_ids_via_rest(client, object_name), "rest"


def fetch_live_field_names(client: SalesforceClient, object_name: str) -> list[str]:
    describe = client.describe(object_name)
    fields = describe.get("fields", [])
    if not isinstance(fields, list):
        return []
    result: list[str] = []
    for field in fields:
        if not isinstance(field, dict):
            continue
        name = str(field.get("name") or "").strip()
        if name:
            result.append(name)
    return result


def enrich_object_config(config_item: dict[str, Any], export_object_map: dict[str, dict[str, Any]]) -> dict[str, Any]:
    object_name = str(config_item.get("name", "")).strip()
    export_item = export_object_map.get(object_name, {})
    merged = dict(config_item)
    merged["name"] = object_name
    if "target_table" not in merged and export_item.get("target_table"):
        merged["target_table"] = export_item.get("target_table")
    if "priority" not in merged and export_item.get("priority"):
        merged["priority"] = export_item.get("priority")
    if "category" not in merged and export_item.get("category"):
        merged["category"] = export_item.get("category")
    if "required_fields" not in merged and export_item.get("required_fields"):
        merged["required_fields"] = export_item.get("required_fields")
    return merged


def select_reconcile_configs(
    configured_objects: list[dict[str, Any]],
    requested_objects: set[str],
) -> list[dict[str, Any]]:
    available_objects = {
        str(item.get("name", "")).strip()
        for item in configured_objects
        if str(item.get("name", "")).strip()
    }
    unknown_objects = sorted(requested_objects - available_objects)
    if unknown_objects:
        raise ValueError(
            "unknown reconcile object(s): "
            f"{', '.join(unknown_objects)}; "
            f"allowed: {', '.join(sorted(available_objects))}"
        )

    selected_configs = [
        item
        for item in configured_objects
        if not requested_objects
        or str(item.get("name", "")).strip() in requested_objects
    ]
    if not selected_configs:
        raise ValueError("reconcile selected zero monitored objects")
    return selected_configs


def reconcile_object(
    object_config: dict[str, Any],
    client: SalesforceClient,
    mysql_client: str,
    db: DbConfig,
    sample_limit: int,
) -> dict[str, Any]:
    object_name = str(object_config.get("name", "")).strip()
    priority = str(object_config.get("priority", "")).strip() or "UNSPECIFIED"
    category = str(object_config.get("category", "")).strip() or "uncategorized"
    required_fields = [
        str(item).strip()
        for item in object_config.get("required_fields", [])
        if str(item).strip()
    ]
    allowed_extra_ids = {
        str(item).strip()
        for item in object_config.get("allowed_target_extra_ids", [])
        if str(item).strip()
    }
    ignore_unexpected_target_extras = bool(object_config.get("ignore_unexpected_target_extras"))
    target_table = str(object_config.get("target_table", "")).strip()

    log(f"{object_name}: start target_table={target_table} priority={priority}")

    table_columns = fetch_table_columns(mysql_client, db, target_table)
    if "Id" not in table_columns:
        log(f"{object_name}: target table missing or Id column absent")
        return {
            "object": object_name,
            "priority": priority,
            "category": category,
            "status": "missing_table",
            "target_table": target_table,
        }

    log(f"{object_name}: describe fields")
    source_fields = fetch_live_field_names(client, object_name)
    missing_required_fields = sorted(field for field in required_fields if field not in source_fields)
    log(f"{object_name}: fetch live source ids")
    source_ids, source_transport = fetch_live_source_ids(client, object_name)
    log(f"{object_name}: fetch target ids")
    target_ids = fetch_target_ids(mysql_client, db, target_table)
    missing_ids = sorted(source_ids - target_ids)
    extra_ids = sorted(target_ids - source_ids)
    unexpected_extra_ids = sorted(set(extra_ids) - allowed_extra_ids)
    if ignore_unexpected_target_extras:
        unexpected_extra_ids = []
    allowed_extra_ids_present = sorted(set(extra_ids) & allowed_extra_ids)
    missing_allowed_extra_ids = sorted(allowed_extra_ids - target_ids)

    source_count = len(source_ids)
    target_count = len(target_ids)
    effective_target_count = target_count - len(allowed_extra_ids_present)
    if ignore_unexpected_target_extras:
        effective_target_count -= len(set(extra_ids) - allowed_extra_ids)

    status = "ok"
    if missing_ids or unexpected_extra_ids or effective_target_count != source_count or missing_required_fields:
        status = "mismatch"

    log(
        f"{object_name}: done status={status} source={source_count} "
        f"target={target_count} missing={len(missing_ids)} extra={len(unexpected_extra_ids)} transport={source_transport}"
    )

    return {
        "object": object_name,
        "priority": priority,
        "category": category,
        "status": status,
        "source_mode": "live_salesforce",
        "source_transport": source_transport,
        "target_table": target_table,
        "source_field_count": len(source_fields),
        "required_fields_count": len(required_fields),
        "missing_required_fields_count": len(missing_required_fields),
        "source_count": source_count,
        "target_count": target_count,
        "effective_target_count": effective_target_count,
        "missing_ids_count": len(missing_ids),
        "unexpected_extra_ids_count": len(unexpected_extra_ids),
        "allowed_target_extra_ids_present_count": len(allowed_extra_ids_present),
        "missing_required_fields_sample": missing_required_fields[:sample_limit],
        "missing_ids_sample": missing_ids[:sample_limit],
        "unexpected_extra_ids_sample": unexpected_extra_ids[:sample_limit],
        "allowed_target_extra_ids_present_sample": allowed_extra_ids_present[:sample_limit],
        "missing_allowed_target_extra_ids_sample": missing_allowed_extra_ids[:sample_limit],
    }


def main() -> int:
    args = parse_args()
    db = require_db_config(args)
    sf = require_salesforce_runtime_config(args)
    monitor_config = load_json(Path(args.monitor_config).expanduser().resolve())
    export_config = load_export_config(Path(args.export_config).expanduser().resolve())
    export_object_map = build_export_object_map(export_config)
    configured_objects = [
        enrich_object_config(item, export_object_map)
        for item in monitor_config.get("objects", [])
        if isinstance(item, dict) and str(item.get("name", "")).strip()
    ]
    requested_objects = {
        item.strip()
        for item in (args.object or [])
        if item and item.strip()
    }
    selected_configs = select_reconcile_configs(configured_objects, requested_objects)

    readonly_config = load_kv_config(sf.config_path)
    client = SalesforceClient(readonly_config, api_version=sf.api_version)
    log(f"login start org_config={sf.config_path} api_version={sf.api_version}")
    client.login()
    log(f"login complete instance_url={client.auth.instance_url if client.auth else ''}")

    log(f"selected_objects={len(selected_configs)} requested={sorted(requested_objects) if requested_objects else 'all'}")
    results = []
    for index, item in enumerate(selected_configs, start=1):
        object_name = str(item.get("name", "")).strip()
        log(f"progress {index}/{len(selected_configs)} object={object_name}")
        results.append(reconcile_object(item, client, args.mysql_client, db, args.sample_limit))

    priority_order = ["P0", "P1", "P2"]
    grouped: dict[str, dict[str, Any]] = {}
    for result in results:
        priority = str(result.get("priority", "UNSPECIFIED")).strip() or "UNSPECIFIED"
        entry = grouped.setdefault(priority, {"priority": priority, "object_count": 0, "ok_count": 0, "mismatch_count": 0, "objects": []})
        entry["object_count"] += 1
        if result.get("status") == "ok":
            entry["ok_count"] += 1
        else:
            entry["mismatch_count"] += 1
        entry["objects"].append(result)
    priority_summary = [grouped[key] for key in priority_order if key in grouped] + [
        grouped[key] for key in sorted(grouped.keys()) if key not in priority_order
    ]

    missing_only = [
        item for item in results
        if int(item.get("missing_ids_count") or 0) > 0
        and int(item.get("unexpected_extra_ids_count") or 0) == 0
        and int(item.get("missing_required_fields_count") or 0) == 0
    ]
    extra_only = [
        item for item in results
        if int(item.get("missing_ids_count") or 0) == 0
        and int(item.get("unexpected_extra_ids_count") or 0) > 0
    ]
    both_missing_and_extra = [
        item for item in results
        if int(item.get("missing_ids_count") or 0) > 0
        and int(item.get("unexpected_extra_ids_count") or 0) > 0
    ]
    mismatch_count = sum(1 for item in results if item.get("status") != "ok")

    payload = {
        "status": "ok" if all(item.get("status") == "ok" for item in results) else "mismatch",
        "source_mode": "live_salesforce",
        "source_org": client.auth.instance_url if client.auth else "",
        "sf_config": str(sf.config_path),
        "sf_api_version": sf.api_version,
        "monitor_config": str(Path(args.monitor_config).expanduser().resolve()),
        "object_count": len(results),
        "mismatch_count": mismatch_count,
        "missing_only": missing_only,
        "extra_only": extra_only,
        "both_missing_and_extra": both_missing_and_extra,
        "priority_summary": priority_summary,
        "objects": results,
    }
    if args.output_json:
        write_json(Path(args.output_json).expanduser().resolve(), payload)
    log(f"complete status={payload['status']} object_count={len(results)}")
    print(json.dumps(payload, indent=2))
    return 0 if payload["status"] == "ok" else 1


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