"""Example sandbox environment server generated by `orwd init`."""
from typing import List

from openreward import AsyncOpenReward, SandboxBucketConfig, SandboxSettings
from openreward.environments import (Environment, JSONObject, TextBlock,
                                     ToolOutput, tool)
from pydantic import BaseModel


class BashParams(BaseModel, extra="forbid"):
    command: str

class EnvironmentSpec(BaseModel):
    task_id: str


class SandboxEnv(Environment):
    def __init__(self, task_spec: JSONObject, secrets: dict[str, str] = {}) -> None:
        super().__init__(task_spec)

        self.validated = EnvironmentSpec.model_validate(task_spec)
        self.task_id = self.validated.task_id

        if not secrets.get("api_key"):
            raise ValueError("OpenReward API key is required")

        self.sandbox_settings = SandboxSettings(
            environment="YourUsername/SandboxEnv",
            image="generalreasoning/python-ds:3.12-tools",
            machine_size="0.5:1",
            block_network=False,
            bucket_config=SandboxBucketConfig(
                mount_path="/tmp/sandbox/",
                read_only=True
            )
        )

        or_client = AsyncOpenReward(api_key=secrets.get("api_key"))
        self.sandbox = or_client.sandbox(self.sandbox_settings)

    async def setup(self) -> None:
        await self.sandbox.start()

    async def teardown(self) -> None:
        await self.sandbox.stop()

    @tool
    async def bash(self, params: BashParams) -> ToolOutput:
        """Executes a bash command in the environment."""

        result = await self.sandbox.run(params.command.strip())
        output, code = result
        truncated = result.truncated

        if truncated:
            output = f"...(truncated, output exceeded limit)\n{output}"

        return ToolOutput(
            blocks=[TextBlock(text=f"{output}\n\n(exit {code})")],
            metadata={"output": output, "exit_code": code, "truncated": truncated},
            reward=0.0,
            finished=False,
        )

    async def get_prompt(self) -> List[TextBlock]:
        """Return the challenge prompt."""
        full_prompt = f"""You are participating in task {self.task_id}. Please explore the file system."""

        return [TextBlock(text=full_prompt)]

    @classmethod
    def list_tasks(cls, split: str) -> List[JSONObject]:
        """Get all available CTF challenge tasks."""

        if split == "test":
            return [{"task_id": "0"}]
        elif split == "train":
            return [{"task_id": "1"}]
        else:
            raise ValueError(f"Unknown split: {split}")

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