"""Unit tests for FabricLivyService."""

import itertools
import json
from unittest.mock import Mock, patch

import pytest

from ms_fabric_mcp_server.client.exceptions import (
    FabricAPIError,
    FabricLivyError,
    FabricLivySessionError,
    FabricLivyStatementError,
    FabricLivyTimeoutError,
)


def _make_response(status_code=200, json_data=None, text_data: str | None = None):
    response = Mock()
    response.status_code = status_code
    response.ok = 200 <= status_code < 300
    response.json.return_value = json_data or {}
    response.text = text_data or ""
    response.content = (text_data or "").encode("utf-8")
    response.headers = {}
    return response


@pytest.fixture
def livy_service(mock_fabric_client):
    from ms_fabric_mcp_server.services.livy import FabricLivyService

    return FabricLivyService(
        mock_fabric_client,
        session_wait_timeout=2,
        statement_wait_timeout=2,
        poll_interval=0.01,
    )


@pytest.mark.unit
class TestFabricLivyService:
    """Test suite for FabricLivyService."""

    def test_create_session_payload_includes_environment(self, livy_service, mock_fabric_client):
        """Create session payload includes environment details."""
        response = _make_response(200, {"id": 1, "state": "starting"})
        mock_fabric_client.make_api_request.return_value = response

        result = livy_service.create_session(
            workspace_id="ws-1",
            lakehouse_id="lh-1",
            environment_id="env-1",
            kind="pyspark",
            conf={"spark.executor.memory": "4g"},
            with_wait=False,
        )

        assert result["id"] == 1
        call_kwargs = mock_fabric_client.make_api_request.call_args.kwargs
        payload = call_kwargs["payload"]
        assert payload["kind"] == "pyspark"
        assert payload["conf"]["spark.executor.memory"] == "4g"
        env_details = json.loads(payload["conf"]["spark.fabric.environmentDetails"])
        assert env_details["id"] == "env-1"

    def test_create_session_includes_fallback_info(self, livy_service, mock_fabric_client):
        """Create session surfaces fallback warnings from tags."""
        response = _make_response(
            200,
            {
                "id": 1,
                "state": "starting",
                "tags": {
                    "FallbackReasons": ["env"],
                    "FallbackMessages": ["Using base environment"],
                },
            },
        )
        mock_fabric_client.make_api_request.return_value = response

        result = livy_service.create_session(
            workspace_id="ws-1",
            lakehouse_id="lh-1",
            with_wait=False,
        )

        assert result["fallback_reasons"] == ["env"]
        assert result["fallback_messages"] == ["Using base environment"]

    def test_create_session_with_wait_uses_wait_for_session(self, livy_service, mock_fabric_client):
        """with_wait True delegates to wait_for_session."""
        mock_fabric_client.make_api_request.return_value = _make_response(200, {"id": 7})
        livy_service.wait_for_session = Mock(return_value={"id": 7, "state": "idle"})

        result = livy_service.create_session(
            workspace_id="ws-1",
            lakehouse_id="lh-1",
            with_wait=True,
        )

        assert result["state"] == "idle"
        livy_service.wait_for_session.assert_called_once_with(
            workspace_id="ws-1",
            lakehouse_id="lh-1",
            session_id=7,
            timeout_seconds=None,
        )

    def test_create_session_api_error(self, livy_service, mock_fabric_client):
        """API errors are mapped to FabricLivySessionError."""
        mock_fabric_client.make_api_request.side_effect = FabricAPIError(500, "boom")

        with pytest.raises(FabricLivySessionError):
            livy_service.create_session("ws-1", "lh-1", with_wait=False)

    def test_list_sessions_success(self, livy_service, mock_fabric_client):
        """List sessions returns response payload."""
        response = _make_response(200, {"sessions": [{"id": 1}], "total": 1})
        mock_fabric_client.make_api_request.return_value = response

        result = livy_service.list_sessions("ws-1", "lh-1")

        assert result["total"] == 1
        assert result["sessions"][0]["id"] == 1

    def test_list_sessions_api_error(self, livy_service, mock_fabric_client):
        """API errors raise FabricLivyError."""
        mock_fabric_client.make_api_request.side_effect = FabricAPIError(500, "boom")

        with pytest.raises(FabricLivyError):
            livy_service.list_sessions("ws-1", "lh-1")

    def test_get_session_status_success(self, livy_service, mock_fabric_client):
        """Get session status returns session data."""
        response = _make_response(200, {"id": 1, "state": "idle"})
        mock_fabric_client.make_api_request.return_value = response

        result = livy_service.get_session_status("ws-1", "lh-1", "1")

        assert result["state"] == "idle"

    def test_get_session_status_api_error(self, livy_service, mock_fabric_client):
        """API errors raise FabricLivySessionError."""
        mock_fabric_client.make_api_request.side_effect = FabricAPIError(404, "missing")

        with pytest.raises(FabricLivySessionError):
            livy_service.get_session_status("ws-1", "lh-1", "1")

    def test_close_session_success(self, livy_service, mock_fabric_client):
        """Close session returns JSON body when available."""
        response = _make_response(200, {"status": "ok"})
        mock_fabric_client.make_api_request.return_value = response

        result = livy_service.close_session("ws-1", "lh-1", "1")

        assert result["status"] == "ok"

    def test_close_session_api_error(self, livy_service, mock_fabric_client):
        """API errors raise FabricLivySessionError."""
        mock_fabric_client.make_api_request.side_effect = FabricAPIError(500, "boom")

        with pytest.raises(FabricLivySessionError):
            livy_service.close_session("ws-1", "lh-1", "1")

    def test_get_session_log_success(self, livy_service, mock_fabric_client):
        """Get session log returns log payload."""
        response = _make_response(200, text_data="a\nb")
        mock_fabric_client.make_api_request.return_value = response

        result = livy_service.get_session_log("ws-1", "lh-1", "1", start=0, size=2)

        assert result["status"] == "success"
        assert result["log_content"] == "a\nb"

    def test_get_session_log_api_error(self, livy_service, mock_fabric_client):
        """API errors raise FabricLivySessionError."""
        mock_fabric_client.make_api_request.side_effect = FabricAPIError(500, "boom")

        with pytest.raises(FabricLivySessionError):
            livy_service.get_session_log("ws-1", "lh-1", "1")

    def test_run_statement_with_wait_false(self, livy_service, mock_fabric_client):
        """Run statement returns immediate response when with_wait False."""
        response = _make_response(200, {"id": 3, "state": "running"})
        mock_fabric_client.make_api_request.return_value = response

        result = livy_service.run_statement(
            workspace_id="ws-1",
            lakehouse_id="lh-1",
            session_id="1",
            code="print(1)",
            kind="pyspark",
            with_wait=False,
        )

        assert result["id"] == 3
        call_kwargs = mock_fabric_client.make_api_request.call_args.kwargs
        payload = call_kwargs["payload"]
        assert payload["code"] == "print(1)"
        assert payload["kind"] == "pyspark"

    def test_run_statement_with_wait_true(self, livy_service, mock_fabric_client):
        """Run statement waits for completion when with_wait True."""
        mock_fabric_client.make_api_request.return_value = _make_response(200, {"id": 4})
        livy_service.wait_for_statement = Mock(return_value={"id": 4, "state": "available"})

        result = livy_service.run_statement(
            workspace_id="ws-1",
            lakehouse_id="lh-1",
            session_id="1",
            code="print(1)",
            with_wait=True,
        )

        assert result["state"] == "available"
        livy_service.wait_for_statement.assert_called_once_with(
            workspace_id="ws-1",
            lakehouse_id="lh-1",
            session_id="1",
            statement_id=4,
            timeout_seconds=None,
        )

    def test_run_statement_api_error(self, livy_service, mock_fabric_client):
        """API errors raise FabricLivyStatementError."""
        mock_fabric_client.make_api_request.side_effect = FabricAPIError(500, "boom")

        with pytest.raises(FabricLivyStatementError):
            livy_service.run_statement("ws-1", "lh-1", "1", "print(1)", with_wait=False)

    def test_get_statement_status_success(self, livy_service, mock_fabric_client):
        """Get statement status returns payload."""
        response = _make_response(200, {"id": 1, "state": "available"})
        mock_fabric_client.make_api_request.return_value = response

        result = livy_service.get_statement_status("ws-1", "lh-1", "1", "1")

        assert result["state"] == "available"

    def test_get_statement_status_api_error(self, livy_service, mock_fabric_client):
        """API errors raise FabricLivyStatementError."""
        mock_fabric_client.make_api_request.side_effect = FabricAPIError(404, "missing")

        with pytest.raises(FabricLivyStatementError):
            livy_service.get_statement_status("ws-1", "lh-1", "1", "1")

    def test_cancel_statement_success(self, livy_service, mock_fabric_client):
        """Cancel statement returns response."""
        response = _make_response(200, {"msg": "canceled"})
        mock_fabric_client.make_api_request.return_value = response

        result = livy_service.cancel_statement("ws-1", "lh-1", "1", "1")

        assert result["msg"] == "canceled"
        call_kwargs = mock_fabric_client.make_api_request.call_args.kwargs
        assert call_kwargs["payload"] == {}

    def test_cancel_statement_api_error(self, livy_service, mock_fabric_client):
        """API errors raise FabricLivyStatementError."""
        mock_fabric_client.make_api_request.side_effect = FabricAPIError(500, "boom")

        with pytest.raises(FabricLivyStatementError):
            livy_service.cancel_statement("ws-1", "lh-1", "1", "1")

    def test_wait_for_session_idle_returns(self, livy_service):
        """Idle session returns immediately."""
        livy_service.get_session_status = Mock(return_value={"state": "idle"})

        with patch("time.sleep", return_value=None):
            result = livy_service.wait_for_session("ws-1", "lh-1", "1", timeout_seconds=2)

        assert result["state"] == "idle"

    @pytest.mark.parametrize("state", ["error", "dead", "killed"])
    def test_wait_for_session_error_state_with_logs(self, livy_service, state):
        """Error states raise FabricLivySessionError with log details."""
        livy_service.get_session_status = Mock(return_value={"state": state, "log": ["a", "b"]})

        with patch("time.sleep", return_value=None):
            with pytest.raises(FabricLivySessionError):
                livy_service.wait_for_session("ws-1", "lh-1", "1", timeout_seconds=2)

    def test_wait_for_session_error_state_no_logs(self, livy_service):
        """Error session raises FabricLivySessionError when logs empty."""
        livy_service.get_session_status = Mock(return_value={"state": "error", "log": []})

        with patch("time.sleep", return_value=None):
            with pytest.raises(FabricLivySessionError):
                livy_service.wait_for_session("ws-1", "lh-1", "1", timeout_seconds=2)

    def test_wait_for_session_timeout(self, livy_service):
        """Timeout raises FabricLivyTimeoutError."""
        livy_service.get_session_status = Mock()

        with patch("time.time", side_effect=itertools.chain([0], itertools.repeat(10))):
            with pytest.raises(FabricLivyTimeoutError):
                livy_service.wait_for_session("ws-1", "lh-1", "1", timeout_seconds=1)

    def test_wait_for_statement_available_returns(self, livy_service):
        """Available statement returns immediately."""
        livy_service.get_statement_status = Mock(return_value={"state": "available"})

        with patch("time.sleep", return_value=None):
            result = livy_service.wait_for_statement("ws-1", "lh-1", "1", "1", timeout_seconds=2)

        assert result["state"] == "available"

    def test_wait_for_statement_error_state(self, livy_service):
        """Error statement raises FabricLivyStatementError."""
        livy_service.get_statement_status = Mock(
            return_value={"state": "error", "output": {"evalue": "boom"}}
        )

        with patch("time.sleep", return_value=None):
            with pytest.raises(FabricLivyStatementError):
                livy_service.wait_for_statement("ws-1", "lh-1", "1", "1", timeout_seconds=2)

    def test_wait_for_statement_cancelled_state(self, livy_service):
        """Cancelled statement raises FabricLivyStatementError."""
        livy_service.get_statement_status = Mock(
            return_value={"state": "cancelled", "output": {"evalue": "cancelled"}}
        )

        with patch("time.sleep", return_value=None):
            with pytest.raises(FabricLivyStatementError):
                livy_service.wait_for_statement("ws-1", "lh-1", "1", "1", timeout_seconds=2)

    def test_wait_for_statement_timeout(self, livy_service):
        """Timeout raises FabricLivyTimeoutError."""
        livy_service.get_statement_status = Mock()

        with patch("time.time", side_effect=itertools.chain([0], itertools.repeat(10))):
            with pytest.raises(FabricLivyTimeoutError):
                livy_service.wait_for_statement("ws-1", "lh-1", "1", "1", timeout_seconds=1)

    def test_wait_for_session_retries_on_transient_404(self, livy_service):
        """Transient 404 errors are retried during session polling."""
        # Simulate 2x 404 errors, then success
        livy_service.get_session_status = Mock(
            side_effect=[
                FabricLivySessionError("1", "Failed to get session: API error 404: ..."),
                FabricLivySessionError("1", "Failed to get session: API error 404: ..."),
                {"id": 1, "state": "idle"},
            ]
        )

        with patch("time.sleep", return_value=None):
            result = livy_service.wait_for_session("ws-1", "lh-1", "1", timeout_seconds=10)

        assert result["state"] == "idle"
        assert livy_service.get_session_status.call_count == 3

    def test_wait_for_session_fails_after_max_404_retries(self, livy_service):
        """Exceeding max 404 retries raises FabricLivySessionError."""
        # Simulate repeated 404 errors exceeding max retries (5)
        livy_service.get_session_status = Mock(
            side_effect=FabricLivySessionError("1", "Failed to get session: API error 404: ...")
        )

        with patch("time.sleep", return_value=None):
            with pytest.raises(FabricLivySessionError):
                livy_service.wait_for_session("ws-1", "lh-1", "1", timeout_seconds=10)

        # Should have retried max_404_retries + 1 times
        assert livy_service.get_session_status.call_count == 6
