import pytest
import asyncio
import uvicorn
import aiohttp
from threading import Thread
from typing import Generator
from openreward import AsyncOpenReward
from openreward.environments import Environment, Server, tool, ToolOutput
from openreward.environments.types import Blocks, TextBlock, JSONObject


class Foo(Environment):
    def setup(self):
        pass

    def teardown(self):
        pass

    def get_prompt(self) -> Blocks:
        return [TextBlock(text=str(self.task_spec["foo"]))]

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

    @classmethod
    def list_tasks(cls, split: str) -> list[JSONObject]:
        assert split == "train"
        return [{"foo": "bar"}]

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


class Bar(Environment):
    def setup(self):
        pass

    def teardown(self):
        pass

    def get_prompt(self) -> Blocks:
        return [TextBlock(text=str(self.task_spec["bar"]))]

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

    @classmethod
    def list_tasks(cls, split: str) -> list[JSONObject]:
        assert split == "test"
        return [{"bar": "baz"}]

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


class AsyncBaz(Environment):
    """Environment with async list_tasks to test maybe_await in the server."""

    def setup(self):
        pass

    def teardown(self):
        pass

    def get_prompt(self) -> Blocks:
        return [TextBlock(text=str(self.task_spec["baz"]))]

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

    @classmethod
    async def list_tasks(cls, split: str) -> list[JSONObject]:
        assert split == "train"
        return [{"baz": "qux"}]

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

class LargeEnv(Environment):
    """Environment that overrides num_tasks/get_task to avoid materializing all tasks."""

    def setup(self):
        pass

    def teardown(self):
        pass

    def get_prompt(self) -> Blocks:
        return [TextBlock(text=f"task_{self.task_spec['id']}")]

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

    @classmethod
    def list_tasks(cls, split: str) -> list[JSONObject]:
        raise NotImplementedError("LargeEnv doesn't support listing all tasks")

    @classmethod
    async def num_tasks(cls, split: str) -> int:
        """Return a fixed count without materializing tasks."""
        return {"train": 1000, "test": 200}[split]

    @classmethod
    async def get_task(cls, split: str, index: int) -> JSONObject:
        """Generate a task on-the-fly by index."""
        limit = {"train": 1000, "test": 200}[split]
        if index < 0 or index >= limit:
            raise IndexError(f"index {index} out of range for split {split}")
        return {"id": index, "split": split}

    @tool
    async def submit(self) -> ToolOutput:
        return ToolOutput(
            blocks=[TextBlock(text=f"result_{self.task_spec['id']}")],
            reward=self.task_spec["id"] / 1000,
            finished=True,
        )

async def wait_for_server(base_url: str, timeout: float = 5.0):
    """Wait for server to be ready using aiohttp."""
    import time
    start = time.monotonic()
    async with aiohttp.ClientSession() as session:
        while time.monotonic() - start < timeout:
            try:
                async with session.get(f"{base_url}/health", timeout=aiohttp.ClientTimeout(total=0.5)) as resp:
                    if resp.status == 200:
                        return
            except Exception:
                pass
            await asyncio.sleep(0.1)
    pytest.fail("Server failed to start")

@pytest.fixture(scope="module")
def server() -> Generator[str, None, None]:
    """Start the server in a background thread and yield the base URL."""
    host = "localhost"
    port = 8080
    app = Server(environments=[Foo, Bar, AsyncBaz, LargeEnv]).app
    config = uvicorn.Config(app, host=host, port=port, log_level="warning")
    server_instance = uvicorn.Server(config)

    thread = Thread(target=server_instance.run, daemon=True)
    thread.start()

    base_url = f"http://{host}:{port}"
    asyncio.run(wait_for_server(base_url))

    yield base_url
    server_instance.should_exit = True


@pytest.fixture
def client() -> AsyncOpenReward:
    """Create an async client."""
    return AsyncOpenReward(api_key="test")


# Tests for default environment (no variant specified, uses first env)

@pytest.mark.asyncio
async def test_default_variant_splits(client: AsyncOpenReward, server: str):
    """Test that default variant works - redirects to first environment."""
    environment = client.environments.get("foo", base_url=server)
    splits = await environment.list_splits()
    assert splits == ["train"]


@pytest.mark.asyncio
async def test_default_variant_tools(client: AsyncOpenReward, server: str):
    """Test listing tools with default variant."""
    environment = client.environments.get("foo", base_url=server)
    tools = await environment.list_tools()
    tool_names = [t.name for t in tools]
    assert "submit" in tool_names


