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

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


DEFAULT_LOGIN_DOMAIN = "login.salesforce.com"
DEFAULT_API_VERSION = "v63.0"
DEFAULT_MANIFEST_NAME = "manifest_content_document_link.json"
DEFAULT_OUTPUT_FILE = "ContentDocumentLink.csv.gz"
DEFAULT_CHUNK_SIZE = 200

SAFE_SF_ID_RE = re.compile(r"^[A-Za-z0-9]{15,18}$")


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 query_pages(self, soql: str) -> list[dict[str, Any]]:
        url = self._api("/query") + "?q=" + urllib.parse.quote(soql, safe="")
        records: list[dict[str, Any]] = []
        while url:
            payload = self._request_json(url)
            for record in payload.get("records", []):
                if isinstance(record, dict):
                    record.pop("attributes", None)
                    records.append(record)
            next_url = payload.get("nextRecordsUrl")
            url = f"{self.instance_url}{next_url}" if next_url else ""
        return records


def parse_args() -> argparse.Namespace:
    script_dir = Path(__file__).resolve().parent
    parser = argparse.ArgumentParser(
        description="Export ContentDocumentLink in filtered chunks using ContentDocument IDs from a daily snapshot."
    )
    parser.add_argument("--config", required=True, help="Path to readonly Salesforce config file.")
    parser.add_argument(
        "--snapshot-dir",
        required=True,
        help="Existing daily snapshot directory containing ContentDocument.csv.gz and manifest.json.",
    )
    parser.add_argument(
        "--content-document-file",
        default=DEFAULT_OUTPUT_FILE.replace("ContentDocumentLink", "ContentDocument"),
        help="ContentDocument source file name inside snapshot-dir.",
    )
    parser.add_argument(
        "--linked-entity-file",
        action="append",
        default=[],
        help=(
            "CSV.gz file inside snapshot-dir whose Id values should seed "
            "ContentDocumentLink.LinkedEntityId queries. May be repeated."
        ),
    )
    parser.add_argument(
        "--output-file",
        default=DEFAULT_OUTPUT_FILE,
        help="Output file name written inside snapshot-dir.",
    )
    parser.add_argument(
        "--manifest-name",
        default=DEFAULT_MANIFEST_NAME,
        help="Special-path manifest file name written inside snapshot-dir.",
    )
    parser.add_argument(
        "--base-manifest-name",
        default="manifest.json",
        help="Base snapshot manifest file name inside snapshot-dir.",
    )
    parser.add_argument("--api-version", default=DEFAULT_API_VERSION, help="Salesforce API version.")
    parser.add_argument("--chunk-size", type=int, default=DEFAULT_CHUNK_SIZE, help="SOQL IN-clause chunk size.")
    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 valid_sf_id(value: str) -> bool:
    return bool(SAFE_SF_ID_RE.match(value))


def read_salesforce_ids(path: Path, *, field_name: str = "Id") -> list[str]:
    with gzip.open(path, "rt", encoding="utf-8-sig", newline="") as fh:
        reader = csv.DictReader(fh)
        ids = {
            str(row.get(field_name, "")).strip()
            for row in reader
            if valid_sf_id(str(row.get(field_name, "")).strip())
        }
    return sorted(ids)


def read_content_document_ids(path: Path) -> list[str]:
    return read_salesforce_ids(path, field_name="Id")


def linked_entity_file_names(values: list[str]) -> list[str]:
    result: list[str] = []
    seen: set[str] = set()
    for raw in values:
        value = str(raw or "").strip()
        if not value or value in seen:
            continue
        result.append(value)
        seen.add(value)
    return result


def pick_fields(describe: dict[str, Any]) -> list[str]:
    required = [
        "Id",
        "ContentDocumentId",
        "LinkedEntityId",
        "ShareType",
        "Visibility",
        "SystemModstamp",
        "CreatedDate",
        "LastModifiedDate",
    ]
    available: set[str] = set()
    for field in describe.get("fields", []):
        name = field.get("name")
        if not name or field.get("deprecatedAndHidden"):
            continue
        available.add(str(name))
    fields = [name for name in required if name in available]
    if not fields:
        raise ValueError("no usable ContentDocumentLink fields discovered")
    return fields


def sql_quote(value: str) -> str:
    return "'" + value.replace("\\", "\\\\").replace("'", "\\'") + "'"


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


def row_key(row: dict[str, Any]) -> str:
    row_id = str(row.get("Id") or "").strip()
    if row_id:
        return f"id:{row_id}"
    return "\x1f".join(
        [
            str(row.get("ContentDocumentId") or "").strip(),
            str(row.get("LinkedEntityId") or "").strip(),
            str(row.get("ShareType") or "").strip(),
            str(row.get("Visibility") or "").strip(),
        ]
    )


