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


class MemoryEnvironment:
    def __init__(self):
        self.config = {"name": "memory"}
        self.actions: list[str] = []

    def execute(self, action: dict, cwd: str = "") -> dict:
        command = action["command"]
        self.actions.append(command)
        if command == "finish":
            raise Submitted(
                {
                    "role": "exit",
                    "content": "offline complete",
                    "extra": {"exit_status": "Submitted", "submission": "offline complete"},
                }
            )
        return {"output": f"recorded: {command}", "returncode": 0, "exception_info": ""}

    def get_template_vars(self, **kwargs) -> dict:
        return {"workspace": "memory"}

    def serialize(self) -> dict:
        return {"info": {"config": {"environment": self.config}}}


env = MemoryEnvironment()
model = DeterministicModel(
    outputs=[
        make_output("I will inspect", [{"command": "read files"}], cost=0.25),
        make_output("I will finish", [{"command": "finish"}], cost=0.25),
    ],
    cost_per_call=0.0,
)
agent = DefaultAgent(
    model=model,
    env=env,
    system_template="You are offline. Calls={{n_model_calls}}.",
    instance_template="Task={{task}}; workspace={{workspace}}.",
    step_limit=3,
    cost_limit=1.0,
)

print(agent.run("learn the loop"))
print(env.actions)
print([message.get("role", message.get("type")) for message in agent.messages])
print(agent.n_calls, agent.cost)