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

import argparse
import csv
import gzip
import hashlib
import json
import mimetypes
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_METADATA_MANIFEST_NAME = "manifest_document.json"
DEFAULT_METADATA_FILE = "Document.csv.gz"
DEFAULT_BLOB_PREFIX = "salesforce/document/blobs"


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 download_document_body(self, document_id: str) -> bytes:
        req = urllib.request.Request(self._api(f"/sobjects/Document/{document_id}/Body"), headers=self._headers())
        with urllib.request.urlopen(req, timeout=300) as resp:
            return resp.read()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Export Document metadata and blob bodies to S3.")
    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("--metadata-file", default=DEFAULT_METADATA_FILE, help="Metadata CSV file inside snapshot-dir.")
    parser.add_argument(
        "--metadata-manifest-name",
        default=DEFAULT_METADATA_MANIFEST_NAME,
        help="Metadata 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("--blob-prefix", default=DEFAULT_BLOB_PREFIX, help="S3 prefix for Document bodies.")
    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.")
    parser.add_argument("--skip-blob-upload", action="store_true", help="Skip uploading Document bodies.")
    return parser.parse_args()


def pick_fields(describe: dict[str, Any]) -> list[str]:
    preferred = [
        "Id",
        "Name",
        "FolderId",
        "IsDeleted",
        "DeveloperName",
        "NamespacePrefix",
        "ContentType",
        "Type",
        "IsPublic",
        "BodyLength",
        "Url",
        "Description",
        "Keywords",
        "IsInternalUseOnly",
        "AuthorId",
        "CreatedDate",
        "CreatedById",
        "LastModifiedDate",
        "LastModifiedById",
        "SystemModstamp",
        "LastViewedDate",
        "LastReferencedDate",
    ]
    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", "base64"}:
            continue
        if str(name) == "Body":
            continue
        available.add(str(name))
    fields = [name for name in preferred if name in available]
    for name in sorted(available):
        if name not in fields:
            fields.append(name)
    if not fields:
        raise ValueError("no usable Document metadata 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 upload_bytes_to_s3(data: bytes, s3_uri: str, content_type: str | None = None) -> 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")
    cmd = [aws_path, "s3", "cp", "-", s3_uri]
    if content_type:
        cmd.extend(["--content-type", content_type])
    subprocess.run(cmd, input=data, 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 guess_content_type(content_type: str, name: str) -> str:
    if content_type:
        return content_type
    guessed, _ = mimetypes.guess_type(name)
    return guessed or "application/octet-stream"


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("Document")
    fields = pick_fields(describe)
    soql = f"SELECT {', '.join(fields)} FROM Document ORDER BY SystemModstamp, Id"
    expected_count = client.count("Document")
    print(f"[ok] expected_count={expected_count}", flush=True)

    metadata_path = snapshot_dir / args.metadata_file
    metadata_rows: list[dict[str, Any]] = []
    page_count = 0
    with gzip.open(metadata_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_rows = 0
            for record in page.get("records", []):
                if not isinstance(record, dict):
                    continue
                record.pop("attributes", None)
                writer.writerow({name: record.get(name, "") for name in fields})
                metadata_rows.append(record)
                page_rows += 1
            print(f"[ok] page={page_count} rows={page_rows} done={len(metadata_rows)}", flush=True)

    metadata_uri = None
    snapshot_s3_dir = resolve_s3_target(snapshot_dir, args, base_manifest)
    if snapshot_s3_dir and not args.skip_s3_upload:
        metadata_uri = snapshot_s3_dir.rstrip("/") + "/" + metadata_path.name
        upload_file_to_s3(metadata_path, metadata_uri)
        print(f"[ok] uploaded={metadata_uri}", flush=True)

    blob_uploaded = 0
    blob_skipped = 0
    blob_failed = 0
    blob_prefix = args.blob_prefix.strip("/")
    if not args.skip_blob_upload:
        if not args.s3_bucket:
            raise ValueError("--s3-bucket is required for Document blob upload")
        for idx, row in enumerate(metadata_rows, start=1):
            document_id = str(row.get("Id") or "").strip()
            if not document_id:
                blob_failed += 1
                continue
            blob_uri = f"s3://{args.s3_bucket}/{blob_prefix}/{document_id}"
            try:
                body = client.download_document_body(document_id)
                upload_bytes_to_s3(
                    body,
                    blob_uri,
                    guess_content_type(str(row.get("ContentType") or ""), str(row.get("Name") or "")),
                )
                blob_uploaded += 1
                if idx % 10 == 0:
                    print(f"[ok] blob_progress uploaded={blob_uploaded} skipped={blob_skipped} failed={blob_failed}", flush=True)
            except Exception as exc:
                blob_failed += 1
                print(f"[warn] blob_failed document_id={document_id} error={exc}", flush=True)
    else:
        blob_skipped = len(metadata_rows)

    manifest = {
        "manifest_version": 1,
        "object": "Document",
        "export_type": "document_layered",
        "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": len(metadata_rows),
        "metadata_file": metadata_path.name,
        "metadata_sha256": sha256_file(metadata_path),
        "metadata_s3_uri": metadata_uri,
        "blob_prefix": blob_prefix,
        "blob_uploaded": blob_uploaded,
        "blob_skipped": blob_skipped,
        "blob_failed": blob_failed,
        "fields": fields,
        "base_manifest": base_manifest_path.name if base_manifest_path.exists() else None,
    }

    manifest_path = snapshot_dir / args.metadata_manifest_name
    manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    if snapshot_s3_dir and not args.skip_s3_upload:
        manifest_uri = snapshot_s3_dir.rstrip("/") + "/" + manifest_path.name
        upload_file_to_s3(manifest_path, manifest_uri)
        print(f"[ok] uploaded={manifest_uri}", flush=True)

    print(f"[ok] metadata_file={metadata_path}", flush=True)
    print(f"[ok] metadata_manifest={manifest_path}", flush=True)
    return 0


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