def dedupe_link_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    deduped: dict[str, dict[str, Any]] = {}
    for row in rows:
        key = row_key(row)
        if key and key not in deduped:
            deduped[key] = row
    return sorted(
        deduped.values(),
        key=lambda row: (
            str(row.get("ContentDocumentId") or ""),
            str(row.get("LinkedEntityId") or ""),
            str(row.get("Id") or ""),
        ),
    )


def query_content_document_links(
    *,
    client: SalesforceClient,
    fields: list[str],
    filter_field: str,
    ids: list[str],
    chunk_size: int,
    log_label: str,
) -> tuple[list[dict[str, Any]], int]:
    rows: list[dict[str, Any]] = []
    total_chunks = 0
    for idx, chunk in enumerate(chunked(ids, max(1, chunk_size)), start=1):
        total_chunks += 1
        quoted_ids = ", ".join(sql_quote(item) for item in chunk)
        soql = (
            f"SELECT {', '.join(fields)} FROM ContentDocumentLink "
            f"WHERE {filter_field} IN ({quoted_ids})"
        )
        chunk_rows = client.query_pages(soql)
        rows.extend(chunk_rows)
        print(f"[ok] {log_label}_chunk={idx} ids={len(chunk)} rows={len(chunk_rows)}")
        time.sleep(0.05)
    return rows, total_chunks


def write_rows_gzip(path: Path, fieldnames: list[str], rows: list[dict[str, Any]]) -> None:
    with gzip.open(path, "wt", encoding="utf-8", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=fieldnames)
        writer.writeheader()
        for row in rows:
            writer.writerow({name: row.get(name, "") for name in fieldnames})


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 {}

    content_document_path = snapshot_dir / args.content_document_file
    if not content_document_path.exists():
        raise FileNotFoundError(f"ContentDocument source file not found: {content_document_path}")

    doc_ids = read_content_document_ids(content_document_path)
    linked_sources: list[dict[str, Any]] = []
    for file_name in linked_entity_file_names(args.linked_entity_file):
        linked_path = snapshot_dir / file_name
        if not linked_path.exists():
            raise FileNotFoundError(f"linked entity source file not found: {linked_path}")
        linked_ids = read_salesforce_ids(linked_path)
        linked_sources.append({"file": file_name, "id_count": len(linked_ids), "ids": linked_ids})

    if not doc_ids and not any(source["ids"] for source in linked_sources):
        raise ValueError(
            f"no ContentDocument IDs found in {content_document_path} "
            "and no linked entity IDs found"
        )

    describe = client.describe("ContentDocumentLink")
    fields = pick_fields(describe)

    all_rows: list[dict[str, Any]] = []
    total_chunks = 0
    rows, chunks = query_content_document_links(
        client=client,
        fields=fields,
        filter_field="ContentDocumentId",
        ids=doc_ids,
        chunk_size=args.chunk_size,
        log_label="content_document",
    )
    all_rows.extend(rows)
    total_chunks += chunks
    for source in linked_sources:
        rows, chunks = query_content_document_links(
            client=client,
            fields=fields,
            filter_field="LinkedEntityId",
            ids=source["ids"],
            chunk_size=args.chunk_size,
            log_label=f"linked_entity_{Path(source['file']).stem}",
        )
        all_rows.extend(rows)
        total_chunks += chunks
        source["chunk_count"] = chunks
        source["record_count_before_dedupe"] = len(rows)

    deduped_rows = dedupe_link_rows(all_rows)

    output_path = snapshot_dir / args.output_file
    write_rows_gzip(output_path, fields, deduped_rows)

    manifest = {
        "manifest_version": 1,
        "export_type": "content_document_link_filtered",
        "snapshot_time": utc_now().isoformat(),
        "snapshot_dir": str(snapshot_dir),
        "source_org": client.instance_url,
        "api_version": args.api_version,
        "source_file": str(content_document_path.name),
        "linked_entity_sources": [
            {key: value for key, value in source.items() if key != "ids"}
            for source in linked_sources
        ],
        "doc_id_count": len(doc_ids),
        "chunk_size": args.chunk_size,
        "chunk_count": total_chunks,
        "record_count_before_dedupe": len(all_rows),
        "record_count": len(deduped_rows),
        "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)
    if s3_uri:
        target_file_uri = s3_uri.rstrip("/") + "/" + output_path.name
        manifest_uri = s3_uri.rstrip("/") + "/" + args.manifest_name
        upload_file_to_s3(output_path, target_file_uri)
        manifest["s3_uri"] = target_file_uri
        manifest_path = snapshot_dir / args.manifest_name
        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}")
        print(f"[ok] uploaded={manifest_uri}")
    else:
        manifest_path = snapshot_dir / args.manifest_name
        manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8")

    print(f"[ok] file={output_path}")
    print(f"[ok] manifest={snapshot_dir / args.manifest_name}")
    return 0


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