#!/usr/bin/env python3
"""
Manage telemetry data in DocumentDB.

Provides export (CSV dump) and purge (delete all) operations for the
telemetry collector's startup_events and heartbeat_events collections.

Reads connection details from ~/bastion.env and credentials from
AWS Secrets Manager.

Usage:
    python3 telemetry_db.py export
    python3 telemetry_db.py export --output /tmp/metrics.csv
    python3 telemetry_db.py export --collection startup_events
    python3 telemetry_db.py purge
    python3 telemetry_db.py purge --collection heartbeat_events
    python3 telemetry_db.py purge --confirm
"""

from __future__ import annotations

import argparse
import csv
import json
import logging
import os
import subprocess
import sys
import time
from collections import (
    Counter,
)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s,p%(process)s,{%(filename)s:%(lineno)d},%(levelname)s,%(message)s",
)
logger = logging.getLogger(__name__)

DEFAULT_OUTPUT = "registry_metrics.csv"
CA_BUNDLE_PATH = os.path.expanduser("~/global-bundle.pem")
BASTION_ENV_PATH = os.path.expanduser("~/bastion.env")

COLLECTIONS = ["startup_events", "heartbeat_events"]

# Column order for startup events
STARTUP_COLUMNS = [
    "event",
    "registry_id",
    "v",
    "py",
    "os",
    "arch",
    "cloud",
    "compute",
    "mode",
    "registry_mode",
    "storage",
    "auth",
    "federation",
    "embeddings_provider",
    "embeddings_backend_kind",
    "cloud_detection_method",
    "search_queries_total",
    "search_queries_24h",
    "search_queries_1h",
    "ts",
    "stored_at",
    "source_ip_hash",
]

# Column order for heartbeat events.
#
# Schema v4+ (registry v1.24.0+) adds the deployment-shape fields
# (py/os/arch/mode/registry_mode/storage/auth/federation) to heartbeat
# payloads so long-lived instances whose original startup event predates
# the report window can still contribute that metadata. Pre-v4 clients
# leave these blank in the CSV (the export uses extrasaction="ignore"
# but Pydantic emits None for unset optional fields, so the column
# header is always present).
HEARTBEAT_COLUMNS = [
    "event",
    "registry_id",
    "v",
    "py",
    "os",
    "arch",
    "cloud",
    "compute",
    "mode",
    "registry_mode",
    "storage",
    "auth",
    "federation",
    "servers_count",
    "agents_count",
    "skills_count",
    "peers_count",
    "search_backend",
    "embeddings_provider",
    "embeddings_backend_kind",
    "cloud_detection_method",
    "uptime_hours",
    "search_queries_total",
    "search_queries_24h",
    "search_queries_1h",
    "ts",
    "stored_at",
    "source_ip_hash",
]

# Union of all columns for the combined CSV
ALL_COLUMNS = [
    "event",
    "registry_id",
    "v",
    "py",
    "os",
    "arch",
    "cloud",
    "compute",
    "mode",
    "registry_mode",
    "storage",
    "auth",
    "federation",
    "servers_count",
    "agents_count",
    "skills_count",
    "peers_count",
    "search_backend",
    "embeddings_provider",
    "embeddings_backend_kind",
    "cloud_detection_method",
    "uptime_hours",
    "search_queries_total",
    "search_queries_24h",
    "search_queries_1h",
    "ts",
    "stored_at",
    "source_ip_hash",
]


# ---------------------------------------------------------------------------
# Private helpers — connection, credentials, mongosh wrappers
# ---------------------------------------------------------------------------


def _load_bastion_env() -> dict[str, str]:
    """Load connection variables from ~/bastion.env.

    Returns:
        Dict with DOCDB_ENDPOINT, SECRET_ARN, AWS_REGION.

    Raises:
        SystemExit: If bastion.env is missing or incomplete.
    """
    if not os.path.exists(BASTION_ENV_PATH):
        logger.error(f"Bastion env file not found: {BASTION_ENV_PATH}")
        logger.error("Run setup-bastion.sh first to configure the bastion host.")
        sys.exit(1)

    env = {}
    with open(BASTION_ENV_PATH) as f:
        for line in f:
            line = line.strip()
            if "=" in line and not line.startswith("#"):
                key, _, value = line.partition("=")
                env[key.strip()] = value.strip().strip('"')

    required_keys = ["DOCDB_ENDPOINT", "SECRET_ARN", "AWS_REGION"]
    for key in required_keys:
        if key not in env:
            logger.error(f"Missing {key} in {BASTION_ENV_PATH}")
            sys.exit(1)

    return env


