from dataclasses import asdict, dataclass
from typing import Any

from minisweagent.agents.default import DefaultAgent
from minisweagent.exceptions import Submitted
from minisweagent.models.test_models import DeterministicModel, make_output


AGENT_CONFIG = {
    "system_template": "You are an offline test agent.",
    "instance_template": "Task: {{ task }}",
    "step_limit": 0,
    "cost_limit": 0,
}


def test_recording_environment_contract():
    env = RecordingEnvironment(prefix="seen")

    assert env.execute({"command": "inspect repository"}, cwd="/virtual") == {
        "output": "seen: inspect repository",
        "returncode": 0,
        "exception_info": "",
    }
    assert env.records == [{"command": "inspect repository", "cwd": "/virtual"}]


def test_default_agent_uses_the_extension_end_to_end():
    model = DeterministicModel(
        outputs=[
            make_output("I will inspect.", [{"command": "inspect repository"}], cost=0),
            make_output("I will submit.", [{"command": "submit: offline extension complete"}], cost=0),
        ],
        cost_per_call=0,
    )
    env = RecordingEnvironment()
    agent = DefaultAgent(model, env, **AGENT_CONFIG)

    assert agent.run("Build a recording environment") == {
        "exit_status": "Submitted",
        "submission": "offline extension complete",
    }
    assert env.records == [
        {"command": "inspect repository", "cwd": ""},
        {"command": "submit: offline extension complete", "cwd": ""},
    ]
    assert [message.get("role") for message in agent.messages] == [
        "system",
        "user",
        "assistant",
        "user",
        "assistant",
        "exit",
    ]


def test_step_limit_prevents_the_next_action():
    model = DeterministicModel(
        outputs=[
            make_output("First.", [{"command": "first"}], cost=0),
            make_output("Second.", [{"command": "second"}], cost=0),
        ],
        cost_per_call=0,
    )
    env = RecordingEnvironment()
    agent = DefaultAgent(model, env, **(AGENT_CONFIG | {"step_limit": 1}))

    assert agent.run("Stop after one model call") == {"exit_status": "LimitsExceeded", "submission": ""}
    assert env.records == [{"command": "first", "cwd": ""}]
    assert agent.n_calls == 1


@dataclass
class RecordingEnvironmentConfig:
    prefix: str = "recorded"


class RecordingEnvironment:
    def __init__(self, *, config_class=RecordingEnvironmentConfig, **kwargs):
        self.config = config_class(**kwargs)
        self.records: list[dict[str, str]] = []

    def execute(self, action: dict, cwd: str = "") -> dict[str, Any]:
        command = action.get("command", "")
        self.records.append({"command": command, "cwd": cwd})
        if command.startswith("submit:"):
            submission = command.removeprefix("submit:").strip()
            raise Submitted(
                {
                    "role": "exit",
                    "content": submission,
                    "extra": {"exit_status": "Submitted", "submission": submission},
                }
            )
        return {"output": f"{self.config.prefix}: {command}", "returncode": 0, "exception_info": ""}

    def get_template_vars(self, **kwargs) -> dict[str, Any]:
        return {**asdict(self.config), **kwargs}

    def serialize(self) -> dict:
        return {
            "info": {
                "config": {
                    "environment": asdict(self.config),
                    "environment_type": f"{self.__class__.__module__}.{self.__class__.__name__}",
                }
            }
        }