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


def _make_mock_databricks():
    """Create a minimal mock databricks.sql module tree."""
    db_top = types.ModuleType("databricks")
    db_sql = types.ModuleType("databricks.sql")
    db_sql.connect = MagicMock()
    db_top.sql = db_sql
    return db_top, db_sql


@pytest.fixture(autouse=True)
def _mock_databricks(monkeypatch):
    """Inject mock databricks.sql so imports succeed."""
    db_top, db_sql = _make_mock_databricks()
    monkeypatch.setitem(sys.modules, "databricks", db_top)
    monkeypatch.setitem(sys.modules, "databricks.sql", db_sql)
    yield db_sql


class TestDetectWarehouseScheme:
    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_snowflake_scheme(self):
        from kanoniv.cloud_io import detect_warehouse_scheme

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

    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 TestParseDatabricksUrl:
    def test_full_url(self):
        from kanoniv.cloud_io import _parse_databricks_url

        url = "databricks://token:dapi_secret@adb-123.azuredatabricks.net?http_path=/sql/1.0/warehouses/abc&catalog=main&schema=default"
        params = _parse_databricks_url(url)

        assert params["server_hostname"] == "adb-123.azuredatabricks.net"
        assert params["access_token"] == "dapi_secret"
        assert params["http_path"] == "/sql/1.0/warehouses/abc"
        assert params["catalog"] == "main"
        assert params["schema"] == "default"

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

        url = "databricks://token:dapi_test@host?http_path=/sql/1.0/warehouses/x"
        params = _parse_databricks_url(url)

        assert params["server_hostname"] == "host"
        assert params["access_token"] == "dapi_test"
        assert params["http_path"] == "/sql/1.0/warehouses/x"
        assert "catalog" not in params
        assert "schema" not in params

    def test_special_chars_in_token(self):
        from kanoniv.cloud_io import _parse_databricks_url

        url = "databricks://token:dapi%2Fslash%40at@host?http_path=/x"
        params = _parse_databricks_url(url)

        assert params["access_token"] == "dapi/slash@at"


class TestReadArrowDatabricks:
    def test_calls_fetchall_arrow(self, _mock_databricks):
        from kanoniv.cloud_io import read_arrow_databricks

        mock_conn = MagicMock()
        mock_cursor = MagicMock()
        mock_arrow_table = MagicMock()
        mock_cursor.fetchall_arrow.return_value = mock_arrow_table
        mock_conn.cursor.return_value = mock_cursor
        _mock_databricks.connect.return_value = mock_conn

        url = "databricks://token:dapi_test@host?http_path=/sql/1.0/warehouses/x"
        result = read_arrow_databricks("catalog.schema.table", url)

        assert result is mock_arrow_table
        mock_cursor.execute.assert_called_once_with("SELECT * FROM catalog.schema.table")
        mock_cursor.fetchall_arrow.assert_called_once()
        mock_cursor.close.assert_called_once()
        mock_conn.close.assert_called_once()

    def test_passes_parsed_params(self, _mock_databricks):
        from kanoniv.cloud_io import read_arrow_databricks

        mock_conn = MagicMock()
        mock_cursor = MagicMock()
        mock_cursor.fetchall_arrow.return_value = MagicMock()
        mock_conn.cursor.return_value = mock_cursor
        _mock_databricks.connect.return_value = mock_conn

        url = "databricks://token:dapi_tok@adb-123.net?http_path=/sql/1.0/warehouses/abc&catalog=main"
        read_arrow_databricks("t", url)

        call_kwargs = _mock_databricks.connect.call_args[1]
        assert call_kwargs["server_hostname"] == "adb-123.net"
        assert call_kwargs["access_token"] == "dapi_tok"
        assert call_kwargs["http_path"] == "/sql/1.0/warehouses/abc"
        assert call_kwargs["catalog"] == "main"


class TestFromWarehouseDatabricksUrl:
    """Verify that Source.from_warehouse() with a databricks:// 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

        db_url = "databricks://token:dapi_test@host?http_path=/sql/1.0/warehouses/x"
        src = Source.from_warehouse(
            "crm", "catalog.schema.customers",
            connection_string=db_url,
            primary_key="customer_id",
        )

        assert src.connection_string == db_url

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