from __future__ import annotations

import subprocess
from datetime import UTC, datetime

from .config import SourceConfig
from .source_db import selected_table_names


def start_export(
    config: SourceConfig,
    run_date: str | None = None,
    run_id: str | None = None,
    tables: list[str] | None = None,
) -> dict:
    """Extract each selected BigQuery table to its own GCS prefix as Parquet.

    ``bq extract`` runs a synchronous extract job per table (BigQuery has no whole-dataset
    export), so unlike the async TiDB export there is no separate poll/wait step — when
    this returns, the objects are in GCS and the manifest can be written.
    """
    run_date = run_date or datetime.now(UTC).strftime("%Y-%m-%d")
    run_id = run_id or datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
    names = tables or selected_table_names(config)
    results = []
    for table in names:
        destination = f"{config.table_gcs_uri(table, run_date, run_id)}{table}-*.parquet"
        extract_table(config, table, destination)
        results.append({"table": table, "gcs_uri": config.table_gcs_uri(table, run_date, run_id)})
    return {
        "run_date": run_date,
        "run_id": run_id,
        "gcs_uri": config.gcs_run_uri(run_date, run_id),
        "tables": results,
    }


def extract_table(config: SourceConfig, table: str, destination_uri: str) -> None:
    subprocess.run(
        [
            "bq",
            f"--project_id={config.gcloud_project}",
            "extract",
            "--destination_format=PARQUET",
            f"--compression={config.compression}",
            f"{config.project}:{config.dataset}.{table}",
            destination_uri,
        ],
        check=True,
        capture_output=True,
        text=True,
    )