# #region Test.SupersetSqlLabExecutor [C:3] [TYPE Module] [SEMANTICS test, superset, sqllab, executor, poll]
# @BRIEF Verify SupersetSqlLabExecutor contracts — _get_client, resolve_database_id, execute_sql,
#         poll_execution_status, execute_and_poll, get_query_results.
# @RELATION BINDS_TO -> [SupersetSqlLabExecutor]
# @TEST_EDGE: env_not_found -> Raises ValueError
# @TEST_EDGE: target_database_id_resolution -> Resolves via target_database_id
# @TEST_EDGE: database_name_resolution -> Resolves by database_name
# @TEST_EDGE: databases_empty -> Raises ValueError
# @TEST_EDGE: database_name_not_found -> Raises ValueError
# @TEST_EDGE: execute_sql_no_database_id -> Falls back to resolve_database_id
# @TEST_EDGE: execute_sql_missing_query_id -> Returns result with None query_id
# @TEST_EDGE: polling_success -> Returns success status
# @TEST_EDGE: polling_failed -> Returns failed status
# @TEST_EDGE: polling_timeout -> Returns timeout status
# @TEST_EDGE: polling_error_retry -> Retries on network error
# @TEST_EDGE: execute_and_poll_sync_success -> Returns directly from sync mode
# @TEST_EDGE: execute_and_poll_no_query_id -> Returns failed
# @TEST_EDGE: get_query_results_success -> Returns results
# @TEST_EDGE: get_query_results_failure -> Returns error
import pytest
from unittest.mock import AsyncMock, MagicMock, patch

from src.core.config_manager import ConfigManager
from src.plugins.translate.superset_executor import SupersetSqlLabExecutor


def _make_config_manager(environments=None):
    """Create a mock ConfigManager with the given environments."""
    config = MagicMock(spec=ConfigManager)
    if environments is not None:
        config.get_environments.return_value = environments
        for env in environments:
            config.get_environment.return_value = env
    else:
        config.get_environments.return_value = []
        config.get_environment.return_value = None
    return config


def _make_env(id="env-1", name="Test Env"):
    env = MagicMock()
    env.id = id
    env.name = name
    return env


class TestInit:
    """Verify initialization."""

    def test_init(self):
        config = _make_config_manager()
        executor = SupersetSqlLabExecutor(config, "env-1")
        assert executor.config_manager is config
        assert executor.env_id == "env-1"
        assert executor._client is None
        assert executor._database_id is None