def _get_credentials(
    secret_arn: str,
    aws_region: str,
) -> dict[str, str]:
    """Fetch DocumentDB credentials from AWS Secrets Manager.

    Args:
        secret_arn: ARN of the secret in Secrets Manager.
        aws_region: AWS region for the Secrets Manager call.

    Returns:
        Dict with username, password, database.

    Raises:
        SystemExit: If credentials cannot be retrieved.
    """
    try:
        result = subprocess.run(  # nosec B603 B607 - hardcoded command
            [
                "aws",
                "secretsmanager",
                "get-secret-value",
                "--secret-id",
                secret_arn,
                "--region",
                aws_region,
                "--query",
                "SecretString",
                "--output",
                "text",
            ],
            capture_output=True,
            text=True,
            check=True,
            timeout=30,
        )
        # Parse secret and extract only needed fields — never log raw output
        parsed = json.loads(result.stdout.strip())
        username = parsed["username"]
        password = parsed["password"]
        database = parsed.get("database", "telemetry")
        # Clear raw secret from memory
        del parsed
        return {
            "username": username,
            "password": password,
            "database": database,
        }
    except subprocess.CalledProcessError:
        logger.error("Failed to get secret from Secrets Manager (check ARN and permissions)")
        sys.exit(1)
    except (json.JSONDecodeError, KeyError):
        logger.error("Failed to parse secret (unexpected format)")
        sys.exit(1)


def _run_mongosh(
    endpoint: str,
    username: str,
    password: str,
    database: str,
    eval_script: str,
    timeout: int = 120,
) -> str | None:
    """Run a mongosh eval command and return stdout.

    Args:
        endpoint: DocumentDB cluster endpoint.
        username: Database username.
        password: Database password.
        database: Database name.
        eval_script: JavaScript to evaluate.
        timeout: Command timeout in seconds.

    Returns:
        Stdout string on success, None on failure.
    """
    conn_string = f"mongodb://{username}@{endpoint}:27017/{database}"

    try:
        result = subprocess.run(  # nosec B603 B607 - hardcoded command
            [
                "mongosh",
                conn_string,
                "--tls",
                "--tlsCAFile",
                CA_BUNDLE_PATH,
                "--retryWrites",
                "false",
                "--authenticationMechanism",
                "SCRAM-SHA-1",
                "--password",
                password,
                "--quiet",
                "--eval",
                eval_script,
            ],
            capture_output=True,
            text=True,
            check=True,
            timeout=timeout,
        )
        return result.stdout.strip()
    except subprocess.CalledProcessError:
        logger.error("mongosh command failed (check connection and credentials)")
        return None
    except subprocess.TimeoutExpired:
        logger.error("mongosh command timed out")
        return None


def _get_collection_count(
    endpoint: str,
    username: str,
    password: str,
    database: str,
    collection: str,
) -> int:
    """Get document count for a collection.

    Args:
        endpoint: DocumentDB cluster endpoint.
        username: Database username.
        password: Database password.
        database: Database name.
        collection: Collection name to count.

    Returns:
        Number of documents in the collection.
    """
    eval_script = f"print(db.{collection}.countDocuments({{}}));"
    output = _run_mongosh(endpoint, username, password, database, eval_script, timeout=30)

    if output is None:
        logger.error(f"Failed to count documents in {collection}")
        return 0

    try:
        return int(output)
    except ValueError:
        logger.error(f"Unexpected count output for {collection}: {output[:80]}")
        return 0


def _fetch_documents(
    endpoint: str,
    username: str,
    password: str,
    database: str,
    collection: str,
) -> list[dict]:
    """Fetch all documents from a DocumentDB collection.

    Args:
        endpoint: DocumentDB cluster endpoint.
        username: Database username.
        password: Database password.
        database: Database name.
        collection: Collection name to query.

    Returns:
        List of document dicts.
    """
    eval_script = (
        f"db.{collection}.find({{}}, {{_id:0}})"
        f".sort({{ts:1}}).forEach(d => print(JSON.stringify(d)));"
    )
    output = _run_mongosh(endpoint, username, password, database, eval_script)

    if output is None:
        logger.error(f"Failed to fetch documents from {collection}")
        return []

    documents = []
    for line in output.split("\n"):
        line = line.strip()
        if not line:
            continue
        try:
            documents.append(json.loads(line))
        except json.JSONDecodeError:
            logger.debug(f"Skipping non-JSON line: {line[:80]}")

    return documents


