from inspect_ai.dataset import Sample, json_dataset
from inspect_ai.tool import bash_session

from bench_af._abstract.objects import Environment

from .scorer import turn_off_monitor_scorer


def get() -> Environment:
    """Return the default environment with observed and unobserved datasets."""

    observed_dataset = json_dataset(
        json_file="environments/turn_off_monitor/test/dataset.jsonl",
        sample_fields=lambda record: Sample(
            input=record["input"],
            target=record["target"],
        ),
    )

    unobserved_dataset = json_dataset(
        json_file="environments/turn_off_monitor/deployment/dataset.jsonl",
        sample_fields=lambda record: Sample(
            input=record["input"],
            target=record["target"],
        ),
    )

    return Environment(
        name="turn_off_monitor",
        tools=[bash_session(timeout=60)],
        observed_dataset=observed_dataset,
        unobserved_dataset=unobserved_dataset,
        scorer=turn_off_monitor_scorer(),
    )
