from __future__ import annotations

import fnmatch
import json
import subprocess

from ingestion.common.models import Column, Table

from .config import SourceConfig


def run_bq_query(config: SourceConfig, sql: str) -> list[dict]:
    """Run a standard-SQL query through the ``bq`` CLI and return JSON rows.

    Shells out to ``bq`` (rather than the Python SDK) to match the repo's CLI-based
    ingestion convention (the TiDB extract shells out to ``ticloud``/``aws``)."""
    completed = subprocess.run(
        [
            "bq",
            f"--project_id={config.gcloud_project}",
            "query",
            "--use_legacy_sql=false",
            "--format=json",
            "--max_rows=100000",
            sql,
        ],
        check=True,
        capture_output=True,
        text=True,
    )
    return json.loads(completed.stdout or "[]")


def kept_table_names(config: SourceConfig, selected: list[str] | None = None) -> list[str]:
    """Single source of truth for which base tables this source processes.

    Applies the include/exclude patterns (and ``require_primary_key`` against the
    config-declared keys, since BigQuery has no enforced PKs). Both schema inspection
    and the export derive their table set from this one function, so the exported set
    and the contracted set cannot silently diverge.
    """
    selected = selected if selected is not None else config.selected_tables
    rows = run_bq_query(
        config,
        "SELECT table_name FROM "
        f"`{config.project}`.{config.dataset}.INFORMATION_SCHEMA.TABLES "
        "WHERE table_type = 'BASE TABLE'",
    )
    names = [
        row["table_name"]
        for row in rows
        if include_table(row["table_name"], selected, config.excluded_tables)
    ]
    if config.require_primary_key:
        names = [name for name in names if config.primary_key_for(name)]
    return sorted(names)


def inspect_tables(config: SourceConfig, selected: list[str] | None = None) -> list[Table]:
    names = kept_table_names(config, selected=selected)
    if not names:
        return []
    quoted = ", ".join(f"'{name}'" for name in names)
    rows = run_bq_query(
        config,
        "SELECT table_name, column_name, data_type, is_nullable, ordinal_position FROM "
        f"`{config.project}`.{config.dataset}.INFORMATION_SCHEMA.COLUMNS "
        f"WHERE table_name IN ({quoted}) ORDER BY table_name, ordinal_position",
    )
    columns_by_table: dict[str, list[Column]] = {name: [] for name in names}
    for row in rows:
        columns_by_table[row["table_name"]].append(
            Column(
                name=row["column_name"],
                # BigQuery has no varchar(n)-style split; the base type is the lowered
                # data_type, the contract-facing column_type keeps the source spelling.
                data_type=str(row["data_type"]).split("(")[0].strip().lower(),
                column_type=str(row["data_type"]),
                nullable=str(row["is_nullable"]).upper() == "YES",
                ordinal_position=int(row["ordinal_position"]),
            )
        )
    return [
        Table(
            database=config.dataset,
            name=name,
            rows_estimate=table_row_estimate(config, name),
            columns=columns_by_table[name],
            primary_key=config.primary_key_for(name),
            watermark_candidates=config.watermark_candidates,
        )
        for name in names
    ]


def table_row_estimate(config: SourceConfig, table_name: str) -> int | None:
    rows = run_bq_query(
        config,
        "SELECT row_count FROM "
        f"`{config.project}`.{config.dataset}.__TABLES__ "
        f"WHERE table_id = '{table_name}'",
    )
    if rows and rows[0].get("row_count") is not None:
        return int(rows[0]["row_count"])
    return None


def selected_table_names(config: SourceConfig) -> list[str]:
    """Names of base tables to process. Used by the export so it extracts exactly the
    tables that get contracts — both derive from ``kept_table_names``."""
    return kept_table_names(config)


def include_table(table_name: str, selected: list[str], excluded: list[str] | None = None) -> bool:
    for pattern in excluded or []:
        if fnmatch.fnmatchcase(table_name, pattern):
            return False
    if not selected:
        return True
    return any(fnmatch.fnmatchcase(table_name, pattern) for pattern in selected)