@pytest.mark.asyncio
async def test_default_variant_list_tasks(client: AsyncOpenReward, server: str):
    """Test listing tasks with default variant."""
    environment = client.environments.get("foo", base_url=server)
    tasks = await environment.list_tasks(split="train")
    assert len(tasks) == 1
    assert tasks[0].task_spec == {"foo": "bar"}


@pytest.mark.asyncio
async def test_default_variant_call_tool(client: AsyncOpenReward, server: str):
    """Test calling tool with default variant."""
    environment = client.environments.get("foo", base_url=server)
    tasks = await environment.list_tasks(split="train")
    async with environment.session(tasks[0]) as session:
        res = await session.call_tool("submit")
        assert res.reward == 1.0
        assert res.finished is True
        assert len(res.blocks) == 1
        assert res.blocks[0].type == "text"
        assert res.blocks[0].text == "foo_result"


# Tests for explicit variants

@pytest.mark.asyncio
async def test_explicit_variant_foo_splits(client: AsyncOpenReward, server: str):
    """Test Foo environment with explicit variant."""
    environment = client.environments.get("foo", variant="foo", base_url=server)
    splits = await environment.list_splits()
    assert splits == ["train"]


@pytest.mark.asyncio
async def test_explicit_variant_bar_splits(client: AsyncOpenReward, server: str):
    """Test Bar environment with explicit variant."""
    environment = client.environments.get("bar", variant="bar", base_url=server)
    splits = await environment.list_splits()
    assert splits == ["test"]


@pytest.mark.asyncio
async def test_explicit_variant_bar_tools(client: AsyncOpenReward, server: str):
    """Test Bar environment tools with explicit variant."""
    environment = client.environments.get("bar", variant="bar", base_url=server)
    tools = await environment.list_tools()
    tool_names = [t.name for t in tools]
    assert "submit" in tool_names


@pytest.mark.asyncio
async def test_explicit_variant_bar_list_tasks(client: AsyncOpenReward, server: str):
    """Test Bar environment task listing with explicit variant."""
    environment = client.environments.get("bar", variant="bar", base_url=server)
    tasks = await environment.list_tasks(split="test")
    assert len(tasks) == 1
    assert tasks[0].task_spec == {"bar": "baz"}


@pytest.mark.asyncio
async def test_explicit_variant_bar_call_tool(client: AsyncOpenReward, server: str):
    """Test Bar environment tool call with explicit variant."""
    environment = client.environments.get("bar", variant="bar", base_url=server)
    tasks = await environment.list_tasks(split="test")
    async with environment.session(tasks[0]) as session:
        res = await session.call_tool("submit")
        assert res.reward == 0.5
        assert res.finished is True
        assert len(res.blocks) == 1
        assert res.blocks[0].type == "text"
        assert res.blocks[0].text == "bar_result"


# Tests for async list_tools

@pytest.mark.asyncio
async def test_async_list_tools_returns_correct_schema(client: AsyncOpenReward, server: str):
    """Test that async list_tools returns tool specs with correct structure."""
    environment = client.environments.get("foo", base_url=server)
    tools = await environment.list_tools()
    assert len(tools) >= 1
    submit_tool = next(t for t in tools if t.name == "submit")
    assert submit_tool.name == "submit"
    assert submit_tool.input_schema is None  # submit takes no input model


@pytest.mark.asyncio
async def test_async_list_tools_with_provider_format(client: AsyncOpenReward, server: str):
    """Test that async list_tools works with provider-specific formats."""
    environment = client.environments.get("foo", base_url=server)

    openai_tools = await environment.list_tools(format="openai")
    assert len(openai_tools) >= 1
    assert openai_tools[0]["type"] == "function"
    assert openai_tools[0]["name"] == "submit"

    anthropic_tools = await environment.list_tools(format="anthropic")
    assert len(anthropic_tools) >= 1
    assert anthropic_tools[0]["type"] == "custom"
    assert anthropic_tools[0]["name"] == "submit"
    assert anthropic_tools[0]["input_schema"] == {"type": "object", "properties": {}}

    google_tools = await environment.list_tools(format="google")
    assert len(google_tools) >= 1
    assert google_tools[0]["name"] == "submit"


