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

import argparse
import csv
import gzip
import hashlib
import json
import shutil
import subprocess
import urllib.parse
import urllib.request
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterator


DEFAULT_LOGIN_DOMAIN = "login.salesforce.com"
DEFAULT_API_VERSION = "v63.0"
DEFAULT_MANIFEST_NAME = "manifest_task_who_relation.json"
DEFAULT_OUTPUT_FILE = "TaskWhoRelation.csv.gz"


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


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as fh:
        for chunk in iter(lambda: fh.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def ensure_dir(path: Path) -> None:
    path.mkdir(parents=True, exist_ok=True)


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


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


class SalesforceClient:
    def __init__(self, config: dict[str, str], api_version: str) -> None:
        self.login_domain = config.get("tokenDomain") or DEFAULT_LOGIN_DOMAIN
        self.client_id = config["clientId"]
        self.client_secret = config["clientSecret"]
        self.username = config["username"]
        self.password = config["password"]
        self.api_version = api_version
        self.instance_url = ""
        self.access_token = ""

    def login(self) -> None:
        token_url = f"https://{self.login_domain}/services/oauth2/token"
        form = urllib.parse.urlencode(
            {
                "grant_type": "password",
                "client_id": self.client_id,
                "client_secret": self.client_secret,
                "username": self.username,
                "password": self.password,
            }
        ).encode("utf-8")
        req = urllib.request.Request(token_url, data=form, method="POST")
        with urllib.request.urlopen(req, timeout=30) as resp:
            data = json.loads(resp.read().decode("utf-8"))
        self.instance_url = data["instance_url"]
        self.access_token = data["access_token"]

    def _headers(self) -> dict[str, str]:
        if not self.access_token:
            self.login()
        return {"Authorization": f"Bearer {self.access_token}"}

    def _api(self, path: str) -> str:
        if not self.instance_url:
            self.login()
        return f"{self.instance_url}/services/data/{self.api_version}{path}"

    def _request_json(self, url: str) -> dict[str, Any]:
        req = urllib.request.Request(url, headers=self._headers())
        with urllib.request.urlopen(req, timeout=120) as resp:
            return json.loads(resp.read().decode("utf-8"))

    def describe(self, object_name: str) -> dict[str, Any]:
        return self._request_json(self._api(f"/sobjects/{object_name}/describe"))

    def iter_query_pages(self, soql: str) -> Iterator[dict[str, Any]]:
        url = self._api("/query") + "?q=" + urllib.parse.quote(soql, safe="")
        while url:
            payload = self._request_json(url)
            yield payload
            next_url = payload.get("nextRecordsUrl")
            url = f"{self.instance_url}{next_url}" if next_url else ""

    def count(self, object_name: str) -> int:
        payload = self._request_json(self._api("/query") + "?q=" + urllib.parse.quote(f"SELECT count() FROM {object_name}", safe=""))
        return int(payload.get("totalSize", 0))


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Export TaskWhoRelation using Salesforce REST query pagination.")
    parser.add_argument("--config", required=True, help="Path to Salesforce config file.")
    parser.add_argument("--snapshot-dir", required=True, help="Existing daily snapshot directory.")
    parser.add_argument("--output-file", default=DEFAULT_OUTPUT_FILE, help="Output CSV file inside snapshot-dir.")
    parser.add_argument("--manifest-name", default=DEFAULT_MANIFEST_NAME, help="Manifest file name inside snapshot-dir.")
    parser.add_argument("--base-manifest-name", default="manifest.json", help="Base snapshot manifest name.")
    parser.add_argument("--api-version", default=DEFAULT_API_VERSION, help="Salesforce API version.")
    parser.add_argument("--s3-bucket", help="Optional S3 bucket override.")
    parser.add_argument("--s3-prefix", help="Optional S3 prefix override.")
    parser.add_argument("--skip-s3-upload", action="store_true", help="Skip S3 upload.")
    return parser.parse_args()


def pick_fields(describe: dict[str, Any]) -> list[str]:
    ordered_defaults = [
        "Id",
        "TaskId",
        "RelationId",
        "AccountId",
        "Type",
        "IsDeleted",
        "CreatedDate",
        "CreatedById",
        "LastModifiedDate",
        "LastModifiedById",
        "SystemModstamp",
    ]
    available: set[str] = set()
    for field in describe.get("fields", []):
        name = field.get("name")
        if not name or field.get("deprecatedAndHidden"):
            continue
        if field.get("type") in {"address", "location"}:
            continue
        available.add(str(name))
    fields = [name for name in ordered_defaults if name in available]
    for name in sorted(available):
        if name not in fields:
            fields.append(name)
    if not fields:
        raise ValueError("no usable TaskWhoRelation fields discovered")
    return fields


def upload_file_to_s3(local_path: Path, s3_uri: str) -> None:
    aws_path = shutil.which("aws")
    if not aws_path:
        raise RuntimeError("aws CLI is required for S3 upload but was not found in PATH")
    subprocess.run([aws_path, "s3", "cp", str(local_path), s3_uri], check=True)


def resolve_s3_target(snapshot_dir: Path, args: argparse.Namespace, base_manifest: dict[str, Any]) -> str | None:
    if args.skip_s3_upload:
        return None
    if args.s3_bucket:
        prefix = (args.s3_prefix or "salesforce/full").strip("/")
        parts = snapshot_dir.parts[-2:]
        return f"s3://{args.s3_bucket}/{prefix}/{parts[0]}/{parts[1]}/"
    s3_uri = base_manifest.get("s3_uri")
    return str(s3_uri) if s3_uri else None


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

    readonly_config = load_kv_config(Path(args.config).expanduser().resolve())
    client = SalesforceClient(readonly_config, api_version=args.api_version)
    client.login()

    base_manifest_path = snapshot_dir / args.base_manifest_name
    base_manifest = load_json(base_manifest_path) if base_manifest_path.exists() else {}

    describe = client.describe("TaskWhoRelation")
    fields = pick_fields(describe)
    output_path = snapshot_dir / args.output_file
    soql = f"SELECT {', '.join(fields)} FROM TaskWhoRelation ORDER BY SystemModstamp, Id"
    expected_count = client.count("TaskWhoRelation")
    print(f"[ok] expected_count={expected_count}", flush=True)

    page_count = 0
    record_count = 0
    with gzip.open(output_path, "wt", encoding="utf-8", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=fields)
        writer.writeheader()
        for page in client.iter_query_pages(soql):
            page_count += 1
            page_records = page.get("records", [])
            page_row_count = 0
            for record in page_records:
                if not isinstance(record, dict):
                    continue
                record.pop("attributes", None)
                writer.writerow({name: record.get(name, "") for name in fields})
                record_count += 1
                page_row_count += 1
            print(f"[ok] page={page_count} rows={page_row_count} done={record_count}", flush=True)

    manifest = {
        "manifest_version": 1,
        "object": "TaskWhoRelation",
        "export_type": "task_who_relation_rest",
        "snapshot_time": utc_now().isoformat(),
        "snapshot_dir": str(snapshot_dir),
        "source_org": client.instance_url,
        "api_version": args.api_version,
        "expected_count": expected_count,
        "query": soql,
        "page_count": page_count,
        "record_count": record_count,
        "file": output_path.name,
        "sha256": sha256_file(output_path),
        "fields": fields,
        "base_manifest": base_manifest_path.name if base_manifest_path.exists() else None,
    }

    s3_uri = resolve_s3_target(snapshot_dir, args, base_manifest)
    manifest_path = snapshot_dir / args.manifest_name
    if s3_uri:
        target_file_uri = s3_uri.rstrip("/") + "/" + output_path.name
        manifest_uri = s3_uri.rstrip("/") + "/" + manifest_path.name
        upload_file_to_s3(output_path, target_file_uri)
        manifest["s3_uri"] = target_file_uri
        manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8")
        upload_file_to_s3(manifest_path, manifest_uri)
        print(f"[ok] uploaded={target_file_uri}", flush=True)
        print(f"[ok] uploaded={manifest_uri}", flush=True)
    else:
        manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8")

    print(f"[ok] file={output_path}", flush=True)
    print(f"[ok] manifest={manifest_path}", flush=True)
    return 0


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