from unittest.mock import AsyncMock, patch

import pytest

from haiku.rag.config import Config
from haiku.rag.config.models import AppConfig, LanceDBConfig
from haiku.rag.store.engine import ConnectionMode, Store, connect_lancedb


class TestConnectionMode:
    def test_local_when_uri_empty(self):
        config = AppConfig(lancedb=LanceDBConfig(uri=""))
        assert ConnectionMode.from_config(config) == ConnectionMode.LOCAL

    def test_cloud_when_db_uri(self):
        config = AppConfig(
            lancedb=LanceDBConfig(
                uri="db://my-database", api_key="key", region="us-east-1"
            )
        )
        assert ConnectionMode.from_config(config) == ConnectionMode.CLOUD

    def test_object_storage_s3(self):
        config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path"))
        assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE

    def test_object_storage_gs(self):
        config = AppConfig(lancedb=LanceDBConfig(uri="gs://bucket/path"))
        assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE

    def test_object_storage_az(self):
        config = AppConfig(lancedb=LanceDBConfig(uri="az://container/path"))
        assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE

    def test_object_storage_hdfs(self):
        config = AppConfig(lancedb=LanceDBConfig(uri="hdfs://namenode/path"))
        assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE

    def test_unknown_uri_treated_as_object_storage(self):
        config = AppConfig(lancedb=LanceDBConfig(uri="custom://something"))
        assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE


class TestConnectLancedb:
    @pytest.mark.asyncio
    async def test_local_passes_db_path(self, temp_db_path):
        config = AppConfig(lancedb=LanceDBConfig(uri=""))
        with patch(
            "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
        ) as mock_connect:
            await connect_lancedb(config, db_path=temp_db_path)
            mock_connect.assert_called_once_with(temp_db_path.resolve())

    @pytest.mark.asyncio
    async def test_local_resolves_relative_db_path(self, tmp_path, monkeypatch):
        from pathlib import Path

        monkeypatch.chdir(tmp_path)
        relative = Path("db/rag/rag.lancedb")
        config = AppConfig(lancedb=LanceDBConfig(uri=""))
        with patch(
            "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
        ) as mock_connect:
            await connect_lancedb(config, db_path=relative)
            mock_connect.assert_called_once_with((tmp_path / relative).resolve())

    @pytest.mark.asyncio
    async def test_cloud_passes_uri_api_key_region(self):
        config = AppConfig(
            lancedb=LanceDBConfig(
                uri="db://my-database", api_key="test-key", region="us-west-2"
            )
        )
        with patch(
            "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
        ) as mock_connect:
            await connect_lancedb(config)
            mock_connect.assert_called_once_with(
                uri="db://my-database", api_key="test-key", region="us-west-2"
            )

    @pytest.mark.asyncio
    async def test_object_storage_passes_uri_and_storage_options(self):
        config = AppConfig(
            lancedb=LanceDBConfig(
                uri="s3://bucket/path",
                storage_options={
                    "endpoint": "http://minio:9000",
                    "region": "us-east-1",
                },
            )
        )
        with patch(
            "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
        ) as mock_connect:
            await connect_lancedb(config)
            mock_connect.assert_called_once_with(
                uri="s3://bucket/path",
                storage_options={
                    "endpoint": "http://minio:9000",
                    "region": "us-east-1",
                },
            )

    @pytest.mark.asyncio
    async def test_object_storage_without_storage_options(self):
        config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path"))
        with patch(
            "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
        ) as mock_connect:
            await connect_lancedb(config)
            mock_connect.assert_called_once_with(uri="s3://bucket/path")

    @pytest.mark.asyncio
    async def test_local_without_db_path_raises(self):
        config = AppConfig(lancedb=LanceDBConfig(uri=""))
        with pytest.raises(
            ValueError, match="No lancedb.uri configured and no db_path provided"
        ):
            await connect_lancedb(config)


class TestStoreConnectionMode:
    @pytest.mark.asyncio
    async def test_store_connection_mode_local(self, temp_db_path):
        async with Store(temp_db_path, create=True) as store:
            assert store._connection_mode == ConnectionMode.LOCAL

    @pytest.mark.asyncio
    async def test_store_connection_mode_cloud(self, temp_db_path):
        async with Store(temp_db_path, create=True) as store:
            with (
                patch.object(Config.lancedb, "uri", "db://test-database"),
                patch.object(Config.lancedb, "api_key", "test-api-key"),
                patch.object(Config.lancedb, "region", "us-east-1"),
            ):
                assert store._connection_mode == ConnectionMode.CLOUD

    @pytest.mark.asyncio
    async def test_store_connection_mode_object_storage(self, temp_db_path):
        async with Store(temp_db_path, create=True) as store:
            with patch.object(Config.lancedb, "uri", "s3://bucket/path"):
                assert store._connection_mode == ConnectionMode.OBJECT_STORAGE


