"""Tests for BigQuery cloud I/O functions in cloud_io.py."""
import sys
import types
import pytest
from unittest.mock import MagicMock


def _make_mock_bigquery():
    """Create a minimal mock google.cloud.bigquery module tree."""
    google = types.ModuleType("google")
    google_cloud = types.ModuleType("google.cloud")
    google_cloud_bq = types.ModuleType("google.cloud.bigquery")

    # Mock classes used by write_parquet_to_bigquery
    google_cloud_bq.Client = MagicMock()
    google_cloud_bq.DatasetReference = MagicMock()
    google_cloud_bq.Dataset = MagicMock()
    google_cloud_bq.LoadJobConfig = MagicMock()
    google_cloud_bq.SourceFormat = MagicMock()
    google_cloud_bq.WriteDisposition = MagicMock()

    google.cloud = google_cloud
    return google, google_cloud, google_cloud_bq


@pytest.fixture(autouse=True)
def _mock_bigquery(monkeypatch):
    """Inject mock google.cloud.bigquery so imports succeed."""
    google, google_cloud, google_cloud_bq = _make_mock_bigquery()
    monkeypatch.setitem(sys.modules, "google", google)
    monkeypatch.setitem(sys.modules, "google.cloud", google_cloud)
    monkeypatch.setitem(sys.modules, "google.cloud.bigquery", google_cloud_bq)
    yield google_cloud_bq


class TestDetectWarehouseScheme:
    def test_bigquery_scheme(self):
        from kanoniv.cloud_io import detect_warehouse_scheme

        assert detect_warehouse_scheme("bigquery://my-project/my_dataset?location=US") == "bigquery"

    def test_snowflake_scheme(self):
        from kanoniv.cloud_io import detect_warehouse_scheme

        assert detect_warehouse_scheme("snowflake://user:pass@account/db/schema") == "snowflake"

    def test_databricks_scheme(self):
        from kanoniv.cloud_io import detect_warehouse_scheme

        assert detect_warehouse_scheme("databricks://token:x@host?http_path=/y") == "databricks"

    def test_unknown_defaults_to_snowflake(self):
        from kanoniv.cloud_io import detect_warehouse_scheme

        assert detect_warehouse_scheme("postgresql://user:pass@host/db") == "snowflake"


class TestParseBigQueryUrl:
    def test_full_url(self):
        from kanoniv.cloud_io import _parse_bigquery_url

        url = "bigquery://my-gcp-project/analytics_dataset?location=US"
        params = _parse_bigquery_url(url)

        assert params["project"] == "my-gcp-project"
        assert params["dataset"] == "analytics_dataset"
        assert params["location"] == "US"

    def test_minimal_url(self):
        from kanoniv.cloud_io import _parse_bigquery_url

        url = "bigquery://my-project/dataset"
        params = _parse_bigquery_url(url)

        assert params["project"] == "my-project"
        assert params["dataset"] == "dataset"
        assert "location" not in params

    def test_project_only(self):
        from kanoniv.cloud_io import _parse_bigquery_url

        url = "bigquery://my-project"
        params = _parse_bigquery_url(url)

        assert params["project"] == "my-project"
        assert "dataset" not in params

    def test_special_chars_in_project(self):
        from kanoniv.cloud_io import _parse_bigquery_url

        url = "bigquery://my-project-123/ds?location=EU"
        params = _parse_bigquery_url(url)

        assert params["project"] == "my-project-123"
        assert params["location"] == "EU"


class TestReadArrowBigQuery:
    def test_calls_list_rows_to_arrow(self, _mock_bigquery):
        from kanoniv.cloud_io import read_arrow_bigquery

        mock_client = MagicMock()
        mock_table_ref = MagicMock()
        mock_rows = MagicMock()
        mock_arrow_table = MagicMock()

        mock_rows.to_arrow.return_value = mock_arrow_table
        mock_client.get_table.return_value = mock_table_ref
        mock_client.list_rows.return_value = mock_rows
        _mock_bigquery.Client.return_value = mock_client

        url = "bigquery://my-project/my_dataset?location=US"
        result = read_arrow_bigquery("my_table", url)

        assert result is mock_arrow_table
        mock_client.get_table.assert_called_once_with("my-project.my_dataset.my_table")
        mock_client.list_rows.assert_called_once_with(mock_table_ref)
        mock_rows.to_arrow.assert_called_once()
        mock_client.close.assert_called_once()

    def test_fully_qualified_table_name(self, _mock_bigquery):
        from kanoniv.cloud_io import read_arrow_bigquery

        mock_client = MagicMock()
        mock_table_ref = MagicMock()
        mock_rows = MagicMock()
        mock_rows.to_arrow.return_value = MagicMock()
        mock_client.get_table.return_value = mock_table_ref
        mock_client.list_rows.return_value = mock_rows
        _mock_bigquery.Client.return_value = mock_client

        url = "bigquery://my-project/my_dataset"
        read_arrow_bigquery("other-project.other_dataset.other_table", url)

        # Should use the fully-qualified name as-is (has dots)
        mock_client.get_table.assert_called_once_with(
            "other-project.other_dataset.other_table"
        )

    def test_passes_project_and_location(self, _mock_bigquery):
        from kanoniv.cloud_io import read_arrow_bigquery

        mock_client = MagicMock()
        mock_table_ref = MagicMock()
        mock_rows = MagicMock()
        mock_rows.to_arrow.return_value = MagicMock()
        mock_client.get_table.return_value = mock_table_ref
        mock_client.list_rows.return_value = mock_rows
        _mock_bigquery.Client.return_value = mock_client

        url = "bigquery://my-project/ds?location=EU"
        read_arrow_bigquery("t", url)

        _mock_bigquery.Client.assert_called_once_with(
            project="my-project", location="EU"
        )

    def test_closes_client_on_error(self, _mock_bigquery):
        from kanoniv.cloud_io import read_arrow_bigquery

        mock_client = MagicMock()
        mock_client.get_table.side_effect = RuntimeError("not found")
        _mock_bigquery.Client.return_value = mock_client

        with pytest.raises(RuntimeError, match="not found"):
            read_arrow_bigquery("bad_table", "bigquery://proj/ds")

        mock_client.close.assert_called_once()


