"""Working Memory Compressor"""
import logging
from typing import Optional, Callable

logger = logging.getLogger(__name__)


class MemoryCompressor:
    """Compress working memories when they exceed threshold."""

    def __init__(self, mysql_store, redis_cache, embedding_service):
        self.mysql = mysql_store
        self.redis = redis_cache
        self.embedder = embedding_service

    def compress(self, agent_id: str, max_items: int = 20,
                 summary_callback: Optional[Callable] = None) -> str:
        """
        Compress working memories into a single long-term memory.

        If working memory exceeds max_items, take the oldest items,
        summarize them (via callback or simple concatenation), and
        store as a personal long-term memory.

        Returns: Summary text of compressed memories.
        """
        # Get agent info to find team_id
        agent = self.mysql.get_agent(agent_id)
        if not agent:
            raise ValueError(f"Agent {agent_id} not found")

        items = self.redis.get_working_memories(agent_id, limit=100)
        if len(items) <= max_items:
            return {
                "compressed": 0,
                "remaining": len(items),
                "total_before": len(items),
                "target_count": max_items,
                "summary": None,
            }

        # Take items beyond max_items (oldest)
        to_compress = items[max_items:]
        remaining = items[:max_items]

        # Build summary
        contents = [item.get("content", "") for item in to_compress]
        if summary_callback:
            summary = summary_callback(contents)
        else:
            summary = " | ".join(contents[:50])  # Simple concat, limit 50 items

        # Store as personal long-term memory
        embedding = self.embedder.embed(summary)
        from storage.vector_search import embedding_to_bytes
        emb_bytes = embedding_to_bytes(embedding)

        self.mysql.add_personal_memory(
            agent_id=agent_id,
            team_id=agent["team_id"],
            content=summary,
            embedding=emb_bytes,
            importance=0.6,  # Slightly elevated importance for compressed memories
            metadata={
                "source": "compression",
                "item_count": len(to_compress),
                "source_items": [
                    {"content": item.get("content", "")[:200], "timestamp": item.get("timestamp")}
                    for item in to_compress
                ],
            },
        )

        # Replace working memory with remaining items
        self.redis.clear_working_memory(agent_id)
        for item in reversed(remaining):  # lpush reverses order
            self.redis.add_working_memory(
                agent_id, item["content"],
                metadata=item.get("metadata", {}),
            )

        logger.info(f"Compressed {len(to_compress)} working memories for agent {agent_id}")
        return {
            "compressed": len(to_compress),
            "remaining": len(remaining),
            "total_before": len(items),
            "target_count": max_items,
            "summary": summary,
        }