from pathlib import Path

from google.cloud import bigquery

from pingcap_datapipelines.bq_load import (
    ensure_table,
    join_predicate,
    record_table_load,
    watermark_expression,
    watermark_order,
)
from pingcap_datapipelines.config import PipelineConfig


def test_watermark_expression_uses_available_columns_case_insensitively() -> None:
    table = {"watermark_columns": ["updated_at", "created_at"]}
    assert watermark_expression(["ID", "Updated_At"], table) == "COALESCE(SAFE_CAST(`Updated_At` AS TIMESTAMP))"


def test_watermark_order_falls_back_without_watermark() -> None:
    assert watermark_order(["id"], {"watermark_columns": ["updated_at"]}) == "1 DESC"


def test_join_predicate_quotes_keys() -> None:
    assert join_predicate("T", "S", ["id", "tenant_id"]) == "T.`id` = S.`id` AND T.`tenant_id` = S.`tenant_id`"


def test_record_table_load_allows_raw_only_final_fields() -> None:
    queries = []

    class FakeQueryJob:
        def result(self):
            return []

    class FakeClient:
        def query(self, sql, *, location):
            queries.append((sql, location))
            return FakeQueryJob()

    config = PipelineConfig(
        path=Path("source.yaml"),
        data={
            "source": {"name": "example_source", "type": "bigquery"},
            "target": {
                "project": "gcp-digital-bi",
                "location": "us-west2",
                "datasets": {"raw": "raw", "control": "control"},
            },
            "gcs": {
                "bucket": "bucket",
                "snapshot_prefix": "snapshots",
                "manifest_prefix": "manifests",
                "runtime_log_prefix": "logs",
            },
        },
    )

    record_table_load(
        FakeClient(),
        config,
        {"run_id": "run-1"},
        {"table_name": "example_table"},
        raw_table="raw.example_table_raw",
        final_table=None,
        loaded_rows=10,
        final_rows=None,
        duplicate_primary_key_rows=0,
        status="loaded",
    )

    sql, location = queries[0]
    assert location == "us-west2"
    assert "NULL,\n  10,\n  NULL," in sql


def test_ensure_table_relaxes_required_fields_and_adds_missing_columns() -> None:
    updates = []

    class FakeTable:
        schema = [
            bigquery.SchemaField("Id", "STRING", mode="REQUIRED"),
            bigquery.SchemaField("Vertical__c", "STRING", mode="REQUIRED"),
        ]

    table = FakeTable()

    class FakeClient:
        def get_table(self, table_ref):
            assert table_ref == "project.dataset.table"
            return table

        def update_table(self, updated_table, fields):
            updates.append((updated_table, fields))

    ensure_table(
        FakeClient(),
        "project.dataset.table",
        [
            bigquery.SchemaField("Id", "STRING", mode="REQUIRED"),
            bigquery.SchemaField("Vertical__c", "STRING", mode="NULLABLE"),
            bigquery.SchemaField("Region__c", "STRING", mode="REQUIRED"),
        ],
    )

    assert updates == [(table, ["schema"])]
    modes = {field.name: field.mode for field in table.schema}
    assert modes == {"Id": "REQUIRED", "Vertical__c": "NULLABLE", "Region__c": "NULLABLE"}