import asyncio
import os
from typing import Any
from unittest.mock import AsyncMock, Mock, patch

import httpx
import pytest

from veris_ai import veris
from veris_ai.tool_mock import _session_id_context


@pytest.fixture
def tool_mock():
    return veris


# Test mock decorator
@pytest.mark.asyncio
async def test_mock_decorator_simulation_mode(simulation_env):
    @veris.mock
    async def test_func(param1: str, param2: int) -> dict[str, Any]:
        return {"result": "real"}

    mock_response = {"result": {"mocked": True}}

    with patch("httpx.AsyncClient") as mock_client:
        mock_response_obj = Mock()
        mock_response_obj.json.return_value = mock_response
        mock_response_obj.raise_for_status.return_value = None

        mock_client.return_value.__aenter__.return_value.post = AsyncMock(
            return_value=mock_response_obj,
        )

        result = await test_func("test", 42)
        assert result == {"mocked": True}


@pytest.mark.asyncio
async def test_mock_decorator_production_mode(production_env):
    @veris.mock
    async def test_func(param1: str, param2: int) -> dict:
        return {"result": "real"}

    result = await test_func("test", 42)
    assert result == {"result": "real"}


@pytest.mark.asyncio
async def test_mock_with_context(simulation_env):
    @veris.mock
    async def test_func() -> dict:
        return {"result": "real"}

    mock_response = {"result": {"mocked": True}}

    with patch("veris_ai.tool_mock.httpx.AsyncClient") as mock_client:
        mock_response_obj = Mock()
        mock_response_obj.json.return_value = mock_response
        mock_response_obj.raise_for_status.return_value = None

        mock_client.return_value.__aenter__.return_value.post = AsyncMock(
            return_value=mock_response_obj,
        )
        veris.set_session_id("test-session-id")
        result = await test_func()
        first_call = mock_client.return_value.__aenter__.return_value.post.call_args
        assert first_call.kwargs["json"]["session_id"] == "test-session-id"
        assert result == {"mocked": True}


@pytest.mark.asyncio
async def test_mock_without_context(simulation_env):
    @veris.mock
    async def test_func() -> dict:
        return {"result": "real"}

    mock_response = {"result": {"mocked": True}}

    with patch("veris_ai.tool_mock.httpx.AsyncClient") as mock_client:
        mock_response_obj = Mock()
        mock_response_obj.json.return_value = mock_response
        mock_response_obj.raise_for_status.return_value = None

        mock_client.return_value.__aenter__.return_value.post = AsyncMock(
            return_value=mock_response_obj,
        )
        veris.clear_session_id()
        result = await test_func()
        first_call = mock_client.return_value.__aenter__.return_value.post.call_args
        assert first_call.kwargs["json"]["session_id"] == None
        assert result == {"mocked": True}


# Test error handling
@pytest.mark.asyncio
async def test_mock_http_error(simulation_env):
    @veris.mock
    async def test_func(param1: str) -> dict:
        return {"result": "real"}

    with patch("httpx.AsyncClient") as mock_client:
        mock_client.return_value.__aenter__.return_value.post = AsyncMock(
            side_effect=httpx.HTTPError("Mock HTTP Error"),
        )

        with pytest.raises(httpx.HTTPError):
            await test_func("test")


@pytest.mark.asyncio
async def test_mock_missing_endpoint():
    with (
        patch.dict(os.environ, {"VERIS_MOCK_ENDPOINT_URL": ""}),
        pytest.raises(ValueError, match="VERIS_MOCK_ENDPOINT_URL environment variable is not set"),
    ):

        @veris.mock
        async def test_func():
            return {"result": "real"}


@pytest.mark.asyncio
async def test_mock_invalid_endpoint(simulation_env):
    with patch.dict(os.environ, {"VERIS_MOCK_ENDPOINT_URL": ""}), pytest.raises(ValueError):

        @veris.mock
        async def test_func():
            return {"result": "real"}

        await test_func()


