#!/usr/bin/env python3

import json
import sys
from pathlib import Path
from unittest.mock import patch

from minisweagent.models.test_models import DeterministicModel
from minisweagent.run.mini import DEFAULT_CONFIG_FILE, main


def update_trajectory():
    traj_path = Path(__file__).parent / "local.traj.json"
    trajectory = json.loads(traj_path.read_text())

    task = "Blah blah blah"

    model_responses = [msg["content"] for msg in trajectory[2:] if msg["role"] == "assistant"]
    print(f"Got {len(model_responses)} model responses")

    with patch("minisweagent.run.mini.get_model") as mock_get_model:
        mock_get_model.return_value = DeterministicModel(outputs=model_responses)
        main(model_name="tardis", config_spec=DEFAULT_CONFIG_FILE, output=traj_path, task=task, yolo=True, model_class=None)

if __name__ == "__main__":
    update_trajectory()