@pytest.mark.asyncio
async def test_async_list_tools_on_session(client: AsyncOpenReward, server: str):
    """Test that list_tools works on an active async session."""
    environment = client.environments.get("foo", base_url=server)
    tasks = await environment.list_tasks(split="train")
    async with environment.session(tasks[0]) as session:
        tools = await session.list_tools()
        tool_names = [t.name for t in tools]
        assert "submit" in tool_names


@pytest.mark.asyncio
async def test_async_list_tools_session_with_provider_format(client: AsyncOpenReward, server: str):
    """Test that session list_tools works with provider-specific formats."""
    environment = client.environments.get("bar", variant="bar", base_url=server)
    tasks = await environment.list_tasks(split="test")
    async with environment.session(tasks[0]) as session:
        openai_tools = await session.list_tools(format="openai")
        assert len(openai_tools) >= 1
        assert openai_tools[0]["type"] == "function"


# Tests for environment with async classmethods

@pytest.mark.asyncio
async def test_async_env_list_tools(client: AsyncOpenReward, server: str):
    """Test that list_tools works on an environment with async list_tasks."""
    environment = client.environments.get("asyncbaz", variant="asyncbaz", base_url=server)
    tools = await environment.list_tools()
    tool_names = [t.name for t in tools]
    assert "submit" in tool_names


@pytest.mark.asyncio
async def test_async_env_list_splits(client: AsyncOpenReward, server: str):
    """Test that list_splits works on an env with async list_tasks."""
    environment = client.environments.get("asyncbaz", variant="asyncbaz", base_url=server)
    splits = await environment.list_splits()
    assert splits == ["train"]


@pytest.mark.asyncio
async def test_async_env_list_tasks(client: AsyncOpenReward, server: str):
    """Test that async list_tasks is handled correctly."""
    environment = client.environments.get("asyncbaz", variant="asyncbaz", base_url=server)
    tasks = await environment.list_tasks(split="train")
    assert len(tasks) == 1
    assert tasks[0].task_spec == {"baz": "qux"}


@pytest.mark.asyncio
async def test_async_env_call_tool(client: AsyncOpenReward, server: str):
    """Test full session flow on an environment with async classmethods."""
    environment = client.environments.get("asyncbaz", variant="asyncbaz", base_url=server)
    tasks = await environment.list_tasks(split="train")
    async with environment.session(tasks[0]) as session:
        res = await session.call_tool("submit")
        assert res.reward == 0.75
        assert res.finished is True
        assert res.blocks[0].text == "baz_result"


# Tests for multiple variants interaction

@pytest.mark.asyncio
async def test_multiple_variants_different_splits(client: AsyncOpenReward, server: str):
    """Test that different variants return different splits."""
    foo_env = client.environments.get("foo", variant="foo", base_url=server)
    bar_env = client.environments.get("bar", variant="bar", base_url=server)

    foo_splits = await foo_env.list_splits()
    bar_splits = await bar_env.list_splits()

    assert foo_splits == ["train"]
    assert bar_splits == ["test"]
    assert foo_splits != bar_splits


@pytest.mark.asyncio
async def test_multiple_variants_different_tasks(client: AsyncOpenReward, server: str):
    """Test that different variants return different tasks."""
    foo_env = client.environments.get("foo", variant="foo", base_url=server)
    bar_env = client.environments.get("bar", variant="bar", base_url=server)

    foo_tasks = await foo_env.list_tasks(split="train")
    bar_tasks = await bar_env.list_tasks(split="test")

    assert foo_tasks[0].task_spec == {"foo": "bar"}
    assert bar_tasks[0].task_spec == {"bar": "baz"}


@pytest.mark.asyncio
async def test_multiple_variants_concurrent_sessions(client: AsyncOpenReward, server: str):
    """Test running sessions on multiple variants concurrently."""
    foo_env = client.environments.get("foo", variant="foo", base_url=server)
    bar_env = client.environments.get("bar", variant="bar", base_url=server)

    foo_tasks = await foo_env.list_tasks(split="train")
    bar_tasks = await bar_env.list_tasks(split="test")

    async with foo_env.session(foo_tasks[0]) as foo_session:
        async with bar_env.session(bar_tasks[0]) as bar_session:
            foo_res = await foo_session.call_tool("submit")
            bar_res = await bar_session.call_tool("submit")

            assert foo_res.reward == 1.0
            assert foo_res.blocks[0].text == "foo_result"

            assert bar_res.reward == 0.5
            assert bar_res.blocks[0].text == "bar_result"


