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

import argparse
import csv
import gzip
import json
import re
import urllib.parse
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any


DEFAULT_SOURCE_FIELDS = [
    "Unit_Net_Price__c",
    "Service_Start_Date__c",
]
DEFAULT_ID_BATCH_SIZE = 100
DEFAULT_APPLY_BATCH_SIZE = 25
SALESFORCE_ID_RE = re.compile(r"^[A-Za-z0-9]{15,18}$")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Refresh selected OpportunityLineItem fields on a target org from source Salesforce rows with the same Id."
    )
    parser.add_argument("--source-config", required=True, help="Path to source Salesforce auth config.")
    parser.add_argument("--target-base-url", required=True, help="Target org base URL.")
    parser.add_argument("--target-username", required=True, help="Target org username.")
    parser.add_argument("--target-password", required=True, help="Target org password.")
    parser.add_argument("--target-client-id", default="nexuscrm-ui", help="Target org connected app client id.")
    parser.add_argument(
        "--target-client-secret",
        default="nexuscrm-ui-internal",
        help="Target org connected app client secret.",
    )
    parser.add_argument(
        "--fields",
        default=",".join(DEFAULT_SOURCE_FIELDS),
        help="Comma-separated OpportunityLineItem fields to copy from source to target.",
    )
    selector = parser.add_mutually_exclusive_group(required=True)
    selector.add_argument("--opportunity-id", help="Opportunity Id whose OLI rows should be refreshed.")
    selector.add_argument(
        "--snapshot-file",
        help="Path to an exported OpportunityLineItem csv.gz; all row Ids in the file are refreshed.",
    )
    parser.add_argument(
        "--query-batch-size",
        type=int,
        default=DEFAULT_ID_BATCH_SIZE,
        help="Ids per source/target query batch when --snapshot-file is used.",
    )
    parser.add_argument(
        "--apply-batch-size",
        type=int,
        default=DEFAULT_APPLY_BATCH_SIZE,
        help="Records per composite PATCH batch when --apply is used.",
    )
    parser.add_argument("--apply", action="store_true", help="Apply PATCH writes to the target org.")
    return parser.parse_args()


def load_kv_config(path: Path) -> dict[str, str]:
    payload: dict[str, str] = {}
    for line in path.read_text(encoding="utf-8").splitlines():
        stripped = line.strip()
        if not stripped or stripped.startswith("#") or ":" not in stripped:
            continue
        key, value = stripped.split(":", 1)
        payload[key.strip()] = value.strip()
    return payload


def request_json(url: str, *, method: str = "GET", headers: dict[str, str] | None = None, body: dict[str, Any] | None = None) -> Any:
    encoded_body = None
    req_headers = dict(headers or {})
    if body is not None:
        encoded_body = json.dumps(body).encode("utf-8")
        req_headers.setdefault("Content-Type", "application/json")
    req = urllib.request.Request(url, data=encoded_body, headers=req_headers, method=method)
    try:
        with urllib.request.urlopen(req, timeout=60) as resp:
            raw = resp.read()
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"{method} {url} failed with HTTP {exc.code}: {detail}") from exc
    if not raw:
        return None
    return json.loads(raw.decode("utf-8"))