class TestVacuumByConnectionMode:
    @pytest.mark.asyncio
    async def test_cloud_skips_vacuum(self, temp_db_path):
        async with Store(temp_db_path, create=True) as store:
            with (
                patch.object(Config.lancedb, "uri", "db://test-database"),
                patch.object(Config.lancedb, "api_key", "test-api-key"),
                patch.object(Config.lancedb, "region", "us-east-1"),
            ):
                with patch.object(
                    store.chunks_table, "optimize", new_callable=AsyncMock
                ) as mock_optimize:
                    await store.vacuum()
                    mock_optimize.assert_not_called()

    @pytest.mark.asyncio
    async def test_object_storage_runs_vacuum(self, temp_db_path):
        async with Store(temp_db_path, create=True) as store:
            with patch.object(Config.lancedb, "uri", "s3://bucket/path"):
                with patch.object(
                    store.chunks_table, "optimize", new_callable=AsyncMock
                ) as mock_optimize:
                    await store.vacuum()
                    mock_optimize.assert_called()

    @pytest.mark.asyncio
    async def test_local_runs_vacuum(self, temp_db_path):
        async with Store(temp_db_path, create=True) as store:
            with patch.object(Config.lancedb, "uri", ""):
                with patch.object(
                    store.chunks_table, "optimize", new_callable=AsyncMock
                ) as mock_optimize:
                    await store.vacuum()
                    mock_optimize.assert_called()


class TestVectorIndexByConnectionMode:
    @pytest.mark.asyncio
    async def test_cloud_skips_index_creation(self, temp_db_path):
        async with Store(temp_db_path, create=True) as store:
            with (
                patch.object(Config.lancedb, "uri", "db://test-database"),
                patch.object(Config.lancedb, "api_key", "test-api-key"),
                patch.object(Config.lancedb, "region", "us-east-1"),
            ):
                with patch.object(
                    store.chunks_table, "count_rows", new_callable=AsyncMock
                ) as mock_count:
                    await store._ensure_vector_index()
                    mock_count.assert_not_called()

    @pytest.mark.asyncio
    async def test_object_storage_runs_index_creation(self, temp_db_path):
        async with Store(temp_db_path, create=True) as store:
            with patch.object(Config.lancedb, "uri", "s3://bucket/path"):
                with patch.object(
                    store.chunks_table,
                    "count_rows",
                    new_callable=AsyncMock,
                    return_value=0,
                ) as mock_count:
                    await store._ensure_vector_index()
                    mock_count.assert_called()


class TestStoreSkipsPathValidationForRemote:
    @pytest.mark.asyncio
    async def test_skips_path_check_for_cloud(self, tmp_path):
        nonexistent = tmp_path / "does_not_exist" / "db.lancedb"
        config = AppConfig(
            lancedb=LanceDBConfig(
                uri="db://test-database", api_key="key", region="us-east-1"
            )
        )
        with patch(
            "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
        ):
            with patch.object(Store, "_init_tables", new_callable=AsyncMock):
                async with Store(
                    nonexistent,
                    config=config,
                    create=True,
                    skip_validation=True,
                    skip_migration_check=True,
                ) as store:
                    assert store is not None

    @pytest.mark.asyncio
    async def test_skips_path_check_for_object_storage(self, tmp_path):
        nonexistent = tmp_path / "does_not_exist" / "db.lancedb"
        config = AppConfig(
            lancedb=LanceDBConfig(
                uri="s3://bucket/path",
                storage_options={"endpoint": "http://localhost:9000"},
            )
        )
        with patch(
            "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
        ):
            with patch.object(Store, "_init_tables", new_callable=AsyncMock):
                async with Store(
                    nonexistent,
                    config=config,
                    create=True,
                    skip_validation=True,
                    skip_migration_check=True,
                ) as store:
                    assert store is not None


class TestInitFailureCleanup:
    @pytest.mark.asyncio
    async def test_store_aenter_closes_connection_on_init_failure(
        self, temp_db_path, monkeypatch
    ):
        """If _initialize raises after connect, __aenter__ must close the
        AsyncConnection so it doesn't leak (no __aexit__ runs in that case)."""
        mock_conn = AsyncMock()
        mock_conn.close = lambda: mock_conn.close_calls.append(True)  # type: ignore[attr-defined]
        mock_conn.close_calls = []  # type: ignore[attr-defined]

        async def fake_connect(*args, **kwargs):
            return mock_conn

        async def failing_init_tables(self):
            raise RuntimeError("simulated table init failure")

        monkeypatch.setattr("haiku.rag.store.engine.connect_lancedb", fake_connect)
        monkeypatch.setattr(Store, "_init_tables", failing_init_tables)

        with pytest.raises(RuntimeError, match="simulated table init failure"):
            async with Store(temp_db_path, create=True) as store:
                assert store is not None

        assert mock_conn.close_calls == [True], (
            "AsyncConnection.close() was not called on init failure"
        )

    @pytest.mark.asyncio
    async def test_client_aenter_closes_store_on_init_failure(
        self, temp_db_path, monkeypatch
    ):
        """HaikuRAG.__aenter__ must close the store if _initialize fails."""
        from haiku.rag.client import HaikuRAG

        close_calls: list[bool] = []

        original_close = Store.close

        def tracking_close(self):
            close_calls.append(True)
            original_close(self)

        async def failing_init(self):
            # Set db so close() has something to close
            self.db = AsyncMock()
            self.db.close = lambda: None
            raise RuntimeError("simulated initialize failure")

        monkeypatch.setattr(Store, "_initialize", failing_init)
        monkeypatch.setattr(Store, "close", tracking_close)

        with pytest.raises(RuntimeError, match="simulated initialize failure"):
            async with HaikuRAG(temp_db_path, create=True):
                pass

        assert close_calls, "Store.close() was not called when _initialize raised"