from benchmark.tau_bench.envs.airline.data import load_data
from benchmark.tau_bench.envs.airline.rules import RULES
from benchmark.tau_bench.envs.airline.tools import ALL_TOOLS
from benchmark.tau_bench.envs.airline.wiki import WIKI
from benchmark.tau_bench.envs.base import Env
from benchmark.tau_bench.envs.user import UserStrategy


class MockAirlineDomainEnv(Env):
    def __init__(
        self,
        user_strategy: str | UserStrategy = UserStrategy.LLM,
        user_model: str = "gpt-4o",
        user_provider: str | None = None,
        task_split: str = "test",
        task_index: int | None = None,
    ) -> None:
        match task_split:
            case "test":
                from benchmark.tau_bench.envs.airline.tasks_test import TASKS as tasks
            case _:
                raise ValueError(f"Unknown task split: {task_split}")
        super().__init__(
            data_load_func=load_data,
            tools=ALL_TOOLS,
            tasks=tasks,
            wiki=WIKI,
            rules=RULES,
            user_strategy=user_strategy,
            user_model=user_model,
            user_provider=user_provider,
            task_index=task_index,
        )
        self.terminate_tools: list[str] = ["transfer_to_human_agents"]