@pytest.mark.asyncio
async def test_mock_string_json_response(simulation_env):
    @veris.mock
    async def test_func() -> dict:
        return {"result": "real"}

    mock_response = {"result": '{"key": "value"}'}

    with patch("httpx.AsyncClient") as mock_client:
        mock_response_obj = Mock()
        mock_response_obj.json.return_value = mock_response
        mock_response_obj.raise_for_status.return_value = None

        mock_client.return_value.__aenter__.return_value.post = AsyncMock(
            return_value=mock_response_obj,
        )

        result = await test_func()
        assert result == {"key": "value"}


# Test context vars behavior
@pytest.mark.asyncio
async def test_context_var_isolation():
    """Test that session_id context vars are isolated between concurrent calls."""

    async def set_and_check_session(session_id: str):
        veris.set_session_id(session_id)
        # Verify the session_id is set correctly
        assert veris.session_id == session_id
        # Simulate some async work
        await asyncio.sleep(0.01)
        # Verify the session_id hasn't changed
        assert veris.session_id == session_id
        return session_id

    # Run multiple concurrent tasks with different session IDs
    tasks = [
        set_and_check_session("session-1"),
        set_and_check_session("session-2"),
        set_and_check_session("session-3"),
    ]

    results = await asyncio.gather(*tasks)
    assert results == ["session-1", "session-2", "session-3"]


@pytest.mark.asyncio
async def test_context_var_persistence_in_call(simulation_env):
    """Test that session_id persists throughout a mock function call."""
    captured_session_ids = []

    @veris.mock
    async def test_func() -> dict:
        # Capture session_id inside the function
        captured_session_ids.append(veris.session_id)
        return {"result": "real"}

    mock_response = {"result": {"mocked": True}}

    with patch("veris_ai.tool_mock.httpx.AsyncClient") as mock_client:
        mock_response_obj = Mock()
        mock_response_obj.json.return_value = mock_response
        mock_response_obj.raise_for_status.return_value = None

        async def mock_post(*args, **kwargs):
            # Capture session_id during HTTP call
            captured_session_ids.append(veris.session_id)
            return mock_response_obj

        mock_client.return_value.__aenter__.return_value.post = AsyncMock(side_effect=mock_post)

        # Set session_id and call function
        veris.set_session_id("test-session-123")
        await test_func()

        # Verify session_id was captured
        # In simulation mode, we should have 2 captures (one in function, one in HTTP call)
        # But if the function ran in production mode, we'd only have 1
        assert len(captured_session_ids) >= 1
        assert captured_session_ids[0] == "test-session-123"

        # Verify the HTTP client was called with correct session_id
        mock_client.return_value.__aenter__.return_value.post.assert_called_once()
        call_args = mock_client.return_value.__aenter__.return_value.post.call_args
        assert call_args[1]["json"]["session_id"] == "test-session-123"


@pytest.mark.asyncio
async def test_context_var_cleanup():
    """Test that clearing session_id works correctly."""
    # Set a session_id
    veris.set_session_id("test-session")
    assert veris.session_id == "test-session"

    # Clear it
    veris.clear_session_id()
    assert veris.session_id is None

    # Verify it stays None
    await asyncio.sleep(0.01)
    assert veris.session_id is None


@pytest.mark.asyncio
async def test_context_var_no_interference():
    """Test that context vars don't interfere between different async contexts."""
    results = []

    async def task1():
        veris.set_session_id("task1-session")
        await asyncio.sleep(0.02)
        results.append(("task1", veris.session_id))

    async def task2():
        await asyncio.sleep(0.01)
        # Should be None since we didn't set it in this context
        results.append(("task2", veris.session_id))
        veris.set_session_id("task2-session")
        results.append(("task2-after", veris.session_id))

    await asyncio.gather(task1(), task2())

    # Verify results
    assert ("task1", "task1-session") in results
    assert ("task2", None) in results
    assert ("task2-after", "task2-session") in results
