"""Unit tests for FabricSQLService."""

import struct
from unittest.mock import Mock, patch

import pytest

from ms_fabric_mcp_server.client.exceptions import FabricConnectionError, FabricError
from ms_fabric_mcp_server.models.item import FabricItem
from ms_fabric_mcp_server.models.results import QueryResult


@pytest.fixture
def sql_module():
    import ms_fabric_mcp_server.services.sql as sql_module

    return sql_module


@pytest.fixture
def sql_service(mock_fabric_client, sql_module, monkeypatch):
    """Create SQL service with patched pyodbc and DefaultAzureCredential."""
    monkeypatch.setattr(sql_module, "PYODBC_AVAILABLE", True)
    monkeypatch.setattr(sql_module, "OTEL_AVAILABLE", False)
    pyodbc_mock = Mock()
    monkeypatch.setattr(sql_module, "pyodbc", pyodbc_mock)
    azure_cred_mock = Mock()
    monkeypatch.setattr(sql_module, "DefaultAzureCredential", azure_cred_mock)

    workspace_service = Mock()
    item_service = Mock()
    service = sql_module.FabricSQLService(
        mock_fabric_client,
        workspace_service,
        item_service,
    )
    return service, workspace_service, item_service, pyodbc_mock, azure_cred_mock


@pytest.mark.unit
class TestFabricSQLService:
    """Test suite for FabricSQLService."""

    def test_init_missing_pyodbc_raises(self, sql_module, mock_fabric_client, monkeypatch):
        """Missing pyodbc raises ImportError."""
        monkeypatch.setattr(sql_module, "PYODBC_AVAILABLE", False)
        with pytest.raises(ImportError):
            sql_module.FabricSQLService(mock_fabric_client, Mock(), Mock())

    def test_get_sql_endpoint_warehouse(self, sql_service, mock_fabric_client):
        """Warehouse SQL endpoint uses properties.connectionString."""
        service, workspace_service, item_service, _, _ = sql_service
        workspace_service.resolve_workspace_id.return_value = "ws-1"
        item_service.get_item_by_name.return_value = FabricItem(
            id="wh-1",
            display_name="Warehouse",
            type="Warehouse",
            workspace_id="ws-1",
        )
        response = Mock()
        response.json.return_value = {"properties": {"connectionString": "wh-endpoint"}}
        mock_fabric_client.make_api_request.return_value = response

        endpoint = service.get_sql_endpoint("Workspace", "Warehouse", "Warehouse")

        assert endpoint == "wh-endpoint"
        mock_fabric_client.make_api_request.assert_called_once_with(
            "GET", "workspaces/ws-1/warehouses/wh-1"
        )

    def test_get_sql_endpoint_lakehouse(self, sql_service, mock_fabric_client):
        """Lakehouse SQL endpoint uses sqlEndpointProperties.connectionString."""
        service, workspace_service, item_service, _, _ = sql_service
        workspace_service.resolve_workspace_id.return_value = "ws-1"
        item_service.get_item_by_name.return_value = FabricItem(
            id="lh-1",
            display_name="Lakehouse",
            type="Lakehouse",
            workspace_id="ws-1",
        )
        response = Mock()
        response.json.return_value = {
            "properties": {"sqlEndpointProperties": {"connectionString": "lh-endpoint"}}
        }
        mock_fabric_client.make_api_request.return_value = response

        endpoint = service.get_sql_endpoint("Workspace", "Lakehouse", "Lakehouse")

        assert endpoint == "lh-endpoint"
        mock_fabric_client.make_api_request.assert_called_once_with(
            "GET", "workspaces/ws-1/lakehouses/lh-1"
        )

    def test_get_sql_endpoint_missing_connection_string(self, sql_service, mock_fabric_client):
        """Missing connection string raises FabricError."""
        service, workspace_service, item_service, _, _ = sql_service
        workspace_service.resolve_workspace_id.return_value = "ws-1"
        item_service.get_item_by_name.return_value = FabricItem(
            id="wh-1",
            display_name="Warehouse",
            type="Warehouse",
            workspace_id="ws-1",
        )
        response = Mock()
        response.json.return_value = {"properties": {}}
        mock_fabric_client.make_api_request.return_value = response

        with pytest.raises(FabricError):
            service.get_sql_endpoint("Workspace", "Warehouse", "Warehouse")

    def test_get_sql_endpoint_invalid_item_type(self, sql_service):
        """Invalid item type raises ValueError."""
        service, *_ = sql_service

        with pytest.raises(ValueError):
            service.get_sql_endpoint("Workspace", "Item", "Invalid")

    def test_get_token_bytes_success(self, sql_service):
        """Token bytes include length prefix."""
        service, _, _, _, azure_cred_mock = sql_service
        token = Mock()
        token.token = "abc"
        azure_cred_instance = Mock()
        azure_cred_instance.get_token.return_value = token
        azure_cred_mock.return_value = azure_cred_instance

        token_bytes = service._get_token_bytes()

        azure_cred_instance.get_token.assert_called_once_with("https://database.windows.net/.default")
        length = struct.unpack("<i", token_bytes[:4])[0]
        assert length == len(token_bytes[4:])

    def test_get_token_bytes_failure(self, sql_service):
        """Token failures raise FabricConnectionError."""
        service, _, _, _, azure_cred_mock = sql_service
        azure_cred_mock.side_effect = RuntimeError("boom")

        with pytest.raises(FabricConnectionError):
            service._get_token_bytes()

    def test_connect_adds_port_and_calls_pyodbc(self, sql_service):
        """Connect adds default port and uses token bytes."""
        service, _, _, pyodbc_mock, _ = sql_service
        service._get_token_bytes = Mock(return_value=b"token")
        pyodbc_mock.connect.return_value = Mock()

        service.connect("server-host")

        args, kwargs = pyodbc_mock.connect.call_args
        assert "Server=server-host,1433" in args[0]
        assert kwargs["attrs_before"][1256] == b"token"
        assert kwargs["autocommit"] is True
        assert service._connection is not None

    def test_connect_failure(self, sql_service):
        """Connection failures raise FabricConnectionError."""
        service, _, _, pyodbc_mock, _ = sql_service
        service._get_token_bytes = Mock(return_value=b"token")
        pyodbc_mock.connect.side_effect = RuntimeError("boom")

        with pytest.raises(FabricConnectionError):
            service.connect("server-host")

    def test_execute_query_requires_connection(self, sql_service):
        """execute_query raises when not connected."""
        service, *_ = sql_service
        service._connection = None

        with pytest.raises(FabricConnectionError):
            service.execute_query("SELECT 1")

    def test_execute_query_success(self, sql_service):
        """execute_query returns QueryResult with rows."""
        service, *_ = sql_service
        cursor = Mock()
        cursor.description = [("id",), ("name",)]
        cursor.fetchall.return_value = [(1, "a"), (2, "b")]
        connection = Mock()
        connection.cursor.return_value = cursor
        service._connection = connection

        result = service.execute_query("SELECT * FROM test")

        assert result.status == "success"
        assert result.row_count == 2
        assert result.columns == ["id", "name"]
        assert result.data == [{"id": 1, "name": "a"}, {"id": 2, "name": "b"}]

    def test_execute_query_failure_returns_error(self, sql_service):
        """execute_query returns error result on failure."""
        service, *_ = sql_service
        cursor = Mock()
        cursor.execute.side_effect = RuntimeError("boom")
        connection = Mock()
        connection.cursor.return_value = cursor
        service._connection = connection

        result = service.execute_query("SELECT * FROM test")

        assert result.status == "error"
        assert "Query execution failed" in result.message

    def test_execute_query_retries_transient_error(self, sql_service):
        """execute_query retries once on transient ODBC errors."""
        service, *_ = sql_service
        cursor = Mock()
        cursor.description = [("id",)]
        cursor.fetchall.return_value = [(1,)]
        cursor.execute.side_effect = [RuntimeError("HYT00: timeout"), None]
        connection = Mock()
        connection.cursor.return_value = cursor
        service._connection = connection
        service._sql_endpoint = None
        service._database = None

        with patch("ms_fabric_mcp_server.services.sql.time.sleep") as sleep_mock, patch(
            "ms_fabric_mcp_server.services.sql.random.uniform", return_value=5.0
        ):
            result = service.execute_query("SELECT 1")

        assert result.status == "success"
        assert cursor.execute.call_count == 2
        sleep_mock.assert_called_once()

    def test_execute_statement_requires_connection(self, sql_service):
        """execute_statement raises when not connected."""
        service, *_ = sql_service
        service._connection = None

        with pytest.raises(FabricConnectionError):
            service.execute_statement("UPDATE test SET a=1")

    def test_execute_statement_success(self, sql_service):
        """execute_statement executes and returns success (autocommit, no manual commit)."""
        service, *_ = sql_service
        cursor = Mock()
        cursor.rowcount = 3
        connection = Mock()
        connection.cursor.return_value = cursor
        service._connection = connection

        result = service.execute_statement("UPDATE test SET a=1")

        assert result["status"] == "success"
        assert result["affected_rows"] == 3
        connection.commit.assert_not_called()

    def test_execute_statement_rejects_non_dml(self, sql_service):
        """execute_statement returns error for non-DML statements."""
        service, *_ = sql_service
        connection = Mock()
        connection.cursor.return_value = Mock()
        service._connection = connection

        result = service.execute_statement("SELECT 1")

        assert result["status"] == "error"
        assert result["affected_rows"] == 0

    def test_execute_statement_rejects_ddl_without_flag(self, sql_service):
        """execute_statement rejects DDL when allow_ddl is False."""
        service, *_ = sql_service
        connection = Mock()
        connection.cursor.return_value = Mock()
        service._connection = connection

        result = service.execute_statement("CREATE TABLE test (id int)")

        assert result["status"] == "error"
        assert result["affected_rows"] == 0

    def test_execute_statement_allows_ddl_with_flag(self, sql_service):
        """execute_statement allows DDL when allow_ddl is True."""
        service, *_ = sql_service
        cursor = Mock()
        cursor.rowcount = 0
        connection = Mock()
        connection.cursor.return_value = cursor
        service._connection = connection

        result = service.execute_statement("CREATE TABLE test (id int)", allow_ddl=True)

        assert result["status"] == "success"
        connection.commit.assert_not_called()

    def test_execute_statement_failure(self, sql_service):
        """execute_statement returns error on failure (no rollback with autocommit)."""
        service, *_ = sql_service
        cursor = Mock()
        cursor.execute.side_effect = RuntimeError("boom")
        connection = Mock()
        connection.cursor.return_value = cursor
        service._connection = connection

        result = service.execute_statement("UPDATE test SET a=1")

        assert result["status"] == "error"
        connection.rollback.assert_not_called()

    def test_execute_sql_query_calls_close_on_error(self, sql_service):
        """execute_sql_query closes connection even on error."""
        service, *_ = sql_service
        service.connect = Mock()
        service.close = Mock()
        service.execute_query = Mock(side_effect=RuntimeError("boom"))

        with pytest.raises(RuntimeError):
            service.execute_sql_query("endpoint", "SELECT 1")

        service.close.assert_called_once()

    def test_execute_sql_query_closes_on_success(self, sql_service):
        """execute_sql_query closes connection on success."""
        service, *_ = sql_service
        service.connect = Mock()
        service.close = Mock()
        service.execute_query = Mock(
            return_value=QueryResult(
                status="success",
                data=[{"id": 1}],
                columns=["id"],
                row_count=1,
            )
        )

        result = service.execute_sql_query("endpoint", "SELECT 1")

        assert result.status == "success"
        service.connect.assert_called_once_with("endpoint", "Metadata")
        service.execute_query.assert_called_once_with("SELECT 1")
        service.close.assert_called_once()

    def test_execute_sql_statement_calls_close_on_error(self, sql_service):
        """execute_sql_statement closes connection even on error."""
        service, *_ = sql_service
        service.connect = Mock()
        service.close = Mock()
        service.execute_statement = Mock(side_effect=RuntimeError("boom"))

        with pytest.raises(RuntimeError):
            service.execute_sql_statement("endpoint", "UPDATE test SET a=1")

        service.close.assert_called_once()

    def test_execute_sql_statement_closes_on_success(self, sql_service):
        """execute_sql_statement closes connection on success."""
        service, *_ = sql_service
        service.connect = Mock()
        service.close = Mock()
        service.execute_statement = Mock(
            return_value={"status": "success", "affected_rows": 1}
        )

        result = service.execute_sql_statement("endpoint", "UPDATE test SET a=1")

        assert result["status"] == "success"
        service.connect.assert_called_once_with("endpoint", "Metadata")
        service.execute_statement.assert_called_once_with(
            "UPDATE test SET a=1", allow_ddl=False
        )
        service.close.assert_called_once()

    def test_get_tables_success(self, sql_service):
        """get_tables returns table list on success."""
        service, *_ = sql_service
        service.execute_query = Mock(
            return_value=QueryResult(
                status="success",
                data=[{"TABLE_NAME": "a"}, {"TABLE_NAME": "b"}],
                columns=["TABLE_NAME"],
                row_count=2,
            )
        )

        tables = service.get_tables("dbo")

        assert tables == ["a", "b"]

    def test_get_tables_error(self, sql_service):
        """get_tables returns empty list on error."""
        service, *_ = sql_service
        service.execute_query = Mock(return_value=QueryResult(status="error"))

        tables = service.get_tables("dbo")

        assert tables == []

    def test_get_table_schema_success(self, sql_service):
        """get_table_schema returns schema details on success."""
        service, *_ = sql_service
        schema_data = [{"COLUMN_NAME": "id", "DATA_TYPE": "int"}]
        service.execute_query = Mock(
            return_value=QueryResult(status="success", data=schema_data)
        )

        result = service.get_table_schema("table", "dbo")

        assert result == schema_data

    def test_get_table_schema_error(self, sql_service):
        """get_table_schema returns empty list on error."""
        service, *_ = sql_service
        service.execute_query = Mock(return_value=QueryResult(status="error"))

        result = service.get_table_schema("table", "dbo")

        assert result == []

    def test_is_connected_false_when_none(self, sql_service):
        """is_connected returns False when no connection."""
        service, *_ = sql_service
        service._connection = None

        assert service.is_connected() is False

    def test_is_connected_true(self, sql_service):
        """is_connected returns True on successful ping."""
        service, *_ = sql_service
        cursor = Mock()
        cursor.fetchone.return_value = (1,)
        connection = Mock()
        connection.cursor.return_value = cursor
        service._connection = connection

        assert service.is_connected() is True

    def test_is_connected_false_on_error(self, sql_service):
        """is_connected returns False on errors."""
        service, *_ = sql_service
        cursor = Mock()
        cursor.execute.side_effect = RuntimeError("boom")
        connection = Mock()
        connection.cursor.return_value = cursor
        service._connection = connection

        assert service.is_connected() is False

    def test_close_resets_connection(self, sql_service):
        """close resets connection even on close error."""
        service, *_ = sql_service
        connection = Mock()
        connection.close.side_effect = RuntimeError("boom")
        service._connection = connection

        service.close()

        assert service._connection is None