class TestWriteParquetToBigQuery:
    def test_creates_dataset_and_loads(self, _mock_bigquery, monkeypatch, tmp_path):
        from kanoniv.cloud_io import write_parquet_to_bigquery

        # Mock pyarrow.parquet so write_table is a no-op
        pq_mod = types.ModuleType("pyarrow.parquet")
        pq_mod.write_table = MagicMock()
        monkeypatch.setitem(sys.modules, "pyarrow.parquet", pq_mod)

        mock_client = MagicMock()
        mock_client.project = "my-project"
        mock_load_job = MagicMock()
        mock_client.load_table_from_file.return_value = mock_load_job
        _mock_bigquery.Client.return_value = mock_client

        mock_arrow_table = MagicMock()
        mock_arrow_table.num_rows = 42

        url = "bigquery://my-project/my_dataset?location=US"
        counts = write_parquet_to_bigquery(
            {"results": mock_arrow_table},
            connection_string=url,
            dataset="output_ds",
        )

        assert counts == {"results": 42}
        # Dataset created
        mock_client.create_dataset.assert_called_once()
        # Load job started and waited on
        mock_client.load_table_from_file.assert_called_once()
        mock_load_job.result.assert_called_once()
        # Client closed
        mock_client.close.assert_called_once()

    def test_uses_client_project_when_url_has_no_project(self, _mock_bigquery, monkeypatch):
        from kanoniv.cloud_io import write_parquet_to_bigquery

        pq_mod = types.ModuleType("pyarrow.parquet")
        pq_mod.write_table = MagicMock()
        monkeypatch.setitem(sys.modules, "pyarrow.parquet", pq_mod)

        mock_client = MagicMock()
        mock_client.project = "adc-default-project"
        mock_load_job = MagicMock()
        mock_client.load_table_from_file.return_value = mock_load_job
        _mock_bigquery.Client.return_value = mock_client

        mock_arrow_table = MagicMock()
        mock_arrow_table.num_rows = 10

        # URL with no project - relies on ADC default
        url = "bigquery:///my_dataset"
        counts = write_parquet_to_bigquery(
            {"t": mock_arrow_table},
            connection_string=url,
        )

        assert counts == {"t": 10}
        # Should use client.project as effective project
        _mock_bigquery.DatasetReference.assert_called_once_with(
            "adc-default-project", "my_dataset"
        )

    def test_closes_client_on_error(self, _mock_bigquery, monkeypatch):
        from kanoniv.cloud_io import write_parquet_to_bigquery

        pq_mod = types.ModuleType("pyarrow.parquet")
        pq_mod.write_table = MagicMock()
        monkeypatch.setitem(sys.modules, "pyarrow.parquet", pq_mod)

        mock_client = MagicMock()
        mock_client.project = "proj"
        mock_client.create_dataset.side_effect = RuntimeError("permission denied")
        _mock_bigquery.Client.return_value = mock_client

        mock_arrow_table = MagicMock()
        mock_arrow_table.num_rows = 5

        with pytest.raises(RuntimeError, match="permission denied"):
            write_parquet_to_bigquery(
                {"t": mock_arrow_table},
                connection_string="bigquery://proj/ds",
            )

        mock_client.close.assert_called_once()


class TestFromWarehouseBigQueryUrl:
    """Verify that Source.from_warehouse() with a bigquery:// URL
    exposes the connection string correctly for cloud path detection."""

    def test_connection_string_preserved(self, monkeypatch):
        """WarehouseAdapter stores _connection_string for cloud_io detection."""
        # Mock sqlalchemy so WarehouseAdapter can be created
        sa = types.ModuleType("sqlalchemy")
        sa.create_engine = MagicMock()
        sa.inspect = MagicMock()
        sa.text = MagicMock(side_effect=lambda s: s)
        monkeypatch.setitem(sys.modules, "sqlalchemy", sa)

        import importlib
        import kanoniv.adapters.warehouse as wh_mod
        importlib.reload(wh_mod)

        from kanoniv.source import Source

        bq_url = "bigquery://my-project/my_dataset?location=US"
        src = Source.from_warehouse(
            "analytics", "events_table",
            connection_string=bq_url,
            primary_key="event_id",
        )

        assert src.connection_string == bq_url

        from kanoniv.cloud_io import detect_warehouse_scheme
        assert detect_warehouse_scheme(src.connection_string) == "bigquery"