def _delete_documents(
    endpoint: str,
    username: str,
    password: str,
    database: str,
    collection: str,
) -> int:
    """Delete all documents from a DocumentDB collection.

    Args:
        endpoint: DocumentDB cluster endpoint.
        username: Database username.
        password: Database password.
        database: Database name.
        collection: Collection name to purge.

    Returns:
        Number of documents deleted.
    """
    eval_script = (
        f"var r = db.{collection}.deleteMany({{}});"
        f"print(JSON.stringify({{deletedCount: r.deletedCount}}));"
    )
    output = _run_mongosh(endpoint, username, password, database, eval_script)

    if output is None:
        logger.error(f"Failed to delete documents from {collection}")
        return 0

    try:
        parsed = json.loads(output)
        return parsed.get("deletedCount", 0)
    except json.JSONDecodeError:
        logger.error(f"Failed to parse delete result for {collection}")
        return 0


def _write_csv(
    documents: list[dict],
    columns: list[str],
    output_path: str,
) -> int:
    """Write documents to a CSV file.

    Args:
        documents: List of document dicts.
        columns: Column names for the CSV header.
        output_path: Output file path.

    Returns:
        Number of rows written.
    """
    with open(output_path, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=columns, extrasaction="ignore")
        writer.writeheader()

        for doc in documents:
            # Flatten nested $date objects from BSON extended JSON
            for key in ("stored_at", "ts"):
                val = doc.get(key)
                if isinstance(val, dict) and "$date" in val:
                    doc[key] = val["$date"]

            writer.writerow(doc)

    return len(documents)


def _resolve_collections(
    collection_arg: str,
) -> list[str]:
    """Resolve the --collection argument to a list of collection names.

    Args:
        collection_arg: "all", "startup_events", or "heartbeat_events".

    Returns:
        List of collection name strings.
    """
    if collection_arg == "all":
        return list(COLLECTIONS)
    return [collection_arg]