def password_grant(base_url: str, client_id: str, client_secret: str, username: str, password: str) -> dict[str, Any]:
    payload = urllib.parse.urlencode(
        {
            "grant_type": "password",
            "client_id": client_id,
            "client_secret": client_secret,
            "username": username,
            "password": password,
        }
    ).encode("utf-8")
    req = urllib.request.Request(
        f"{base_url.rstrip('/')}/services/oauth2/token",
        data=payload,
        headers={"Content-Type": "application/x-www-form-urlencoded"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=60) as resp:
        return json.loads(resp.read().decode("utf-8"))


def source_password_grant(config: dict[str, str]) -> dict[str, Any]:
    payload = urllib.parse.urlencode(
        {
            "grant_type": "password",
            "client_id": config["clientId"],
            "client_secret": config["clientSecret"],
            "username": config["username"],
            "password": config["password"],
        }
    ).encode("utf-8")
    req = urllib.request.Request(
        f"https://{config['tokenDomain']}/services/oauth2/token",
        data=payload,
        headers={"Content-Type": "application/x-www-form-urlencoded"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=60) as resp:
        return json.loads(resp.read().decode("utf-8"))


def run_query(base_url: str, access_token: str, soql: str) -> list[dict[str, Any]]:
    url = f"{base_url.rstrip('/')}/services/data/v63.0/query?q={urllib.parse.quote(soql)}"
    payload = request_json(url, headers={"Authorization": f"Bearer {access_token}"})
    if isinstance(payload, list):
        raise RuntimeError(json.dumps(payload))
    return list(payload.get("records") or [])


def normalize_fields(raw: str) -> list[str]:
    return [item.strip() for item in raw.split(",") if item.strip()]


def read_snapshot_ids(path: Path) -> list[str]:
    ids: list[str] = []
    seen: set[str] = set()
    with gzip.open(path, "rt", encoding="utf-8-sig", newline="") as src:
        for row in csv.DictReader(src):
            row_id = str(row.get("Id") or "").strip()
            if not row_id or row_id in seen:
                continue
            seen.add(row_id)
            ids.append(row_id)
    return ids


def partition_snapshot_ids(row_ids: list[str]) -> tuple[list[str], list[str]]:
    valid: list[str] = []
    skipped: list[str] = []
    for row_id in row_ids:
        if SALESFORCE_ID_RE.fullmatch(row_id):
            valid.append(row_id)
        else:
            skipped.append(row_id)
    return valid, skipped


def build_soql_for_opportunity(fields: list[str], opportunity_id: str) -> str:
    select_fields = ", ".join(["Id", "OpportunityId", *fields])
    return (
        f"SELECT {select_fields} FROM OpportunityLineItem "
        f"WHERE OpportunityId = '{opportunity_id}' ORDER BY CreatedDate DESC"
    )


def build_soql_for_ids(fields: list[str], row_ids: list[str]) -> str:
    select_fields = ", ".join(["Id", "OpportunityId", *fields])
    quoted_ids = ", ".join(f"'{row_id}'" for row_id in row_ids)
    return (
        f"SELECT {select_fields} FROM OpportunityLineItem "
        f"WHERE Id IN ({quoted_ids}) ORDER BY CreatedDate DESC"
    )


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


def build_patch_plan(
    *,
    target_rows: list[dict[str, Any]],
    source_by_id: dict[str, dict[str, Any]],
    fields: list[str],
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    patches: list[dict[str, Any]] = []
    results: list[dict[str, Any]] = []
    for target_row in target_rows:
        row_id = str(target_row.get("Id") or "").strip()
        if not row_id:
            continue
        source_row = source_by_id.get(row_id)
        if not source_row:
            results.append({"id": row_id, "status": "missing_in_source"})
            continue
        patch_fields: dict[str, Any] = {}
        diffs: dict[str, Any] = {}
        for field in fields:
            source_value = source_row.get(field)
            target_value = target_row.get(field)
            if source_value != target_value:
                patch_fields[field] = source_value
                diffs[field] = {"source": source_value, "target": target_value}
        if not patch_fields:
            results.append({"id": row_id, "status": "already_synced"})
            continue
        patches.append(
            {
                "id": row_id,
                "fields": diffs,
                "record": {
                    "attributes": {"type": "OpportunityLineItem"},
                    "Id": row_id,
                    **patch_fields,
                },
            }
        )
        results.append({"id": row_id, "status": "planned", "fields": diffs})
    return patches, results


def apply_patch_batches(
    *,
    base_url: str,
    access_token: str,
    patches: list[dict[str, Any]],
    batch_size: int,
) -> tuple[int, list[dict[str, Any]]]:
    updated = 0
    failures: list[dict[str, Any]] = []
    def apply_chunk(chunk: list[dict[str, Any]]) -> None:
        nonlocal updated, failures
        payload = {
            "allOrNone": False,
            "records": [item["record"] for item in chunk],
        }
        try:
            response = request_json(
                f"{base_url.rstrip('/')}/services/data/v63.0/composite/sobjects",
                method="PATCH",
                headers={"Authorization": f"Bearer {access_token}"},
                body=payload,
            )
        except Exception as exc:
            if len(chunk) > 1:
                midpoint = max(1, len(chunk) // 2)
                apply_chunk(chunk[:midpoint])
                apply_chunk(chunk[midpoint:])
                return
            item = chunk[0]
            try:
                request_json(
                    f"{base_url.rstrip('/')}/services/data/v63.0/sobjects/OpportunityLineItem/{item['id']}",
                    method="PATCH",
                    headers={"Authorization": f"Bearer {access_token}"},
                    body={key: value for key, value in item["record"].items() if key not in {"attributes", "Id"}},
                )
                updated += 1
            except Exception as single_exc:
                failures.append(
                    {
                        "id": item["id"],
                        "status": "patch_failed",
                        "fields": item["fields"],
                        "errors": [str(exc), str(single_exc)],
                    }
                )
            return
        for item, result in zip(chunk, response):
            if result.get("success"):
                updated += 1
                continue
            failures.append(
                {
                    "id": item["id"],
                    "status": "patch_failed",
                    "fields": item["fields"],
                    "errors": result.get("errors") or [],
                }
            )

    for chunk in batched(patches, batch_size):
        apply_chunk(chunk)
    return updated, failures


def apply_results(
    *,
    planned_results: list[dict[str, Any]],
    failures: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    failures_by_id = {str(item["id"]): item for item in failures}
    merged: list[dict[str, Any]] = []
    for item in planned_results:
        if item.get("status") != "planned":
            merged.append(item)
            continue
        row_id = str(item["id"])
        failure = failures_by_id.get(row_id)
        if failure is not None:
            merged.append(failure)
            continue
        merged.append({"id": row_id, "status": "patched", "fields": item["fields"]})
    return merged


def main() -> int:
    args = parse_args()
    fields = normalize_fields(args.fields)
    if not fields:
        raise SystemExit("at least one field is required")
    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")

    source_cfg = load_kv_config(Path(args.source_config).expanduser().resolve())
    source_auth = source_password_grant(source_cfg)
    source_base = str(source_auth["instance_url"])
    source_token = str(source_auth["access_token"])

    target_auth = password_grant(
        args.target_base_url,
        args.target_client_id,
        args.target_client_secret,
        args.target_username,
        args.target_password,
    )
    target_base = str(target_auth["instance_url"])
    target_token = str(target_auth["access_token"])
    skipped_snapshot_ids: list[str] = []

    if args.opportunity_id:
        soql = build_soql_for_opportunity(fields, args.opportunity_id)
        source_rows = run_query(source_base, source_token, soql)
        target_rows = run_query(target_base, target_token, soql)
    else:
        snapshot_path = Path(str(args.snapshot_file)).expanduser().resolve()
        row_ids = read_snapshot_ids(snapshot_path)
        if not row_ids:
            raise SystemExit(f"snapshot file has no OpportunityLineItem Ids: {snapshot_path}")
        row_ids, skipped_snapshot_ids = partition_snapshot_ids(row_ids)
        if not row_ids:
            raise SystemExit(f"snapshot file has no Salesforce-style OpportunityLineItem Ids: {snapshot_path}")
        source_rows = []
        target_rows = []
        for chunk in batched(row_ids, args.query_batch_size):
            soql = build_soql_for_ids(fields, chunk)
            source_rows.extend(run_query(source_base, source_token, soql))
            target_rows.extend(run_query(target_base, target_token, soql))

    source_by_id = {str(row["Id"]): row for row in source_rows if row.get("Id")}
    patches, results = build_patch_plan(
        target_rows=target_rows,
        source_by_id=source_by_id,
        fields=fields,
    )
    failures: list[dict[str, Any]] = []
    updates = 0
    if args.apply and patches:
        updates, failures = apply_patch_batches(
            base_url=target_base,
            access_token=target_token,
            patches=patches,
            batch_size=args.apply_batch_size,
        )
        results = apply_results(planned_results=results, failures=failures)
    for skipped_id in skipped_snapshot_ids:
        results.append({"id": skipped_id, "status": "skipped_invalid_id"})

    summary = {
        "selector": {
            "opportunity_id": args.opportunity_id,
            "snapshot_file": str(Path(args.snapshot_file).expanduser().resolve()) if args.snapshot_file else None,
        },
        "field_list": fields,
        "query_batch_size": args.query_batch_size,
        "apply_batch_size": args.apply_batch_size,
        "source_count": len(source_rows),
        "target_count": len(target_rows),
        "planned_count": len(patches),
        "updated_count": updates,
        "failure_count": len(failures),
        "skipped_invalid_id_count": len(skipped_snapshot_ids),
        "apply": args.apply,
        "results": results,
    }
    print(json.dumps(summary, indent=2))
    return 0


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