@pytest.mark.asyncio
async def test_multiple_variants_prompt_isolation(client: AsyncOpenReward, server: str):
    """Test that prompts are correctly isolated between variants."""
    foo_env = client.environments.get("foo", variant="foo", base_url=server)
    bar_env = client.environments.get("bar", variant="bar", base_url=server)

    foo_tasks = await foo_env.list_tasks(split="train")
    bar_tasks = await bar_env.list_tasks(split="test")

    async with foo_env.session(foo_tasks[0]) as foo_session:
        foo_prompt = await foo_session.get_prompt()
        assert foo_prompt[0].text == "bar"  # from {"foo": "bar"}

    async with bar_env.session(bar_tasks[0]) as bar_session:
        bar_prompt = await bar_session.get_prompt()
        assert bar_prompt[0].text == "baz"  # from {"bar": "baz"}

# Tests for index-based API (num_tasks, get_task)

@pytest.mark.asyncio
async def test_num_tasks_foo(server: str):
    """Test num_tasks returns correct count for Foo."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/foo/num_tasks",
            json={"split": "train"},
        ) as resp:
            assert resp.status == 200
            data = await resp.json()
            assert data["num_tasks"] == 1


@pytest.mark.asyncio
async def test_num_tasks_bar(server: str):
    """Test num_tasks returns correct count for Bar."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/bar/num_tasks",
            json={"split": "test"},
        ) as resp:
            assert resp.status == 200
            data = await resp.json()
            assert data["num_tasks"] == 1


@pytest.mark.asyncio
async def test_num_tasks_invalid_split(server: str):
    """Test num_tasks rejects an invalid split."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/foo/num_tasks",
            json={"split": "nonexistent"},
        ) as resp:
            assert resp.status == 400


@pytest.mark.asyncio
async def test_num_tasks_async_env(server: str):
    """Test num_tasks works for an environment with async list_tasks."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/asyncbaz/num_tasks",
            json={"split": "train"},
        ) as resp:
            assert resp.status == 200
            data = await resp.json()
            assert data["num_tasks"] == 1


@pytest.mark.asyncio
async def test_get_task_foo(server: str):
    """Test get_task returns the correct task by index."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/foo/task",
            json={"split": "train", "index": 0},
        ) as resp:
            assert resp.status == 200
            data = await resp.json()
            assert data["task"] == {"foo": "bar"}


@pytest.mark.asyncio
async def test_get_task_bar(server: str):
    """Test get_task on Bar environment."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/bar/task",
            json={"split": "test", "index": 0},
        ) as resp:
            assert resp.status == 200
            data = await resp.json()
            assert data["task"] == {"bar": "baz"}


@pytest.mark.asyncio
async def test_get_task_invalid_split(server: str):
    """Test get_task rejects an invalid split."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/foo/task",
            json={"split": "nonexistent", "index": 0},
        ) as resp:
            assert resp.status == 400


@pytest.mark.asyncio
async def test_get_task_index_out_of_bounds(server: str):
    """Test get_task returns 400 for an out-of-bounds index."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/foo/task",
            json={"split": "train", "index": 999},
        ) as resp:
            assert resp.status == 400


@pytest.mark.asyncio
async def test_get_task_async_env(server: str):
    """Test get_task works for an environment with async list_tasks."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/asyncbaz/task",
            json={"split": "train", "index": 0},
        ) as resp:
            assert resp.status == 200
            data = await resp.json()
            assert data["task"] == {"baz": "qux"}


@pytest.mark.asyncio
async def test_get_task_unknown_env(server: str):
    """Test get_task returns 404 for a nonexistent environment."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/doesnotexist/task",
            json={"split": "train", "index": 0},
        ) as resp:
            assert resp.status == 404


# Tests for LargeEnv (overridden num_tasks / get_task)

@pytest.mark.asyncio
async def test_large_env_num_tasks_train(server: str):
    """Test overridden num_tasks returns custom count for train."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/largeenv/num_tasks",
            json={"split": "train"},
        ) as resp:
            assert resp.status == 200
            data = await resp.json()
            assert data["num_tasks"] == 1000


@pytest.mark.asyncio
async def test_large_env_num_tasks_test(server: str):
    """Test overridden num_tasks returns custom count for test."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/largeenv/num_tasks",
            json={"split": "test"},
        ) as resp:
            assert resp.status == 200
            data = await resp.json()
            assert data["num_tasks"] == 200


