from dataclasses import asdict, dataclass, field
from typing import Any

from minisweagent.agents import get_agent
from minisweagent.config import get_config_from_spec
from minisweagent.environments import get_environment
from minisweagent.exceptions import Submitted
from minisweagent.models import get_model
from minisweagent.models.test_models import make_output
from minisweagent.utils.serialize import recursive_merge


@dataclass
class MemoryEnvironmentConfig:
    env: dict[str, str] = field(default_factory=dict)


class MemoryEnvironment:
    def __init__(self, **kwargs):
        self.config = MemoryEnvironmentConfig(**kwargs)
        self.commands: list[str] = []

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

    def get_template_vars(self, **kwargs) -> dict[str, Any]:
        return {"environment_name": "memory", **kwargs}

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


base = get_config_from_spec("mini")
config = recursive_merge(
    base,
    get_config_from_spec("agent.mode=yolo"),
    get_config_from_spec("agent.confirm_exit=false"),
    {
        "agent": {
            "system_template": "You are an offline teaching agent.",
            "instance_template": "Task: {{ task }}",
            "step_limit": 4,
            "cost_limit": 0,
        },
        "model": {
            "model_class": "deterministic",
            "model_name": "offline-lesson",
            "cost_per_call": 0,
            "outputs": [
                make_output("I will record one action.", [{"command": "remember merged config"}], cost=0),
                make_output("I will finish.", [{"command": "finish"}], cost=0),
            ],
        },
        "environment": {"environment_class": f"{__name__}.MemoryEnvironment"},
    },
)

model = get_model(config=config["model"])
environment = get_environment(config["environment"])
agent = get_agent(model, environment, config["agent"], default_type="interactive")

print("base mode:", base["agent"]["mode"])
print("merged mode:", config["agent"]["mode"])
print("agent class:", type(agent).__name__)
print("environment class:", type(environment).__name__)
print("result:", agent.run("Learn config, factories, and InteractiveAgent"))
print("recorded commands:", environment.commands)