import textwrap

import pytest
from httpx import AsyncClient
from httpx_sse import aconnect_sse

from progtc.types import MessageType


@pytest.mark.asyncio
async def test_execute_code_basic(test_client: AsyncClient) -> None:
    code = textwrap.dedent(
        """\
        async def main():
            print("main function called")
        await main()
        """,
    )
    messages = []
    async with aconnect_sse(
        test_client,
        "POST",
        "/execute-code",
        json={"tool_names": [], "code": code},
    ) as event_source:
        async for event in event_source.aiter_sse():
            messages.append(event.json())
    assert messages[0]["type"] == MessageType.SUCCESS
    assert messages[0]["stdout"] == "main function called\n"


@pytest.mark.asyncio
async def test_execute_code_with_func_call(test_client: AsyncClient) -> None:
    code = textwrap.dedent(
        """\
        async def say_hello():
            print("say_hello function called")
        async def main():
            print("main function called")
            await say_hello()
        await main()
        """,
    )
    messages = []
    async with aconnect_sse(
        test_client,
        "POST",
        "/execute-code",
        json={"tool_names": [], "code": code},
    ) as event_source:
        async for event in event_source.aiter_sse():
            messages.append(event.json())
    assert messages[0]["type"] == MessageType.SUCCESS
    assert messages[0]["stdout"] == "main function called\nsay_hello function called\n"


@pytest.mark.asyncio
async def test_execute_code_with_tool_call(test_client: AsyncClient) -> None:
    code = textwrap.dedent(
        """\
        from tools import get_weather
        async def main():
            print("main function called")
            return await get_weather()
        print(await main())
        """,
    )
    messages = []
    async with aconnect_sse(
        test_client,
        "POST",
        "/execute-code",
        json={"tool_names": ["get_weather"], "code": code},
        timeout=10,
    ) as event_source:
        async for event in event_source.aiter_sse():
            message = event.json()
            messages.append(message)
            if (
                message["type"] == MessageType.TOOL_CALL
                and message["tool_name"] == "get_weather"
            ):
                await test_client.post(
                    "/tool-results",
                    json={
                        "execution_id": message["execution_id"],
                        "tool_call_id": message["id"],
                        "result": "Sunny",
                    },
                )
    assert messages[0]["type"] == MessageType.TOOL_CALL, messages[0]
    assert messages[0]["tool_name"] == "get_weather"
    assert messages[0]["args"] == []
    assert messages[0]["kwargs"] == {}
    assert messages[1]["type"] == MessageType.SUCCESS
    assert messages[1]["stdout"] == "main function called\nSunny\n"


@pytest.mark.asyncio
async def test_execute_code_with_tool_call_with_args(test_client: AsyncClient) -> None:
    code = textwrap.dedent(
        """\
        from tools import get_weather
        async def main():
            from tools import get_weather
            print("main function called")
            return await get_weather("UK", city="London")
        print(await main())
        """,
    )
    messages = []
    async with aconnect_sse(
        test_client,
        "POST",
        "/execute-code",
        json={"tool_names": ["get_weather"], "code": code},
        timeout=10,
    ) as event_source:
        async for event in event_source.aiter_sse():
            message = event.json()
            messages.append(message)
            if (
                message["type"] == MessageType.TOOL_CALL
                and message["tool_name"] == "get_weather"
            ):
                await test_client.post(
                    "/tool-results",
                    json={
                        "execution_id": message["execution_id"],
                        "tool_call_id": message["id"],
                        "result": "Sunny",
                    },
                )
    assert messages[0]["type"] == MessageType.TOOL_CALL, messages[0]
    assert messages[0]["tool_name"] == "get_weather"
    assert messages[0]["args"] == ["UK"]
    assert messages[0]["kwargs"] == {"city": "London"}
    assert messages[1]["type"] == MessageType.SUCCESS
    assert messages[1]["stdout"] == "main function called\nSunny\n"


@pytest.mark.asyncio
async def test_execute_code_step_basic(test_client: AsyncClient) -> None:
    """Code with no tool calls returns success directly."""
    code = textwrap.dedent("""\
        async def main():
            print("hello")
        await main()
    """)
    response = await test_client.post(
        "/execute-code-step",
        json={"tool_names": [], "code": code},
        timeout=10,
    )
    assert response.status_code == 200
    data = response.json()
    assert data["type"] == MessageType.SUCCESS
    assert data["stdout"] == "hello\n"


@pytest.mark.asyncio
async def test_execute_code_step_syntax_error(test_client: AsyncClient) -> None:
    response = await test_client.post(
        "/execute-code-step",
        json={"tool_names": [], "code": "def ("},
        timeout=10,
    )
    assert response.status_code == 200
    data = response.json()
    assert data["type"] == MessageType.SYNTAX_ERROR


@pytest.mark.asyncio
async def test_execute_code_step_runtime_error(test_client: AsyncClient) -> None:
    response = await test_client.post(
        "/execute-code-step",
        json={"tool_names": [], "code": "raise ValueError('boom')"},
        timeout=10,
    )
    assert response.status_code == 200
    data = response.json()
    assert data["type"] == MessageType.RUNTIME_ERROR