@pytest.mark.asyncio
async def test_large_env_get_task_first(server: str):
    """Test overridden get_task returns correct task at index 0."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/largeenv/task",
            json={"split": "train", "index": 0},
        ) as resp:
            assert resp.status == 200
            data = await resp.json()
            assert data["task"] == {"id": 0, "split": "train"}


@pytest.mark.asyncio
async def test_large_env_get_task_mid(server: str):
    """Test overridden get_task returns correct task at an arbitrary index."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/largeenv/task",
            json={"split": "train", "index": 500},
        ) as resp:
            assert resp.status == 200
            data = await resp.json()
            assert data["task"] == {"id": 500, "split": "train"}


@pytest.mark.asyncio
async def test_large_env_get_task_last(server: str):
    """Test overridden get_task at the last valid index."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/largeenv/task",
            json={"split": "test", "index": 199},
        ) as resp:
            assert resp.status == 200
            data = await resp.json()
            assert data["task"] == {"id": 199, "split": "test"}


@pytest.mark.asyncio
async def test_large_env_get_task_out_of_bounds(server: str):
    """Test overridden get_task rejects out-of-bounds index."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/largeenv/task",
            json={"split": "train", "index": 1000},
        ) as resp:
            assert resp.status == 400


@pytest.mark.asyncio
async def test_large_env_get_task_invalid_split(server: str):
    """Test overridden get_task rejects invalid split."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/largeenv/num_tasks",
            json={"split": "nonexistent"},
        ) as resp:
            assert resp.status == 400


@pytest.mark.asyncio
async def test_large_env_session_from_get_task(server: str):
    """Test full session flow using split/index to create session."""
    sid = str(__import__("uuid").uuid4())
    headers = {"X-Session-ID": sid}
    async with aiohttp.ClientSession() as http:
        # Create session with split/index instead of task_spec
        async with http.post(
            f"{server}/create",
            json={"env_name": "largeenv", "split": "train", "index": 42},
            headers=headers,
        ) as resp:
            assert resp.status == 200

        # Call tool — server resolved task_spec internally
        async with http.post(
            f"{server}/largeenv/call",
            json={"name": "submit", "input": {}},
            headers=headers,
        ) as resp:
            assert resp.status == 200
            body = await resp.text()
            assert "result_42" in body

        # Cleanup
        async with http.post(f"{server}/delete", headers=headers) as resp:
            assert resp.status == 200


@pytest.mark.asyncio
async def test_create_session_rejects_both_spec_and_index(server: str):
    """Test that providing both task_spec and split/index is rejected."""
    sid = str(__import__("uuid").uuid4())
    async with aiohttp.ClientSession() as http:
        async with http.post(
            f"{server}/create",
            json={"env_name": "foo", "task_spec": {"foo": "bar"}, "split": "train", "index": 0},
            headers={"X-Session-ID": sid},
        ) as resp:
            assert resp.status == 422


@pytest.mark.asyncio
async def test_create_session_rejects_neither_spec_nor_index(server: str):
    """Test that providing neither task_spec nor split/index is rejected."""
    sid = str(__import__("uuid").uuid4())
    async with aiohttp.ClientSession() as http:
        async with http.post(
            f"{server}/create",
            json={"env_name": "foo"},
            headers={"X-Session-ID": sid},
        ) as resp:
            assert resp.status == 422


@pytest.mark.asyncio
async def test_create_session_rejects_split_without_index(server: str):
    """Test that providing split alone is rejected."""
    sid = str(__import__("uuid").uuid4())
    async with aiohttp.ClientSession() as http:
        async with http.post(
            f"{server}/create",
            json={"env_name": "foo", "split": "train"},
            headers={"X-Session-ID": sid},
        ) as resp:
            assert resp.status == 422


@pytest.mark.asyncio
async def test_create_session_invalid_split(server: str):
    """Test that an invalid split in create is rejected."""
    sid = str(__import__("uuid").uuid4())
    async with aiohttp.ClientSession() as http:
        async with http.post(
            f"{server}/create",
            json={"env_name": "foo", "split": "nonexistent", "index": 0},
            headers={"X-Session-ID": sid},
        ) as resp:
            assert resp.status == 400


@pytest.mark.asyncio
async def test_create_session_index_out_of_bounds(server: str):
    """Test that an out-of-bounds index in create is rejected."""
    sid = str(__import__("uuid").uuid4())
    async with aiohttp.ClientSession() as http:
        async with http.post(
            f"{server}/create",
            json={"env_name": "largeenv", "split": "train", "index": 9999},
            headers={"X-Session-ID": sid},
        ) as resp:
            assert resp.status == 400


# Tests for index-based API via AsyncOpenReward client

@pytest.mark.asyncio
async def test_client_num_tasks_foo(client: AsyncOpenReward, server: str):
    """Test num_tasks via the client SDK."""
    environment = client.environments.get("foo", base_url=server)
    count = await environment.num_tasks("train")
    assert count == 1


@pytest.mark.asyncio
async def test_client_num_tasks_bar(client: AsyncOpenReward, server: str):
    """Test num_tasks on Bar via the client SDK."""
    environment = client.environments.get("bar", variant="bar", base_url=server)
    count = await environment.num_tasks("test")
    assert count == 1


@pytest.mark.asyncio
async def test_client_get_task_foo(client: AsyncOpenReward, server: str):
    """Test get_task via the client SDK."""
    environment = client.environments.get("foo", base_url=server)
    task = await environment.get_task("train", 0)
    assert task.task_spec == {"foo": "bar"}


@pytest.mark.asyncio
async def test_client_get_task_bar(client: AsyncOpenReward, server: str):
    """Test get_task on Bar via the client SDK."""
    environment = client.environments.get("bar", variant="bar", base_url=server)
    task = await environment.get_task("test", 0)
    assert task.task_spec == {"bar": "baz"}


@pytest.mark.asyncio
async def test_client_get_task_async_env(client: AsyncOpenReward, server: str):
    """Test get_task on an environment with async list_tasks."""
    environment = client.environments.get("asyncbaz", variant="asyncbaz", base_url=server)
    task = await environment.get_task("train", 0)
    assert task.task_spec == {"baz": "qux"}


@pytest.mark.asyncio
async def test_client_session_with_split_index(client: AsyncOpenReward, server: str):
    """Test creating a session via split/index instead of a Task object."""
    environment = client.environments.get("foo", base_url=server)
    async with environment.session(split="train", index=0) as session:
        res = await session.call_tool("submit")
        assert res.reward == 1.0
        assert res.finished is True
        assert res.blocks[0].text == "foo_result"


@pytest.mark.asyncio
async def test_client_session_with_split_index_bar(client: AsyncOpenReward, server: str):
    """Test split/index session on Bar."""
    environment = client.environments.get("bar", variant="bar", base_url=server)
    async with environment.session(split="test", index=0) as session:
        res = await session.call_tool("submit")
        assert res.reward == 0.5
        assert res.blocks[0].text == "bar_result"


@pytest.mark.asyncio
async def test_client_session_with_split_index_prompt(client: AsyncOpenReward, server: str):
    """Test that get_prompt works on a split/index session."""
    environment = client.environments.get("foo", base_url=server)
    async with environment.session(split="train", index=0) as session:
        prompt = await session.get_prompt()
        assert prompt[0].text == "bar"


@pytest.mark.asyncio
async def test_client_session_with_split_index_list_tools(client: AsyncOpenReward, server: str):
    """Test that list_tools works on a split/index session."""
    environment = client.environments.get("foo", base_url=server)
    async with environment.session(split="train", index=0) as session:
        tools = await session.list_tools()
        assert any(t.name == "submit" for t in tools)


@pytest.mark.asyncio
async def test_client_session_rejects_both_task_and_index(client: AsyncOpenReward, server: str):
    """Test that providing both task and split/index raises ValueError."""
    environment = client.environments.get("foo", base_url=server)
    task = await environment.get_task("train", 0)
    with pytest.raises(ValueError, match="either task or both split and index"):
        environment.session(task=task, split="train", index=0)


@pytest.mark.asyncio
async def test_client_session_rejects_neither_task_nor_index(client: AsyncOpenReward, server: str):
    """Test that providing neither task nor split/index raises ValueError."""
    environment = client.environments.get("foo", base_url=server)
    with pytest.raises(ValueError, match="either task or both split and index"):
        environment.session()


@pytest.mark.asyncio
async def test_client_session_rejects_split_without_index(client: AsyncOpenReward, server: str):
    """Test that providing split alone raises ValueError."""
    environment = client.environments.get("foo", base_url=server)
    with pytest.raises(ValueError, match="either task or both split and index"):
        environment.session(split="train")

# Tests for LargeEnv via the client SDK

@pytest.mark.asyncio
async def test_client_large_env_num_tasks(client: AsyncOpenReward, server: str):
    """Test num_tasks on LargeEnv via the client."""
    environment = client.environments.get("largeenv", variant="largeenv", base_url=server)
    assert await environment.num_tasks("train") == 1000
    assert await environment.num_tasks("test") == 200


@pytest.mark.asyncio
async def test_client_large_env_get_task(client: AsyncOpenReward, server: str):
    """Test get_task on LargeEnv via the client."""
    environment = client.environments.get("largeenv", variant="largeenv", base_url=server)
    task = await environment.get_task("train", 42)
    assert task.task_spec == {"id": 42, "split": "train"}


@pytest.mark.asyncio
async def test_client_large_env_session_with_task(client: AsyncOpenReward, server: str):
    """Test LargeEnv session created from a get_task result."""
    environment = client.environments.get("largeenv", variant="largeenv", base_url=server)
    task = await environment.get_task("train", 42)
    async with environment.session(task=task) as session:
        res = await session.call_tool("submit")
        assert "result_42" in res.blocks[0].text


@pytest.mark.asyncio
async def test_client_large_env_session_with_split_index(client: AsyncOpenReward, server: str):
    """Test LargeEnv session created directly from split/index."""
    environment = client.environments.get("largeenv", variant="largeenv", base_url=server)
    async with environment.session(split="train", index=99) as session:
        res = await session.call_tool("submit")
        assert "result_99" in res.blocks[0].text


@pytest.mark.asyncio
async def test_client_large_env_concurrent_index_sessions(client: AsyncOpenReward, server: str):
    """Test multiple concurrent split/index sessions on LargeEnv."""
    environment = client.environments.get("largeenv", variant="largeenv", base_url=server)
    async with environment.session(split="train", index=0) as s1:
        async with environment.session(split="train", index=500) as s2:
            r1 = await s1.call_tool("submit")
            r2 = await s2.call_tool("submit")
            assert "result_0" in r1.blocks[0].text
            assert "result_500" in r2.blocks[0].text

# Tests for optional env_name (server defaults to first environment)

@pytest.mark.asyncio
async def test_create_session_defaults_env_name(server: str):
    """Test that omitting env_name defaults to the first registered environment (Foo)."""
    sid = str(__import__("uuid").uuid4())
    headers = {"X-Session-ID": sid}
    async with aiohttp.ClientSession() as http:
        async with http.post(
            f"{server}/create",
            json={"split": "train", "index": 0},
            headers=headers,
        ) as resp:
            assert resp.status == 200

        # Foo is the first env, so prompt should come from {"foo": "bar"}
        async with http.get(
            f"{server}/foo/prompt",
            headers=headers,
        ) as resp:
            assert resp.status == 200
            data = await resp.json()
            assert data[0]["text"] == "bar"

        async with http.post(f"{server}/delete", headers=headers) as resp:
            assert resp.status == 200


@pytest.mark.asyncio
async def test_create_session_rejects_no_task_source_without_env_name(server: str):
    """Test that omitting everything (no env_name, no task_spec, no split/index) is rejected."""
    sid = str(__import__("uuid").uuid4())
    async with aiohttp.ClientSession() as http:
        async with http.post(
            f"{server}/create",
            json={},
            headers={"X-Session-ID": sid},
        ) as resp:
            assert resp.status == 422


@pytest.mark.asyncio
async def test_create_session_without_env_name_with_split_index(server: str):
    """Test split/index without env_name — server defaults to first env and resolves task."""
    sid = str(__import__("uuid").uuid4())
    headers = {"X-Session-ID": sid}
    async with aiohttp.ClientSession() as http:
        async with http.post(
            f"{server}/create",
            json={"split": "train", "index": 0},
            headers=headers,
        ) as resp:
            assert resp.status == 200

        # Call submit on the default (Foo) environment
        async with http.post(
            f"{server}/foo/call",
            json={"name": "submit", "input": {}},
            headers=headers,
        ) as resp:
            assert resp.status == 200
            body = await resp.text()
            assert "foo_result" in body

        async with http.post(f"{server}/delete", headers=headers) as resp:
            assert resp.status == 200


# Tests for get_task_range server endpoint

@pytest.mark.asyncio
async def test_get_task_range_basic(server: str):
    """Test get_task_range returns tasks for a basic range."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/largeenv/task_range",
            json={"split": "train", "start": 0, "stop": 3},
        ) as resp:
            assert resp.status == 200
            data = await resp.json()
            assert len(data["tasks"]) == 3
            assert data["tasks"][0] == {"id": 0, "split": "train"}
            assert data["tasks"][2] == {"id": 2, "split": "train"}


