from pathlib import Path

from pingcap_datapipelines.config import PipelineConfig, load_config
from pingcap_datapipelines.extract import load_table_metadata, object_table_name, summarize_tables


def config(source_type: str) -> PipelineConfig:
    return PipelineConfig(
        path=Path("source.yaml"),
        data={
            "source": {"name": "example_source", "type": source_type, "database": "finance"},
            "target": {
                "project": "gcp-digital-bi",
                "location": "US",
                "datasets": {"raw": "raw", "final": "final", "control": "control"},
            },
            "gcs": {
                "bucket": "bucket",
                "snapshot_prefix": "snapshots",
                "manifest_prefix": "manifests",
                "runtime_log_prefix": "logs",
            },
        },
    )


def test_bigquery_object_table_name_uses_table_segment() -> None:
    name = "snapshots/source=example_source/run_id=x/table=billing2_invoices/a.parquet"

    assert object_table_name(config("bigquery"), name) == "billing2_invoices"


def test_tidb_object_table_name_uses_export_filename() -> None:
    name = "snapshots/source=x/finance.customer_account.000001.parquet"

    assert object_table_name(config("tidbcloud"), name) == "customer_account"


def test_s3_object_table_name_uses_managed_table_segment() -> None:
    name = "snapshots/source=example_source/run_id=x/table=revenue_c/run-1.parquet"

    assert object_table_name(config("s3"), name) == "revenue_c"


def test_databend_object_table_name_uses_managed_table_segment() -> None:
    name = "snapshots/source=example_source/run_id=x/table=revenue_c/run-1.parquet"

    assert object_table_name(config("databend"), name) == "revenue_c"


def test_mysql_query_object_table_name_uses_managed_table_segment() -> None:
    name = "snapshots/source=example_source/run_id=x/table=cloud_accounts/run-1.parquet"

    assert object_table_name(config("mysql_query"), name) == "cloud_accounts"


def test_helix_manifest_metadata_contains_both_query_table_contracts() -> None:
    metadata = load_table_metadata(load_config(Path("configs/sources/helix.yaml")), selected=None)

    assert list(metadata) == ["cloud_accounts", "users"]
    assert metadata["cloud_accounts"]["primary_key"] == ["scope_id"]
    assert metadata["users"]["primary_key"] == ["user_id"]
    assert metadata["users"]["classifications"] == [
        "internal",
        "employee-directory",
        "contains-work-email",
        "excludes-sensitive-personal-fields",
    ]
    assert len(metadata["users"]["columns"]) == 51


def test_summarize_tables_adds_uris_and_metadata() -> None:
    source = config("bigquery")
    objects = [
        {
            "name": "snapshots/source=example_source/run_id=x/table=billing2_invoices/a.parquet",
            "uri": "gs://b/a.parquet",
            "size": 10,
        },
        {
            "name": "snapshots/source=example_source/run_id=x/table=billing2_invoices/b.parquet",
            "uri": "gs://b/b.parquet",
            "size": 20,
        },
    ]
    metadata = {"billing2_invoices": {"table_name": "billing2_invoices", "primary_key": ["id"], "columns": []}}

    assert summarize_tables(source, run_date="2026-06-23", run_id="run", objects=objects, table_metadata=metadata) == [
        {
            "table_name": "billing2_invoices",
            "primary_key": ["id"],
            "columns": [],
            "file_count": 2,
            "total_bytes": 30,
            "uris": ["gs://b/a.parquet", "gs://b/b.parquet"],
            "gcs_uri": "gs://bucket/snapshots/source=example_source/run_date=2026-06-23/run_id=run/table=billing2_invoices/",
        }
    ]