class TestGetClient:
    """Verify _get_client."""

    @pytest.mark.asyncio
    async def test_lazy_initialization(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")

        with patch('src.plugins.translate.superset_executor.get_superset_client',
                   new=AsyncMock(return_value=MagicMock())) as mock_get:
            client = await executor._get_client()
            assert client is not None
            assert executor._client is not None
            # Second call returns cached client
            client2 = await executor._get_client()
            assert client2 is client
            assert mock_get.call_count == 1

    @pytest.mark.asyncio
    async def test_env_not_found(self):
        config = _make_config_manager([])  # no environments
        executor = SupersetSqlLabExecutor(config, "env-1")

        with pytest.raises(ValueError, match="not found"):
            await executor._get_client()

    @pytest.mark.asyncio
    async def test_env_not_found_empty_list(self):
        config = _make_config_manager()
        config.get_environments.return_value = []
        executor = SupersetSqlLabExecutor(config, "nonexistent")
        with pytest.raises(ValueError, match="not found"):
            await executor._get_client()


class TestResolveDatabaseId:
    """Verify resolve_database_id."""

    @pytest.mark.asyncio
    async def test_resolve_by_target_database_id(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")
        mock_client = AsyncMock()
        mock_client.get_database.return_value = {
            "result": {"id": 42, "backend": "postgresql", "database_name": "Analytics DB"},
        }

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            db_id = await executor.resolve_database_id(target_database_id="42")
            assert db_id == 42
            assert executor._database_backend == "postgresql"
            assert executor._database_name == "Analytics DB"

    @pytest.mark.asyncio
    async def test_resolve_by_target_database_id_fallback(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")
        mock_client = AsyncMock()
        # get_database raises error, fallback to database list
        mock_client.get_database.side_effect = ValueError("API error")
        mock_client.get_databases.return_value = (None, [
            {"id": 1, "database_name": "Main DB", "backend": "postgresql"},
        ])

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            db_id = await executor.resolve_database_id(target_database_id="99")
            assert db_id == 1
            assert executor._database_backend == "postgresql"

    @pytest.mark.asyncio
    async def test_resolve_by_database_name(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")
        mock_client = AsyncMock()
        mock_client.get_databases.return_value = (None, [
            {"id": 1, "database_name": "Main DB", "backend": "postgresql"},
            {"id": 2, "database_name": "Analytics DB", "backend": "clickhouse"},
        ])

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            db_id = await executor.resolve_database_id(database_name="Analytics DB")
            assert db_id == 2
            assert executor._database_backend == "clickhouse"

    @pytest.mark.asyncio
    async def test_resolve_by_database_name_case_insensitive(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")
        mock_client = AsyncMock()
        mock_client.get_databases.return_value = (None, [
            {"id": 1, "database_name": "Main DB", "backend": "postgresql"},
        ])

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            db_id = await executor.resolve_database_id(database_name="main db")
            assert db_id == 1

    @pytest.mark.asyncio
    async def test_resolve_database_name_not_found(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")
        mock_client = AsyncMock()
        mock_client.get_databases.return_value = (None, [
            {"id": 1, "database_name": "Main DB", "backend": "postgresql"},
        ])

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            with pytest.raises(ValueError, match="not found"):
                await executor.resolve_database_id(database_name="Nonexistent DB")

    @pytest.mark.asyncio
    async def test_resolve_default_first_database(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")
        mock_client = AsyncMock()
        mock_client.get_databases.return_value = (None, [
            {"id": 42, "database_name": "Default DB", "backend": "mysql"},
        ])

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            db_id = await executor.resolve_database_id()
            assert db_id == 42
            assert executor._database_backend == "mysql"

    @pytest.mark.asyncio
    async def test_resolve_no_databases_raises(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")
        mock_client = AsyncMock()
        mock_client.get_databases.return_value = (None, [])

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            with pytest.raises(ValueError, match="No databases found"):
                await executor.resolve_database_id()


class TestGetDatabaseBackend:
    """Verify get_database_backend."""

    def test_cached_value(self):
        config = _make_config_manager()
        executor = SupersetSqlLabExecutor(config, "env-1")
        assert executor.get_database_backend() is None
        executor._database_backend = "postgresql"
        assert executor.get_database_backend() == "postgresql"


class TestGetDatabaseName:
    """Verify get_database_name."""

    def test_cached_value(self):
        config = _make_config_manager()
        executor = SupersetSqlLabExecutor(config, "env-1")
        assert executor.get_database_name() is None
        executor._database_name = "Main DB"
        assert executor.get_database_name() == "Main DB"


class TestExecuteSql:
    """Verify execute_sql."""

    @pytest.mark.asyncio
    async def test_execute_with_database_id(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")
        executor._database_id = 42

        mock_client = AsyncMock()
        mock_client.network.request.return_value = {
            "query_id": "q-123",
            "status": "success",
        }

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            result = await executor.execute_sql("SELECT 1", database_id=42)
            assert result["query_id"] == "q-123"
            assert result["status"] == "success"
            assert result["database_id"] == 42

    @pytest.mark.asyncio
    async def test_execute_falls_back_to_resolve(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")

        mock_client = AsyncMock()
        mock_client.network.request.return_value = {
            "query_id": "q-456",
            "status": "success",
        }

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            with patch.object(executor, 'resolve_database_id',
                              new=AsyncMock(return_value=99)):
                result = await executor.execute_sql("SELECT 1")
                assert result["query_id"] == "q-456"
                assert result["database_id"] == 99

    @pytest.mark.asyncio
    async def test_execute_api_error(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")
        executor._database_id = 42

        mock_client = AsyncMock()
        mock_client.network.request.side_effect = ValueError("Connection refused")

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            with pytest.raises(ValueError, match="Superset SQL Lab execute failed"):
                await executor.execute_sql("SELECT 1")

    @pytest.mark.asyncio
    async def test_execute_missing_query_id(self):
        """Response missing query_id — still returns result with None query_id."""
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")
        executor._database_id = 42

        mock_client = AsyncMock()
        mock_client.network.request.return_value = {"status": "pending"}

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            result = await executor.execute_sql("SELECT 1")
            assert result["query_id"] is None
            assert result["status"] == "pending"

    @pytest.mark.asyncio
    async def test_execute_query_id_in_result(self):
        """Query ID nested in result block."""
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")
        executor._database_id = 42

        mock_client = AsyncMock()
        mock_client.network.request.return_value = {
            "result": {"query_id": "q-nested"},
            "status": "success",
        }

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            result = await executor.execute_sql("SELECT 1")
            assert result["query_id"] == "q-nested"

    @pytest.mark.asyncio
    async def test_execute_query_id_in_query_block(self):
        """Query ID nested in query block."""
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")
        executor._database_id = 42

        mock_client = AsyncMock()
        mock_client.network.request.return_value = {
            "query": {"id": "q-from-query"},
            "status": "pending",
        }

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            result = await executor.execute_sql("SELECT 1")
            assert result["query_id"] == "q-from-query"


class TestPollExecutionStatus:
    """Verify poll_execution_status."""

    @pytest.mark.asyncio
    async def test_poll_success(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")

        mock_client = AsyncMock()
        mock_client.network.request.return_value = {
            "result": {"status": "success", "state": "success", "rows": 10},
        }

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            result = await executor.poll_execution_status("q-1", max_polls=1)
            assert result["status"] == "success"
            assert result["rows_affected"] == 10

    @pytest.mark.asyncio
    async def test_poll_failed(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")

        mock_client = AsyncMock()
        mock_client.network.request.return_value = {
            "result": {"status": "failed", "state": "failed", "error_message": "Query error"},
        }

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            result = await executor.poll_execution_status("q-1", max_polls=1)
            assert result["status"] == "failed"
            assert "Query error" in result["error_message"]

    @pytest.mark.asyncio
    async def test_poll_timeout(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")

        mock_client = AsyncMock()
        # Always return pending
        mock_client.network.request.return_value = {
            "result": {"status": "pending", "state": "pending"},
        }

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            result = await executor.poll_execution_status("q-1", max_polls=1, poll_interval_seconds=0.01)
            assert result["status"] == "timeout"

    @pytest.mark.asyncio
    async def test_poll_network_error_retry(self):
        """Network error should be retried."""
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")

        mock_client = AsyncMock()
        # First call fails, second succeeds
        mock_client.network.request = AsyncMock(side_effect=[
            ValueError("Network error"),
            {"result": {"status": "success", "state": "success", "rows": 5}},
        ])

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            result = await executor.poll_execution_status("q-1", max_polls=5, poll_interval_seconds=0.01)
            assert result["status"] == "success"

    @pytest.mark.asyncio
    async def test_poll_unknown_state_continues(self):
        """Unknown state should continue polling."""
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")

        mock_client = AsyncMock()
        mock_client.network.request = AsyncMock(side_effect=[
            {"result": {"status": "weird_state", "state": "weird_state"}},
            {"result": {"status": "success", "state": "success", "rows": 3}},
        ])

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            result = await executor.poll_execution_status("q-1", max_polls=5, poll_interval_seconds=0.01)
            assert result["status"] == "success"

    @pytest.mark.asyncio
    async def test_poll_uses_raw_result_when_no_result_key(self):
        """When response has no 'result' key, use raw dict."""
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")

        mock_client = AsyncMock()
        mock_client.network.request.return_value = {
            "status": "success",
            "state": "success",
            "rows": 7,
        }

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            result = await executor.poll_execution_status("q-1", max_polls=1)
            assert result["status"] == "success"

    @pytest.mark.asyncio
    async def test_poll_not_dict_response(self):
        """When response is not a dict, use empty dict."""
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")

        mock_client = AsyncMock()
        mock_client.network.request.return_value = "not a dict"

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            result = await executor.poll_execution_status("q-1", max_polls=5, poll_interval_seconds=0.01)
            assert result["status"] == "timeout"


class TestExecuteAndPoll:
    """Verify execute_and_poll."""

    @pytest.mark.asyncio
    async def test_sync_mode_success(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")

        with patch.object(executor, 'execute_sql',
                          new=AsyncMock(return_value={
                              "query_id": "q-1",
                              "status": "success",
                              "raw_response": {
                                  "query": {"rows": 5, "endDttm": "2024-01-01"},
                                  "data": [{"col": "val"}],
                              },
                          })):
            result = await executor.execute_and_poll("SELECT 1")
            assert result["status"] == "success"
            assert result["rows_affected"] == 5

    @pytest.mark.asyncio
    async def test_async_mode_poll(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")

        with patch.object(executor, 'execute_sql',
                          new=AsyncMock(return_value={
                              "query_id": "q-1",
                              "status": "pending",
                              "raw_response": {},
                          })):
            with patch.object(executor, 'poll_execution_status',
                              new=AsyncMock(return_value={
                                  "query_id": "q-1",
                                  "status": "success",
                                  "state": "success",
                              })):
                result = await executor.execute_and_poll("SELECT 1")
                assert result["status"] == "success"

    @pytest.mark.asyncio
    async def test_no_query_id_returns_failed(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")

        with patch.object(executor, 'execute_sql',
                          new=AsyncMock(return_value={
                              "query_id": None,
                              "status": "pending",
                              "raw_response": {},
                          })):
            result = await executor.execute_and_poll("SELECT 1")
            assert result["status"] == "failed"
            assert "No query_id" in result["error_message"]


class TestGetQueryResults:
    """Verify get_query_results."""

    @pytest.mark.asyncio
    async def test_success(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")

        mock_client = AsyncMock()
        mock_client.network.request.return_value = {
            "results": [{"col": "val"}],
        }

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            result = await executor.get_query_results("q-1")
            assert result["status"] == "success"
            assert result["results"] == [{"col": "val"}]

    @pytest.mark.asyncio
    async def test_no_results(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")

        mock_client = AsyncMock()
        mock_client.network.request.return_value = {}

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            result = await executor.get_query_results("q-1")
            assert result["status"] == "no_results"

    @pytest.mark.asyncio
    async def test_api_error(self):
        env = _make_env()
        config = _make_config_manager([env])
        executor = SupersetSqlLabExecutor(config, "env-1")

        mock_client = AsyncMock()
        mock_client.network.request.side_effect = ValueError("Not found")

        with patch.object(executor, '_get_client', new=AsyncMock(return_value=mock_client)):
            result = await executor.get_query_results("q-1")
            assert result["status"] == "failed"
            assert "Not found" in result["error_message"]
# #endregion Test.SupersetSqlLabExecutor