@pytest.mark.asyncio
async def test_get_task_range_none_start(server: str):
    """Test get_task_range with None start defaults to 0."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/largeenv/task_range",
            json={"split": "test", "stop": 2},
        ) as resp:
            assert resp.status == 200
            data = await resp.json()
            assert len(data["tasks"]) == 2
            assert data["tasks"][0] == {"id": 0, "split": "test"}


@pytest.mark.asyncio
async def test_get_task_range_none_stop(server: str):
    """Test get_task_range with None stop defaults to num_tasks."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/largeenv/task_range",
            json={"split": "test", "start": 198},
        ) as resp:
            assert resp.status == 200
            data = await resp.json()
            assert len(data["tasks"]) == 2
            assert data["tasks"][0] == {"id": 198, "split": "test"}
            assert data["tasks"][1] == {"id": 199, "split": "test"}


@pytest.mark.asyncio
async def test_get_task_range_negative_start(server: str):
    """Test get_task_range with negative start (relative to end)."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/largeenv/task_range",
            json={"split": "test", "start": -3},
        ) as resp:
            assert resp.status == 200
            data = await resp.json()
            assert len(data["tasks"]) == 3
            assert data["tasks"][0] == {"id": 197, "split": "test"}
            assert data["tasks"][2] == {"id": 199, "split": "test"}


@pytest.mark.asyncio
async def test_get_task_range_negative_stop(server: str):
    """Test get_task_range with negative stop."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/largeenv/task_range",
            json={"split": "test", "start": 0, "stop": -198},
        ) as resp:
            assert resp.status == 200
            data = await resp.json()
            assert len(data["tasks"]) == 2
            assert data["tasks"][0] == {"id": 0, "split": "test"}
            assert data["tasks"][1] == {"id": 1, "split": "test"}