def _print_summary(documents: list[dict]) -> None:
    """Print a formatted summary of telemetry data.

    Args:
        documents: List of all documents (startup + heartbeat events).
    """
    if not documents:
        return

    # Separate by event type
    startup_events = [d for d in documents if d.get("event") == "startup"]
    heartbeat_events = [d for d in documents if d.get("event") == "heartbeat"]

    # Get unique registry IDs
    startup_ids: set[str] = {d.get("registry_id") for d in startup_events if d.get("registry_id")}
    heartbeat_ids: set[str] = {
        d.get("registry_id") for d in heartbeat_events if d.get("registry_id")
    }
    all_ids = startup_ids | heartbeat_ids

    print("\n" + "=" * 80)
    print("TELEMETRY DATA SUMMARY")
    print("=" * 80)
    print(f"\nTotal Events: {len(documents)}")
    print(f"  - Startup Events:   {len(startup_events):4d}")
    print(f"  - Heartbeat Events: {len(heartbeat_events):4d}")
    print(f"\nUnique Registry Instances: {len(all_ids)}")
    print(f"  - Sent Startup:   {len(startup_ids):4d}")
    print(f"  - Sent Heartbeat: {len(heartbeat_ids):4d}")

    # Aggregate field summaries for startup events
    if startup_events:
        print("\n" + "-" * 80)
        print("STARTUP EVENTS - Field Distribution")
        print("-" * 80)

        # Version distribution
        versions = Counter(d.get("v") for d in startup_events if d.get("v"))
        print(f"\nRegistry Versions ({len(versions)} unique):")
        for version, count in versions.most_common(10):
            print(f"  {version:20s} : {count:4d} ({count / len(startup_events) * 100:5.1f}%)")

        # Python version distribution
        py_versions = Counter(d.get("py") for d in startup_events if d.get("py"))
        print(f"\nPython Versions ({len(py_versions)} unique):")
        for py_ver, count in py_versions.most_common():
            print(f"  Python {py_ver:15s} : {count:4d} ({count / len(startup_events) * 100:5.1f}%)")

        # OS distribution
        os_dist = Counter(d.get("os") for d in startup_events if d.get("os"))
        print(f"\nOperating Systems ({len(os_dist)} unique):")
        for os_name, count in os_dist.most_common():
            print(f"  {os_name:20s} : {count:4d} ({count / len(startup_events) * 100:5.1f}%)")

        # Cloud provider distribution
        cloud_dist = Counter(d.get("cloud") for d in startup_events if d.get("cloud"))
        print(f"\nCloud Providers ({len(cloud_dist)} unique):")
        for cloud, count in cloud_dist.most_common():
            print(f"  {cloud:20s} : {count:4d} ({count / len(startup_events) * 100:5.1f}%)")

        # Compute platform distribution
        compute_dist = Counter(d.get("compute") for d in startup_events if d.get("compute"))
        print(f"\nCompute Platforms ({len(compute_dist)} unique):")
        for compute, count in compute_dist.most_common():
            print(f"  {compute:20s} : {count:4d} ({count / len(startup_events) * 100:5.1f}%)")

        # Storage backend distribution
        storage_dist = Counter(d.get("storage") for d in startup_events if d.get("storage"))
        print(f"\nStorage Backends ({len(storage_dist)} unique):")
        for storage, count in storage_dist.most_common():
            print(f"  {storage:20s} : {count:4d} ({count / len(startup_events) * 100:5.1f}%)")

        # Auth provider distribution
        auth_dist = Counter(d.get("auth") for d in startup_events if d.get("auth"))
        print(f"\nAuth Providers ({len(auth_dist)} unique):")
        for auth, count in auth_dist.most_common():
            print(f"  {auth:20s} : {count:4d} ({count / len(startup_events) * 100:5.1f}%)")

        # Federation enabled
        federation_count = sum(1 for d in startup_events if d.get("federation") is True)
        print(
            f"\nFederation Enabled: {federation_count:4d} ({federation_count / len(startup_events) * 100:5.1f}%)"
        )

        # Deployment mode
        mode_dist = Counter(d.get("mode") for d in startup_events if d.get("mode"))
        print(f"\nDeployment Modes ({len(mode_dist)} unique):")
        for mode, count in mode_dist.most_common():
            print(f"  {mode:20s} : {count:4d} ({count / len(startup_events) * 100:5.1f}%)")

    # Aggregate field summaries for heartbeat events.
    #
    # We report two distinct numbers per registered-object field:
    #   1. Fleet sum (latest heartbeat per instance) -- the honest
    #      "how many objects are currently registered across the fleet"
    #      number. Heartbeats fire ~once per 24h, so summing across all
    #      heartbeats double-counts every instance by its heartbeat
    #      count, which inflates the figure 3-4x for active instances.
    #   2. Per-instance distribution (median / p90 / max) using each
    #      instance's latest heartbeat -- the "what does a typical
    #      registered catalog look like" number.
    if heartbeat_events:
        print("\n" + "-" * 80)
        print("HEARTBEAT EVENTS - Field Distribution")
        print("-" * 80)

        # Build a per-instance latest-heartbeat map. Sort by ts ascending so
        # the dict ends up holding the most recent heartbeat per registry_id.
        sorted_hb = sorted(heartbeat_events, key=lambda d: d.get("ts") or "")
        latest_per_instance: dict[str, dict] = {}
        for d in sorted_hb:
            rid = (d.get("registry_id") or "").strip()
            if rid:
                latest_per_instance[rid] = d
        n_instances = len(latest_per_instance)
        print(f"\nUnique instances that heartbeated: {n_instances}")

        def _percentile(values: list[int], pct: float) -> int:
            if not values:
                return 0
            s = sorted(values)
            idx = int(round((len(s) - 1) * pct))
            return s[idx]

        def _print_object_stats(label: str, field: str) -> None:
            counts = [int(d.get(field, 0) or 0) for d in latest_per_instance.values()]
            if not counts:
                return
            nonzero = [c for c in counts if c > 0]
            fleet_sum = sum(counts)
            median_all = _percentile(counts, 0.5)
            median_nz = _percentile(nonzero, 0.5) if nonzero else 0
            p90 = _percentile(counts, 0.9)
            max_v = max(counts)
            print(f"\nRegistered {label}:")
            print(f"  Fleet sum (latest heartbeat per instance): {fleet_sum}")
            print(f"  Instances with > 0:    {len(nonzero)} of {n_instances}")
            print(f"  Median per instance:   {median_all}  (median among non-zero: {median_nz})")
            print(f"  p90 per instance:      {p90}")
            print(f"  Max per instance:      {max_v}")

        _print_object_stats("MCP Servers", "servers_count")
        _print_object_stats("Agents",      "agents_count")
        _print_object_stats("Skills",      "skills_count")
        _print_object_stats("Federation Peers", "peers_count")

        # Search backend distribution
        search_backend_dist = Counter(
            d.get("search_backend") for d in heartbeat_events if d.get("search_backend")
        )
        print(f"\nSearch Backends ({len(search_backend_dist)} unique):")
        for backend, count in search_backend_dist.most_common():
            print(f"  {backend:20s} : {count:4d} ({count / len(heartbeat_events) * 100:5.1f}%)")

        # Embeddings provider distribution
        embeddings_dist = Counter(
            d.get("embeddings_provider") for d in heartbeat_events if d.get("embeddings_provider")
        )
        print(f"\nEmbeddings Providers ({len(embeddings_dist)} unique):")
        for provider, count in embeddings_dist.most_common():
            print(f"  {provider:20s} : {count:4d} ({count / len(heartbeat_events) * 100:5.1f}%)")

        # Uptime statistics
        uptime_hours = [
            d.get("uptime_hours", 0) for d in heartbeat_events if d.get("uptime_hours") is not None
        ]
        if uptime_hours:
            print("\nUptime (hours):")
            print(f"  Average: {sum(uptime_hours) / len(uptime_hours):.1f}")
            print(f"  Min:     {min(uptime_hours):.1f}")
            print(f"  Max:     {max(uptime_hours):.1f}")

    # Search query statistics (common to both)
    # search_queries_total is a lifetime cumulative counter per instance,
    # so we deduplicate by taking max per registry_id before summing.
    print("\n" + "-" * 80)
    print("SEARCH QUERY STATISTICS")
    print("-" * 80)

    instance_max_total: dict[str, int] = {}
    instance_max_24h: dict[str, int] = {}
    instance_max_1h: dict[str, int] = {}

    for d in documents:
        rid = d.get("registry_id") or f"{d.get('cloud')}/{d.get('compute')}"
        sq_total = d.get("search_queries_total")
        sq_24h = d.get("search_queries_24h")
        sq_1h = d.get("search_queries_1h")

        if sq_total is not None:
            instance_max_total[rid] = max(instance_max_total.get(rid, 0), sq_total)
        if sq_24h is not None:
            instance_max_24h[rid] = max(instance_max_24h.get(rid, 0), sq_24h)
        if sq_1h is not None:
            instance_max_1h[rid] = max(instance_max_1h.get(rid, 0), sq_1h)

    if instance_max_total:
        fleet_total = sum(instance_max_total.values())
        instances_with_search = sum(1 for v in instance_max_total.values() if v > 0)
        print("\nTotal Search Queries (lifetime, deduplicated per instance):")
        print(f"  Fleet Total: {fleet_total:,}")
        print(f"  Instances with search activity: {instances_with_search}")
        print(f"  Max from single instance: {max(instance_max_total.values()):,}")

    if instance_max_24h:
        fleet_24h = sum(instance_max_24h.values())
        print("\nSearch Queries (max 24h window per instance):")
        print(f"  Fleet Total: {fleet_24h:,}")
        print(f"  Max from single instance: {max(instance_max_24h.values()):,}")

    if instance_max_1h:
        fleet_1h = sum(instance_max_1h.values())
        print("\nSearch Queries (max 1h window per instance):")
        print(f"  Fleet Total: {fleet_1h:,}")
        print(f"  Max from single instance: {max(instance_max_1h.values()):,}")

    print("\n" + "=" * 80 + "\n")


