#!/usr/bin/env python3
"""Four-tier agent memory store — working, episodic, semantic, procedural.

Demonstrates: in-memory store with write, consolidate, recall, and staleness handling.
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field


@dataclass
class MemoryEntry:
    content: str
    timestamp: float = field(default_factory=time.time)
    source_session: str = ""
    superseded: bool = False


class AgentMemory:
    """Four-tier memory: working, episodic, semantic, procedural."""

    def __init__(self):
        self.working: list[dict] = []
        self.episodic: list[MemoryEntry] = []
        self.semantic: dict[str, MemoryEntry] = {}
        self.procedural: dict[str, str] = {}
        self.session_id: str = ""

    def start_session(self, session_id: str):
        self.session_id = session_id
        self.working = []

    def observe(self, event: str):
        self.working.append({"ts": time.time(), "event": event})

    def end_session(self):
        for entry in self.working:
            self.episodic.append(
                MemoryEntry(content=entry["event"], timestamp=entry["ts"],
                            source_session=self.session_id)
            )

    def consolidate(self, extract_func):
        """Run extraction function over recent episodes to produce semantic facts."""
        recent = [e for e in self.episodic if e.source_session == self.session_id]
        facts = extract_func(recent)
        for key, value in facts.items():
            self._upsert_semantic(key, value)

    def _upsert_semantic(self, key: str, value: str):
        """Insert or update a fact; supersede old value on conflict."""
        if key in self.semantic:
            self.semantic[key].superseded = True
        self.semantic[key] = MemoryEntry(content=value)

    def recall_semantic(self, query_keys: list[str]) -> dict[str, str]:
        return {
            k: self.semantic[k].content
            for k in query_keys
            if k in self.semantic and not self.semantic[k].superseded
        }

    def recall_recent_episodes(self, n: int = 5) -> list[str]:
        return [e.content for e in self.episodic[-n:]]

    def add_procedure(self, name: str, policy: str, version: int = 1):
        self.procedural[f"{name}@v{version}"] = policy

    def get_procedure(self, name: str) -> str | None:
        versions = sorted(
            [k for k in self.procedural if k.startswith(f"{name}@v")], reverse=True)
        return self.procedural[versions[0]] if versions else None

    def prune_episodic(self, max_age_seconds: float):
        cutoff = time.time() - max_age_seconds
        self.episodic = [e for e in self.episodic if e.timestamp > cutoff]

    def prune_working(self, keep_last: int = 10):
        if len(self.working) > keep_last:
            self.working = self.working[-keep_last:]


def demo_extract(episodes: list[MemoryEntry]) -> dict[str, str]:
    """Mock consolidation: extract user preferences from episodes."""
    text = " ".join(e.content for e in episodes)
    facts = {}
    if "kebab" in text.lower():
        facts["naming_preference"] = "kebab-case"
    if "deploy" in text.lower():
        facts["deploy_status"] = "paused this week"
    return facts


if __name__ == "__main__":
    mem = AgentMemory()

    # Session 1
    mem.start_session("s1")
    mem.observe("User asked about function names")
    mem.observe("I prefer kebab-case for all my functions")
    mem.observe("Deploy is paused this week")
    mem.end_session()
    mem.consolidate(demo_extract)

    print("After session 1 — facts:", mem.recall_semantic(["naming_preference", "deploy_status"]))

    # Session 2 — recall context
    mem.start_session("s2")
    context = mem.recall_semantic(["naming_preference", "deploy_status"])
    print("Session 2 recalled:", context)
    assert context.get("naming_preference") == "kebab-case"

    # Contradiction handling
    mem.observe("Actually, I switched to snake_case")
    mem.end_session()
    mem._upsert_semantic("naming_preference", "snake_case")
    print("After update — naming:", mem.recall_semantic(["naming_preference"]))
    assert mem.semantic["naming_preference"].content == "snake_case"
    print("✓ All assertions passed")