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

import argparse
import json
import os
import subprocess
import tempfile
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from run_prod_repair_cycle import build_env, parse_items, sanitize_component, select_mismatches


DEFAULT_EXPORT_CONFIG = "configs/daily_full_export.json"
DEFAULT_MONITOR_CONFIG = "configs/prod_reconcile_monitor.json"


def parse_args() -> argparse.Namespace:
    script_dir = Path(__file__).resolve().parent
    parser = argparse.ArgumentParser(
        description="Run production repair in parallel by dispatching one run_prod_repair_cycle.py process per object."
    )
    parser.add_argument("--readonly-config", required=True, help="Path to Salesforce readonly config.")
    parser.add_argument("--snapshot-root", required=True, help="Snapshot root for reconcile discovery.")
    parser.add_argument("--output-root", help="Export output root. Defaults to snapshot root.")
    parser.add_argument("--monitor-config", default=str(script_dir / DEFAULT_MONITOR_CONFIG))
    parser.add_argument("--export-config", default=str(script_dir / DEFAULT_EXPORT_CONFIG))
    parser.add_argument("--load-plan-config", action="append", default=[], help="Forwarded to child repair cycle.")
    parser.add_argument("--object", action="append", help="Optional object allowlist. Repeatable.")
    parser.add_argument("--skip-object", action="append", help="Optional object denylist. Repeatable.")
    parser.add_argument("--summary-json", help="Optional path for aggregate summary JSON.")
    parser.add_argument("--state-dir", help="Persistent state directory for resumable per-object repair.")
    parser.add_argument("--retry-count", type=int, default=2)
    parser.add_argument("--max-workers", type=int, default=4)
    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"))
    parser.add_argument("--skip-export", action="store_true")
    parser.add_argument("--skip-load", action="store_true")
    parser.add_argument("--skip-post-reconcile", action="store_true")
    parser.add_argument("--truncate-first", action="store_true", default=True)
    return parser.parse_args()


def utc_now() -> datetime:
    return datetime.now(timezone.utc)


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, sort_keys=True), encoding="utf-8")


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 run_cmd(cmd: list[str], env: dict[str, str], *, check: bool = True) -> subprocess.CompletedProcess[str]:
    return subprocess.run(cmd, env=env, text=True, check=check, capture_output=True)


def discover_mismatches(script_dir: Path, args: argparse.Namespace, env: dict[str, str]) -> tuple[dict[str, Any], list[str]]:
    output_root = Path(args.output_root or args.snapshot_root).expanduser().resolve()
    state_dir = Path(args.state_dir or (output_root / "_repair_state")).expanduser().resolve()
    allow = parse_items(args.object)
    deny = parse_items(args.skip_object)
    with tempfile.TemporaryDirectory(prefix="prod-repair-parallel-") as temp_dir:
        pre_json = Path(temp_dir) / "pre_reconcile.json"
        cmd = [
            "python3",
            str(script_dir / "verify_latest_snapshot_vs_tidb.py"),
            "--snapshot-root",
            str(Path(args.snapshot_root).expanduser().resolve()),
            "--monitor-config",
            str(Path(args.monitor_config).expanduser().resolve()),
            "--export-config",
            str(Path(args.export_config).expanduser().resolve()),
            "--output-json",
            str(pre_json),
        ]
        for object_name in sorted(allow):
            cmd.extend(["--object", object_name])
        run_cmd(cmd, env=env, check=False)
        payload = load_json(pre_json)
    summary = {
        "started_at": utc_now().isoformat(),
        "snapshot_root": str(Path(args.snapshot_root).expanduser().resolve()),
        "output_root": str(output_root),
        "state_dir": str(state_dir),
        "pre_reconcile": payload,
    }
    return summary, select_mismatches(payload, allow, deny)


