"""
CivicAI Agent Memory

Stores and retrieves historical trajectories, similar decisions, 
and past consequences to help agents avoid repeating mistakes.
"""
from __future__ import annotations
from typing import Any, Dict, List
import json
import os

class AgentMemory:
    def __init__(self, storage_path: str = "assets/agent_memory.json"):
        self.storage_path = storage_path
        self.records: List[Dict[str, Any]] = []
        self._load()

    def _load(self):
        if os.path.exists(self.storage_path):
            try:
                with open(self.storage_path, "r") as f:
                    self.records = json.load(f)
            except Exception:
                self.records = []

    def save(self):
        os.makedirs(os.path.dirname(self.storage_path), exist_ok=True)
        with open(self.storage_path, "w") as f:
            json.dump(self.records, f, indent=2)

    def add_record(self, turn: int, context: str, action: Dict[str, Any], reward: float, critique: str):
        self.records.append({
            "turn": turn,
            "context": context,
            "action": action,
            "reward": reward,
            "critique": critique
        })
        self.save()

    def get_recent_history(self, k: int = 3) -> List[Dict[str, Any]]:
        return self.records[-k:] if self.records else []

    def get_worst_mistake(self) -> Dict[str, Any] | None:
        if not self.records:
            return None
        return min(self.records, key=lambda x: x["reward"])