@pytest.mark.asyncio
async def test_execute_code_step_tool_call(test_client: AsyncClient) -> None:
    """First call returns tool calls; second call with tool_results returns success."""
    code = textwrap.dedent("""\
        from tools import get_weather
        result = await get_weather("London")
        print(result)
    """)

    # First request — should return tool calls
    response = await test_client.post(
        "/execute-code-step",
        json={"tool_names": ["get_weather"], "code": code},
        timeout=10,
    )
    assert response.status_code == 200
    data = response.json()
    assert data["type"] == MessageType.TOOL_CALL
    assert len(data["tool_calls"]) == 1
    assert data["tool_calls"][0]["tool_name"] == "get_weather"
    assert data["tool_calls"][0]["args"] == ["London"]

    state = data["state"]

    # Second request — send tool results, server appends to state
    response = await test_client.post(
        "/execute-code-step",
        json={
            "tool_names": ["get_weather"],
            "code": code,
            "state": state,
            "tool_results": ["Sunny"],
        },
        timeout=10,
    )
    assert response.status_code == 200
    data = response.json()
    assert data["type"] == MessageType.SUCCESS
    assert data["stdout"] == "Sunny\n"


@pytest.mark.asyncio
@pytest.mark.parametrize("iteration", range(10))
async def test_execute_code_step_gather(
    test_client: AsyncClient, iteration: int
) -> None:
    """Concurrent tool calls via asyncio.gather return as a batch
    and replay deterministically across repeated runs."""
    cities = [f"city_{i}" for i in range(20)]
    gather_args = ", ".join(f'get_weather("{c}")' for c in cities)
    code = textwrap.dedent(f"""\
        import asyncio
        from tools import get_weather
        results = await asyncio.gather({gather_args})
        print(",".join(results))
    """)
    results = [f"weather_{i}" for i in range(20)]

    # First request — should return all tool calls as a batch
    response = await test_client.post(
        "/execute-code-step",
        json={"tool_names": ["get_weather"], "code": code},
        timeout=10,
    )
    data = response.json()
    assert data["type"] == MessageType.TOOL_CALL
    assert len(data["tool_calls"]) == 20
    for i, tc in enumerate(data["tool_calls"]):
        assert tc["args"] == [cities[i]], f"tool call {i} has wrong args"

    state = data["state"]

    # Second request — send all results, expect success
    response = await test_client.post(
        "/execute-code-step",
        json={
            "tool_names": ["get_weather"],
            "code": code,
            "state": state,
            "tool_results": results,
        },
        timeout=10,
    )
    data = response.json()
    assert data["type"] == MessageType.SUCCESS
    assert data["stdout"] == ",".join(results) + "\n"


@pytest.mark.asyncio
async def test_execute_code_step_multiple_tool_calls(
    test_client: AsyncClient,
) -> None:
    """Sequential tool calls: each round-trip sends tool_results back."""
    code = textwrap.dedent("""\
        from tools import get_weather
        a = await get_weather("London")
        b = await get_weather("Paris")
        print(f"{a},{b}")
    """)

    # First request — returns the first tool call
    response = await test_client.post(
        "/execute-code-step",
        json={"tool_names": ["get_weather"], "code": code},
        timeout=10,
    )
    data = response.json()
    assert data["type"] == MessageType.TOOL_CALL
    assert data["tool_calls"][0]["args"] == ["London"]

    state = data["state"]

    # Second request — send first result, get second tool call
    response = await test_client.post(
        "/execute-code-step",
        json={
            "tool_names": ["get_weather"],
            "code": code,
            "state": state,
            "tool_results": ["Sunny"],
        },
        timeout=10,
    )
    data = response.json()
    assert data["type"] == MessageType.TOOL_CALL
    assert data["tool_calls"][0]["args"] == ["Paris"]

    state = data["state"]

    # Third request — send second result, get final success
    response = await test_client.post(
        "/execute-code-step",
        json={
            "tool_names": ["get_weather"],
            "code": code,
            "state": state,
            "tool_results": ["Rainy"],
        },
        timeout=10,
    )
    data = response.json()
    assert data["type"] == MessageType.SUCCESS
    assert data["stdout"] == "Sunny,Rainy\n"


@pytest.mark.asyncio
async def test_execute_code_step_tool_call_with_kwargs(
    test_client: AsyncClient,
) -> None:
    code = textwrap.dedent("""\
        from tools import get_weather
        result = await get_weather("UK", city="London")
        print(result)
    """)

    response = await test_client.post(
        "/execute-code-step",
        json={"tool_names": ["get_weather"], "code": code},
        timeout=10,
    )
    data = response.json()
    assert data["type"] == MessageType.TOOL_CALL
    assert data["tool_calls"][0]["args"] == ["UK"]
    assert data["tool_calls"][0]["kwargs"] == {"city": "London"}

    state = data["state"]

    # Replay with result
    response = await test_client.post(
        "/execute-code-step",
        json={
            "tool_names": ["get_weather"],
            "code": code,
            "state": state,
            "tool_results": ["Sunny"],
        },
        timeout=10,
    )
    data = response.json()
    assert data["type"] == MessageType.SUCCESS
    assert data["stdout"] == "Sunny\n"


@pytest.mark.asyncio
async def test_execute_code_step_determinism_error(
    test_client: AsyncClient,
) -> None:
    """Replaying with different code produces a determinism error."""
    code_v1 = textwrap.dedent("""\
        from tools import get_weather
        result = await get_weather("London")
        print(result)
    """)
    code_v2 = textwrap.dedent("""\
        from tools import get_weather
        result = await get_weather("Paris")
        print(result)
    """)

    # First request with code_v1
    response = await test_client.post(
        "/execute-code-step",
        json={"tool_names": ["get_weather"], "code": code_v1},
        timeout=10,
    )
    data = response.json()
    assert data["type"] == MessageType.TOOL_CALL
    state = data["state"]

    # Second request with code_v2 — tool call args differ from cached hash
    response = await test_client.post(
        "/execute-code-step",
        json={
            "tool_names": ["get_weather"],
            "code": code_v2,
            "state": state,
            "tool_results": ["Sunny"],
        },
        timeout=10,
    )
    data = response.json()
    assert data["type"] == MessageType.DETERMINISM_ERROR