def build_child_cmd(script_dir: Path, args: argparse.Namespace, object_name: str, state_dir: Path) -> tuple[list[str], Path]:
    object_summary = state_dir / f"{sanitize_component(object_name)}.summary.json"
    cmd = [
        "python3",
        str(script_dir / "run_prod_repair_cycle.py"),
        "--readonly-config",
        str(Path(args.readonly_config).expanduser().resolve()),
        "--snapshot-root",
        str(Path(args.snapshot_root).expanduser().resolve()),
        "--output-root",
        str(Path(args.output_root or args.snapshot_root).expanduser().resolve()),
        "--monitor-config",
        str(Path(args.monitor_config).expanduser().resolve()),
        "--export-config",
        str(Path(args.export_config).expanduser().resolve()),
        "--db-host",
        args.db_host,
        "--db-port",
        str(args.db_port),
        "--db-user",
        args.db_user,
        "--db-password",
        args.db_password,
        "--db-database",
        args.db_database,
        "--ssl-mode",
        args.ssl_mode,
        "--state-dir",
        str(state_dir),
        "--summary-json",
        str(object_summary),
        "--retry-count",
        str(args.retry_count),
        "--object",
        object_name,
    ]
    for config in args.load_plan_config:
        cmd.extend(["--load-plan-config", config])
    if args.skip_export:
        cmd.append("--skip-export")
    if args.skip_load:
        cmd.append("--skip-load")
    if args.skip_post_reconcile:
        cmd.append("--skip-post-reconcile")
    if args.truncate_first:
        cmd.append("--truncate-first")
    return cmd, object_summary


def run_one_object(
    script_dir: Path,
    args: argparse.Namespace,
    env: dict[str, str],
    object_name: str,
    state_dir: Path,
) -> dict[str, Any]:
    cmd, object_summary = build_child_cmd(script_dir, args, object_name, state_dir)
    result = run_cmd(cmd, env=env, check=False)
    payload: dict[str, Any] = {
        "object": object_name,
        "exit_code": result.returncode,
        "stdout": result.stdout[-4000:],
        "stderr": result.stderr[-4000:],
    }
    if object_summary.exists():
        payload["summary"] = load_json(object_summary)
    return payload


def main() -> int:
    args = parse_args()
    script_dir = Path(__file__).resolve().parent
    env = build_env(args)
    output_root = Path(args.output_root or args.snapshot_root).expanduser().resolve()
    state_dir = Path(args.state_dir or (output_root / "_repair_state")).expanduser().resolve()
    state_dir.mkdir(parents=True, exist_ok=True)

    if not all([args.db_host, args.db_user, args.db_database, args.db_password]):
        raise ValueError("db host/user/password/database are required")
    if args.max_workers < 1:
        raise ValueError("max-workers must be >= 1")

    summary, mismatches = discover_mismatches(script_dir, args, env)
    summary["mismatch_objects"] = mismatches
    summary["max_workers"] = args.max_workers
    summary["objects"] = []
    if not mismatches:
        summary["status"] = "ok"
        summary["finished_at"] = utc_now().isoformat()
        if args.summary_json:
            write_json(Path(args.summary_json).expanduser().resolve(), summary)
        print(json.dumps(summary, indent=2))
        return 0

    with ThreadPoolExecutor(max_workers=args.max_workers) as executor:
        futures = {
            executor.submit(run_one_object, script_dir, args, env, object_name, state_dir): object_name
            for object_name in mismatches
        }
        for future in as_completed(futures):
            summary["objects"].append(future.result())
            if args.summary_json:
                write_json(Path(args.summary_json).expanduser().resolve(), summary)

    failed = [item["object"] for item in summary["objects"] if item.get("exit_code") != 0]
    summary["failed_objects"] = failed
    summary["finished_at"] = utc_now().isoformat()
    summary["status"] = "ok" if not failed else "mismatch"
    if args.summary_json:
        write_json(Path(args.summary_json).expanduser().resolve(), summary)
    print(json.dumps(summary, indent=2))
    return 0 if not failed else 1


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