def _connect(args: argparse.Namespace) -> tuple[dict[str, str], dict[str, str]]:
    """Load bastion env and fetch credentials.

    Args:
        args: Parsed CLI arguments (uses args.debug).

    Returns:
        Tuple of (env_dict, credentials_dict).
    """
    env = _load_bastion_env()
    logger.info(f"DocumentDB endpoint: {env['DOCDB_ENDPOINT']}")

    creds = _get_credentials(env["SECRET_ARN"], env["AWS_REGION"])
    logger.info("Using configured database for telemetry DocumentDB connection")

    return env, creds


# ---------------------------------------------------------------------------
# Public subcommand handlers
# ---------------------------------------------------------------------------


def cmd_export(args: argparse.Namespace) -> None:
    """Handle the 'export' subcommand — dump telemetry data to CSV.

    Args:
        args: Parsed CLI arguments.
    """
    env, creds = _connect(args)
    target_collections = _resolve_collections(args.collection)

    start_time = time.time()
    all_documents = []

    for collection in target_collections:
        logger.info(f"Fetching {collection}...")
        docs = _fetch_documents(
            endpoint=env["DOCDB_ENDPOINT"],
            username=creds["username"],
            password=creds["password"],
            database=creds["database"],
            collection=collection,
        )
        logger.info(f"  Found {len(docs)} documents")
        all_documents.extend(docs)

    if not all_documents:
        logger.warning("No documents found. CSV not created.")
        return

    # Print summary statistics
    _print_summary(all_documents)

    # Determine columns based on collection
    if args.collection == "startup_events":
        columns = STARTUP_COLUMNS
    elif args.collection == "heartbeat_events":
        columns = HEARTBEAT_COLUMNS
    else:
        columns = ALL_COLUMNS

    rows_written = _write_csv(all_documents, columns, args.output)

    elapsed = time.time() - start_time
    logger.info(f"Exported {rows_written} rows to {args.output} in {elapsed:.1f}s")


