"""BigQuery extract: one table set drives contracts/manifest/export, drift is surfaced.

Mirrors the TiDB ``test_table_set_sot`` for the BigQuery family. BigQuery has no enforced
primary keys, so ``require_primary_key`` filters against the config-declared keys rather
than an information_schema constraint. ``inspect_tables`` and ``selected_table_names`` must
derive the same set from ``kept_table_names``, and ``schema_drift`` must surface both
mismatch directions.

Runnable without a test harness::

    .venv/bin/python tests/test_bigquery_table_set_sot.py
"""

from __future__ import annotations

import pathlib
import sys

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1]))

from ingestion.extract.bigquery.engine import source_db
from ingestion.extract.bigquery.engine.config import SourceConfig
from ingestion.extract.bigquery.engine.manifest import schema_drift, summarize_tables, table_from_key

# Fake INFORMATION_SCHEMA: four base tables; only invoices/payments have a declared PK.
_TABLE_NAMES = ["invoices", "payments", "events", "audit_log"]
_DECLARED_PKS = {"invoices": ["id"], "payments": ["id"]}


def _fake_run_bq_query(config: SourceConfig, sql: str):
    s = " ".join(sql.split())
    if "INFORMATION_SCHEMA.TABLES" in s:
        return [{"table_name": n} for n in _TABLE_NAMES]
    if "INFORMATION_SCHEMA.COLUMNS" in s:
        wanted = [n for n in _TABLE_NAMES if f"'{n}'" in s]
        rows = []
        for n in wanted:
            rows.append(
                {"table_name": n, "column_name": "id", "data_type": "STRING",
                 "is_nullable": "YES", "ordinal_position": 1}
            )
            rows.append(
                {"table_name": n, "column_name": "amount", "data_type": "BIGNUMERIC",
                 "is_nullable": "YES", "ordinal_position": 2}
            )
        return rows
    if "__TABLES__" in s:
        return [{"row_count": 42}]
    return []


def _config(**tables) -> SourceConfig:
    return SourceConfig(
        path=pathlib.Path("source.yaml"),
        data={
            "source": {
                "name": "metadb",
                "project": "proj",
                "dataset": "ds",
                "domain": "billing",
                "owner": "data-platform",
                "primary_keys": _DECLARED_PKS,
                "watermark_candidates": ["updated_at"],
            },
            "tables": tables,
            "gcs": {
                "bucket": "b",
                "location": "us-west2",
                "snapshot_prefix": "snap",
                "manifest_prefix": "m",
                "runtime_log_prefix": "l",
            },
        },
    )


def _with_fake_query(fn):
    original = source_db.run_bq_query
    source_db.run_bq_query = _fake_run_bq_query
    try:
        return fn()
    finally:
        source_db.run_bq_query = original


def test_require_primary_key_filters_against_declared_keys() -> None:
    config = _config(require_primary_key=True)
    names = _with_fake_query(lambda: source_db.kept_table_names(config))
    # events/audit_log have no declared PK and require_primary_key is on, so both drop out.
    assert names == ["invoices", "payments"]


def test_exclude_patterns_apply() -> None:
    config = _config(exclude=["audit_*"], require_primary_key=False)
    names = _with_fake_query(lambda: source_db.kept_table_names(config))
    assert names == ["events", "invoices", "payments"]


def test_inspect_and_selected_share_the_same_set() -> None:
    config = _config(require_primary_key=True)
    inspected = _with_fake_query(lambda: sorted(t.name for t in source_db.inspect_tables(config)))
    selected = _with_fake_query(lambda: source_db.selected_table_names(config))
    assert inspected == selected == ["invoices", "payments"]


def test_inspect_builds_columns_and_declared_pk() -> None:
    config = _config(include=["invoices"], require_primary_key=True)
    tables = _with_fake_query(lambda: source_db.inspect_tables(config))
    assert len(tables) == 1
    table = tables[0]
    assert [c.name for c in table.columns] == ["id", "amount"]
    # BigQuery column_type keeps source spelling; data_type is the lowered base type.
    assert [(c.data_type, c.column_type) for c in table.columns] == [
        ("string", "STRING"),
        ("bignumeric", "BIGNUMERIC"),
    ]
    assert table.primary_key == ["id"]
    assert table.rows_estimate == 42


def test_table_from_key_and_summaries() -> None:
    assert table_from_key("p/dataset=ds/run_id=R/table=invoices/invoices-000.parquet") == "invoices"
    assert table_from_key("p/run_id=R/_metadata.json") is None
    config = _config(include=["invoices"])
    objects = [
        {"key": "p/run_id=R/table=invoices/invoices-000.parquet", "size": 10},
        {"key": "p/run_id=R/table=invoices/invoices-001.parquet", "size": 20},
        {"key": "p/run_id=R/table=invoices/_SUCCESS", "size": 0},
    ]
    summaries = summarize_tables(config, "2026-05-27", "R", objects, {"invoices": {"table_name": "invoices"}})
    assert len(summaries) == 1
    assert summaries[0]["file_count"] == 2 and summaries[0]["total_bytes"] == 30
    assert summaries[0]["gcs_uri"].endswith("table=invoices/")


def test_schema_drift_reports_both_directions() -> None:
    drift = schema_drift(exported={"invoices", "orphan"}, expected={"invoices", "ghost"})
    assert drift == {"files_without_metadata": ["orphan"], "metadata_without_files": ["ghost"]}


if __name__ == "__main__":
    cases = [value for name, value in sorted(globals().items()) if name.startswith("test_") and callable(value)]
    for case in cases:
        case()
        print(f"ok  {case.__name__}")
    print(f"\n{len(cases)} passed")