import pytest
import aiohttp
from openreward import AsyncOpenReward
from openreward.environments import Environment, tool, ToolOutput
from openreward.environments.types import Blocks, TextBlock, JSONObject
from openreward.api.sandboxes.types import SandboxSettings, SandboxBucketConfig, PodTerminatedError


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)


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


@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."""
    import os
    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