"""Conversation memory for multi-turn agent context.

Stores inter-agent messages per workflow run, providing context
for agents that need to reference earlier decisions or outputs.
"""

from __future__ import annotations

import json
import logging
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)


class ConversationMemory:
    """In-memory + optional file-backed conversation storage per run.

    Messages are stored in chronological order and can be filtered
    by agent name for focused context retrieval.
    """

    def __init__(
        self,
        run_id: str = "",
        persist_dir: Optional[str] = None,
    ):
        self.run_id = run_id
        self._messages: List[Dict[str, Any]] = []
        self._persist_dir = Path(persist_dir) if persist_dir else None

        if self._persist_dir:
            self._persist_dir.mkdir(parents=True, exist_ok=True)
            self._load_from_disk()

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def add_message(
        self,
        sender: str,
        recipient: str,
        content: str,
        message_type: str = "info",
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Add a message to the conversation history.

        Returns the stored message dict.
        """
        msg = {
            "sender": sender,
            "recipient": recipient,
            "content": content,
            "message_type": message_type,
            "metadata": metadata or {},
            "timestamp": datetime.utcnow().isoformat() + "Z",
            "run_id": self.run_id,
        }
        self._messages.append(msg)

        if self._persist_dir:
            self._save_to_disk()

        return msg

    def get_context(
        self,
        agent_name: str,
        limit: int = 20,
    ) -> List[Dict[str, Any]]:
        """Retrieve recent messages relevant to an agent.

        Returns messages where the agent is sender or recipient,
        most recent first, up to ``limit``.
        """
        relevant = [
            m for m in self._messages
            if m.get("sender") == agent_name
            or m.get("recipient") == agent_name
            or m.get("recipient") == ""  # broadcasts
        ]
        return list(reversed(relevant[-limit:]))

    def get_all(self, limit: int = 100) -> List[Dict[str, Any]]:
        """Retrieve all messages, most recent first."""
        return list(reversed(self._messages[-limit:]))

    def get_by_type(
        self,
        message_type: str,
        limit: int = 50,
    ) -> List[Dict[str, Any]]:
        """Retrieve messages by type (info, error, decision, etc.)."""
        matches = [m for m in self._messages if m.get("message_type") == message_type]
        return list(reversed(matches[-limit:]))

    def clear(self) -> int:
        """Clear all messages. Returns count of messages cleared."""
        count = len(self._messages)
        self._messages.clear()
        if self._persist_dir:
            self._save_to_disk()
        return count

    @property
    def message_count(self) -> int:
        return len(self._messages)

    def summary(self) -> Dict[str, Any]:
        """Get a summary of the conversation memory."""
        senders = {}
        types = {}
        for m in self._messages:
            s = m.get("sender", "unknown")
            t = m.get("message_type", "unknown")
            senders[s] = senders.get(s, 0) + 1
            types[t] = types.get(t, 0) + 1

        return {
            "run_id": self.run_id,
            "total_messages": len(self._messages),
            "by_sender": senders,
            "by_type": types,
        }

    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------

    def _save_to_disk(self) -> None:
        """Persist messages to a JSON file."""
        if not self._persist_dir:
            return
        path = self._persist_dir / f"memory_{self.run_id}.json"
        path.write_text(
            json.dumps(self._messages, indent=2, default=str),
            encoding="utf-8",
        )

    def _load_from_disk(self) -> None:
        """Load messages from disk if available."""
        if not self._persist_dir:
            return
        path = self._persist_dir / f"memory_{self.run_id}.json"
        if path.exists():
            try:
                self._messages = json.loads(path.read_text(encoding="utf-8"))
            except (json.JSONDecodeError, OSError):
                self._messages = []