@pytest.mark.asyncio
async def test_get_task_range_empty(server: str):
    """Test get_task_range returns empty list when start >= stop."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/largeenv/task_range",
            json={"split": "test", "start": 5, "stop": 5},
        ) as resp:
            assert resp.status == 200
            data = await resp.json()
            assert data["tasks"] == []


@pytest.mark.asyncio
async def test_get_task_range_invalid_split(server: str):
    """Test get_task_range rejects invalid split."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/largeenv/task_range",
            json={"split": "nonexistent", "start": 0, "stop": 1},
        ) as resp:
            assert resp.status == 400


@pytest.mark.asyncio
async def test_get_task_range_foo(server: str):
    """Test get_task_range on a small environment."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{server}/foo/task_range",
            json={"split": "train", "start": 0, "stop": 1},
        ) as resp:
            assert resp.status == 200
            data = await resp.json()
            assert len(data["tasks"]) == 1
            assert data["tasks"][0] == {"foo": "bar"}


# Tests for get_task_range via AsyncOpenReward client

@pytest.mark.asyncio
async def test_client_get_task_range_basic(client: AsyncOpenReward, server: str):
    """Test get_task_range via the client SDK."""
    environment = client.environments.get("largeenv", variant="largeenv", base_url=server)
    tasks = await environment.get_task_range("train", 0, 3)
    assert len(tasks) == 3
    assert tasks[0].task_spec == {"id": 0, "split": "train"}
    assert tasks[2].task_spec == {"id": 2, "split": "train"}


@pytest.mark.asyncio
async def test_client_get_task_range_defaults(client: AsyncOpenReward, server: str):
    """Test get_task_range with default start/stop via client."""
    environment = client.environments.get("largeenv", variant="largeenv", base_url=server)
    tasks = await environment.get_task_range("test")
    assert len(tasks) == 200


@pytest.mark.asyncio
async def test_client_get_task_range_negative(client: AsyncOpenReward, server: str):
    """Test get_task_range with negative indices via client."""
    environment = client.environments.get("largeenv", variant="largeenv", base_url=server)
    tasks = await environment.get_task_range("test", -2)
    assert len(tasks) == 2
    assert tasks[0].task_spec == {"id": 198, "split": "test"}
    assert tasks[1].task_spec == {"id": 199, "split": "test"}