import os

import aiohttp
import httpx
import pytest

from openreward import AsyncOpenReward
from openreward.api.sandboxes.types import (SandboxSettings)
from openreward.environments import Environment, tool, ToolOutput
from openreward.environments.server import Server
from openreward.environments.types import (Blocks, JSONObject, TextBlock)


class SandboxEnv(Environment):
    """Minimal env that starts/stops a sandbox using provided settings."""

    def __init__(self, task_spec: JSONObject, secrets: dict[str, str] = {}, settings: SandboxSettings | None = None) -> None:
        super().__init__(task_spec, secrets=secrets)
        self._client = AsyncOpenReward()
        self.sandbox = self._client.sandbox(settings)

    async def setup(self):
        await self.sandbox.start()

    async def teardown(self):
        await self.sandbox.stop()

    def get_prompt(self) -> Blocks:
        return [TextBlock(text="test")]

    @classmethod
    def list_splits(cls) -> list[str]:
        return ["train"]

    @classmethod
    def list_tasks(cls, split: str) -> list[JSONObject]:
        return [{}]

    @tool
    async def submit(self) -> ToolOutput:
        return ToolOutput(blocks=[TextBlock(text="done")], reward=1.0, finished=True)


class FailingEnv(Environment):
    """Env whose list_tasks always raises, used to trigger unhandled exceptions."""

    def get_prompt(self) -> Blocks:
        return [TextBlock(text="test")]

    @classmethod
    def list_splits(cls) -> list[str]:
        return ["train"]

    @classmethod
    def list_tasks(cls, split: str) -> list[JSONObject]:
        raise RuntimeError("something broke internally")

    @tool
    async def submit(self) -> ToolOutput:
        return ToolOutput(blocks=[TextBlock(text="done")], reward=1.0, finished=True)


def _make_server(return_errors: str) -> Server:
    """Build a Server with FailingEnv and the given return_errors mode."""
    return Server([FailingEnv], return_errors=return_errors)


async def _trigger_error(app) -> httpx.Response:
    """Hit the tasks endpoint which raises RuntimeError in FailingEnv."""
    transport = httpx.ASGITransport(app=app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
        return await client.post("/failingenv/tasks", json={"split": "train"})

@pytest.mark.asyncio
async def test_return_errors_none():
    """With return_errors='none', 500 responses should be opaque."""
    server = _make_server("none")
    resp = await _trigger_error(server.app)

    assert resp.status_code == 500
    body = resp.json()
    assert body["detail"] == "Internal Server Error"
    assert "RuntimeError" not in body["detail"]
    assert "something broke" not in body["detail"]


@pytest.mark.asyncio
async def test_return_errors_exception():
    """With return_errors='exception', 500 responses should include the exception string."""
    server = _make_server("exception")
    resp = await _trigger_error(server.app)

    assert resp.status_code == 500
    body = resp.json()
    assert "RuntimeError" in body["detail"]
    assert "something broke internally" in body["detail"]
    # Should not contain a full traceback
    assert "Traceback" not in body["detail"]


@pytest.mark.asyncio
async def test_return_errors_stacktrace():
    """With return_errors='stacktrace', 500 responses should include the full traceback."""
    server = _make_server("stacktrace")
    resp = await _trigger_error(server.app)

    assert resp.status_code == 500
    body = resp.json()
    assert "Traceback (most recent call last)" in body["detail"]
    assert "RuntimeError: something broke internally" in body["detail"]

async def _start_stop(settings: SandboxSettings):
    env = SandboxEnv(task_spec={}, settings=settings)
    try:
        await env.setup()
    finally:
        await env.teardown()


@pytest.mark.asyncio
async def test_nonexistent_environment():
    """Sandbox referencing a nonexistent OpenReward environment should fail with a clear message."""
    with pytest.raises(aiohttp.ClientResponseError) as exc_info:
        await _start_stop(SandboxSettings(
            environment="GeneralReasoning/idontexist",
            image="python:3.11-slim",
            machine_size="1:2",
        ))
    assert exc_info.value.status == 404
    assert "idontexist" in exc_info.value.message
    assert "GeneralReasoning" in exc_info.value.message


# TODO: This currently hangs indefinitely.
# @pytest.mark.asyncio
# async def test_nonexistent_image():
#     """Sandbox referencing a nonexistent container image should fail with a clear message."""
#     with pytest.raises(RuntimeError) as exc_info:
#         await _start_stop(SandboxSettings(
#             environment="GeneralReasoning/test-env",
#             image="generalreasoning/idontexist:latest",
#             machine_size="1:2",
#         ))
#     assert "idontexist" in str(exc_info.value)
#     assert "ErrImagePull" in str(exc_info.value)


@pytest.mark.asyncio
async def test_missing_api_key():
    """Sandbox creation without an API key should fail with a clear auth message."""
    key = os.environ.pop("OPENREWARD_API_KEY", None)
    try:
        with pytest.raises(Exception) as exc_info:
            await _start_stop(SandboxSettings(
                environment="GeneralReasoning/test-env",
                image="python:3.11-slim",
                machine_size="1:2",
            ))
        assert "Authentication Failed" in str(exc_info.value)
    finally:
        if key is not None:
            os.environ["OPENREWARD_API_KEY"] = key