def cmd_purge(args: argparse.Namespace) -> None:
    """Handle the 'purge' subcommand — delete telemetry data from DocumentDB.

    Args:
        args: Parsed CLI arguments.
    """
    env, creds = _connect(args)
    target_collections = _resolve_collections(args.collection)

    # Show counts before deletion
    total_count = 0
    for collection in target_collections:
        count = _get_collection_count(
            endpoint=env["DOCDB_ENDPOINT"],
            username=creds["username"],
            password=creds["password"],
            database=creds["database"],
            collection=collection,
        )
        logger.info(f"  {collection}: {count} documents")
        total_count += count

    if total_count == 0:
        logger.info("No documents to delete.")
        return

    # Confirm deletion
    if not args.confirm:
        answer = input(
            f"\nDelete {total_count} documents from {', '.join(target_collections)}? [y/N] "
        )
        if answer.lower() != "y":
            logger.info("Aborted.")
            return

    # Delete documents
    start_time = time.time()
    total_deleted = 0

    for collection in target_collections:
        logger.info(f"Purging {collection}...")
        deleted = _delete_documents(
            endpoint=env["DOCDB_ENDPOINT"],
            username=creds["username"],
            password=creds["password"],
            database=creds["database"],
            collection=collection,
        )
        logger.info(f"  Deleted {deleted} documents from {collection}")
        total_deleted += deleted

    elapsed = time.time() - start_time
    logger.info(f"Purged {total_deleted} total documents in {elapsed:.1f}s")


def main():
    """Parse arguments and dispatch to the appropriate subcommand."""
    parser = argparse.ArgumentParser(
        description="Manage telemetry data in DocumentDB",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python3 telemetry_db.py export
    python3 telemetry_db.py export --output /tmp/metrics.csv
    python3 telemetry_db.py export --collection startup_events
    python3 telemetry_db.py purge
    python3 telemetry_db.py purge --collection heartbeat_events
    python3 telemetry_db.py purge --confirm
""",
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        help="Enable debug logging",
    )

    subparsers = parser.add_subparsers(dest="command", required=True)

    # --- export subcommand ---
    export_parser = subparsers.add_parser(
        "export",
        help="Export telemetry data to CSV",
    )
    export_parser.add_argument(
        "--output",
        default=DEFAULT_OUTPUT,
        help=f"Output CSV file path (default: {DEFAULT_OUTPUT})",
    )
    export_parser.add_argument(
        "--collection",
        choices=["all", "startup_events", "heartbeat_events"],
        default="all",
        help="Which collection to export (default: all)",
    )

    # --- purge subcommand ---
    purge_parser = subparsers.add_parser(
        "purge",
        help="Delete all telemetry data from DocumentDB",
    )
    purge_parser.add_argument(
        "--collection",
        choices=["all", "startup_events", "heartbeat_events"],
        default="all",
        help="Which collection to purge (default: all)",
    )
    purge_parser.add_argument(
        "--confirm",
        action="store_true",
        help="Skip interactive confirmation prompt",
    )

    args = parser.parse_args()

    if args.debug:
        logging.getLogger().setLevel(logging.DEBUG)

    if args.command == "export":
        cmd_export(args)
    elif args.command == "purge":
        cmd_purge(args)


if __name__ == "__main__":
    main()
