"""
Memory Store for Breakpilot Agents

Hybrid Valkey + PostgreSQL persistence with TTL, access tracking, and pattern search.
"""

from typing import List, Dict, Any, Optional
from datetime import datetime, timezone, timedelta
import json
import logging
import hashlib

from agent_core.brain.memory_models import Memory

logger = logging.getLogger(__name__)


class MemoryStore:
    """Long-term memory store with TTL, access tracking, and hybrid persistence."""

    def __init__(self, redis_client=None, db_pool=None, namespace: str = "breakpilot"):
        self.redis = redis_client
        self.db_pool = db_pool
        self.namespace = namespace
        self._local_cache: Dict[str, Memory] = {}

    def _redis_key(self, key: str) -> str:
        """Generate Redis key with namespace"""
        return f"{self.namespace}:memory:{key}"

    def _hash_key(self, key: str) -> str:
        """Generate a hash for long keys"""
        if len(key) > 200:
            return hashlib.sha256(key.encode()).hexdigest()[:32]
        return key

    async def remember(
        self, key: str, value: Any, agent_id: str,
        ttl_days: int = 30, metadata: Optional[Dict[str, Any]] = None
    ) -> Memory:
        """Stores a memory with optional TTL and metadata."""
        expires_at = None
        if ttl_days > 0:
            expires_at = datetime.now(timezone.utc) + timedelta(days=ttl_days)

        memory = Memory(
            key=key,
            value=value,
            agent_id=agent_id,
            expires_at=expires_at,
            metadata=metadata or {}
        )

        # Store in all layers
        await self._store_memory(memory, ttl_days)

        logger.debug(f"Agent {agent_id} remembered '{key}' (TTL: {ttl_days} days)")

        return memory

    async def recall(self, key: str) -> Optional[Any]:
        """Retrieves a memory value by key, or None if not found/expired."""
        memory = await self.get_memory(key)
        if memory:
            return memory.value
        return None

    async def get_memory(self, key: str) -> Optional[Memory]:
        """Retrieves a full Memory object by key, updating access count."""
        # Check local cache
        if key in self._local_cache:
            memory = self._local_cache[key]
            if not memory.is_expired():
                await self._update_access(memory)
                return memory
            else:
                del self._local_cache[key]

        # Try Valkey
        memory = await self._get_from_valkey(key)
        if memory and not memory.is_expired():
            self._local_cache[key] = memory
            await self._update_access(memory)
            return memory

        # Try PostgreSQL
        memory = await self._get_from_postgres(key)
        if memory and not memory.is_expired():
            # Re-cache in Valkey
            await self._cache_in_valkey(memory)
            self._local_cache[key] = memory
            await self._update_access(memory)
            return memory

        return None

    async def forget(self, key: str) -> bool:
        """
        Deletes a memory.

        Args:
            key: The memory key to delete

        Returns:
            True if deleted, False if not found
        """
        # Remove from local cache
        self._local_cache.pop(key, None)

        # Remove from Valkey
        if self.redis:
            await self.redis.delete(self._redis_key(key))

        # Mark as deleted in PostgreSQL
        if self.db_pool:
            async with self.db_pool.acquire() as conn:
                result = await conn.execute(
                    """
                    DELETE FROM agent_memory
                    WHERE namespace = $1 AND key = $2
                    """,
                    self.namespace,
                    key
                )
                return "DELETE" in result

        return True

    async def search(
        self,
        pattern: str,
        agent_id: Optional[str] = None,
        limit: int = 100
    ) -> List[Memory]:
        """
        Searches for memories matching a pattern.

        Args:
            pattern: SQL LIKE pattern (e.g., "evaluation:math:%")
            agent_id: Optional filter by agent ID
            limit: Maximum results to return

        Returns:
            List of matching Memory objects
        """
        results = []

        # Search PostgreSQL (primary source for patterns)
        if self.db_pool:
            async with self.db_pool.acquire() as conn:
                query = """
                    SELECT key, value, agent_id, created_at, expires_at,
                           access_count, last_accessed, metadata
                    FROM agent_memory
                    WHERE namespace = $1
                      AND key LIKE $2
                      AND (expires_at IS NULL OR expires_at > NOW())
                """
                params = [self.namespace, pattern.replace("*", "%")]

                if agent_id:
                    query += " AND agent_id = $3"
                    params.append(agent_id)

                query += f" ORDER BY access_count DESC, created_at DESC LIMIT {limit}"

                rows = await conn.fetch(query, *params)

                for row in rows:
                    results.append(self._row_to_memory(row))
        else:
            # Fall back to local cache search
            import fnmatch
            for key, memory in self._local_cache.items():
                if fnmatch.fnmatch(key, pattern):
                    if agent_id is None or memory.agent_id == agent_id:
                        if not memory.is_expired():
                            results.append(memory)
                            if len(results) >= limit:
                                break

        return results

    async def get_by_agent(
        self,
        agent_id: str,
        limit: int = 100
    ) -> List[Memory]:
        """
        Gets all memories for a specific agent.

        Args:
            agent_id: The agent ID
            limit: Maximum results

        Returns:
            List of Memory objects
        """
        return await self.search("*", agent_id=agent_id, limit=limit)

    async def get_recent(
        self,
        hours: int = 24,
        agent_id: Optional[str] = None,
        limit: int = 100
    ) -> List[Memory]:
        """
        Gets recently created memories.

        Args:
            hours: How far back to look
            agent_id: Optional filter by agent
            limit: Maximum results

        Returns:
            List of Memory objects
        """
        cutoff = datetime.now(timezone.utc) - timedelta(hours=hours)

        if self.db_pool:
            async with self.db_pool.acquire() as conn:
                query = """
                    SELECT key, value, agent_id, created_at, expires_at,
                           access_count, last_accessed, metadata
                    FROM agent_memory
                    WHERE namespace = $1
                      AND created_at > $2
                      AND (expires_at IS NULL OR expires_at > NOW())
                """
                params = [self.namespace, cutoff]

                if agent_id:
                    query += " AND agent_id = $3"
                    params.append(agent_id)

                query += f" ORDER BY created_at DESC LIMIT {limit}"

                rows = await conn.fetch(query, *params)
                return [self._row_to_memory(row) for row in rows]

        # Fall back to local cache
        results = []
        for memory in self._local_cache.values():
            if memory.created_at > cutoff:
                if agent_id is None or memory.agent_id == agent_id:
                    if not memory.is_expired():
                        results.append(memory)

        return sorted(results, key=lambda m: m.created_at, reverse=True)[:limit]

    async def get_most_accessed(
        self,
        limit: int = 10,
        agent_id: Optional[str] = None
    ) -> List[Memory]:
        """
        Gets the most frequently accessed memories.

        Args:
            limit: Number of results
            agent_id: Optional filter by agent

        Returns:
            List of Memory objects ordered by access count
        """
        if self.db_pool:
            async with self.db_pool.acquire() as conn:
                query = """
                    SELECT key, value, agent_id, created_at, expires_at,
                           access_count, last_accessed, metadata
                    FROM agent_memory
                    WHERE namespace = $1
                      AND (expires_at IS NULL OR expires_at > NOW())
                """
                params = [self.namespace]

                if agent_id:
                    query += " AND agent_id = $2"
                    params.append(agent_id)

                query += f" ORDER BY access_count DESC LIMIT {limit}"

                rows = await conn.fetch(query, *params)
                return [self._row_to_memory(row) for row in rows]

        # Fall back to local cache
        valid = [m for m in self._local_cache.values() if not m.is_expired()]
        if agent_id:
            valid = [m for m in valid if m.agent_id == agent_id]
        return sorted(valid, key=lambda m: m.access_count, reverse=True)[:limit]

    async def cleanup_expired(self) -> int:
        """
        Removes expired memories.

        Returns:
            Number of memories removed
        """
        count = 0

        # Clean local cache
        expired_keys = [
            key for key, memory in self._local_cache.items()
            if memory.is_expired()
        ]
        for key in expired_keys:
            del self._local_cache[key]
            count += 1

        # Clean PostgreSQL
        if self.db_pool:
            async with self.db_pool.acquire() as conn:
                result = await conn.execute(
                    """
                    DELETE FROM agent_memory
                    WHERE namespace = $1 AND expires_at < NOW()
                    """,
                    self.namespace
                )
                # Parse count from result
                if result:
                    parts = result.split()
                    if len(parts) >= 2:
                        count += int(parts[1])

        logger.info(f"Cleaned up {count} expired memories")
        return count

    async def _store_memory(self, memory: Memory, ttl_days: int) -> None:
        """Stores memory in all persistence layers"""
        self._local_cache[memory.key] = memory
        await self._cache_in_valkey(memory, ttl_days)
        await self._save_to_postgres(memory)

    async def _cache_in_valkey(
        self,
        memory: Memory,
        ttl_days: Optional[int] = None
    ) -> None:
        """Caches memory in Valkey"""
        if not self.redis:
            return

        try:
            ttl_seconds = (ttl_days or 30) * 86400
            await self.redis.setex(
                self._redis_key(memory.key),
                ttl_seconds,
                json.dumps(memory.to_dict())
            )
        except Exception as e:
            logger.warning(f"Failed to cache memory in Valkey: {e}")

    async def _get_from_valkey(self, key: str) -> Optional[Memory]:
        """Retrieves memory from Valkey"""
        if not self.redis:
            return None

        try:
            data = await self.redis.get(self._redis_key(key))
            if data:
                return Memory.from_dict(json.loads(data))
        except Exception as e:
            logger.warning(f"Failed to get memory from Valkey: {e}")

        return None

    async def _save_to_postgres(self, memory: Memory) -> None:
        """Saves memory to PostgreSQL"""
        if not self.db_pool:
            return

        try:
            async with self.db_pool.acquire() as conn:
                await conn.execute(
                    """
                    INSERT INTO agent_memory
                        (namespace, key, value, agent_id, access_count,
                         created_at, expires_at, last_accessed, metadata)
                    VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
                    ON CONFLICT (namespace, key) DO UPDATE SET
                        value = EXCLUDED.value,
                        access_count = agent_memory.access_count + 1,
                        last_accessed = NOW(),
                        metadata = EXCLUDED.metadata
                    """,
                    self.namespace,
                    memory.key,
                    json.dumps(memory.value),
                    memory.agent_id,
                    memory.access_count,
                    memory.created_at,
                    memory.expires_at,
                    memory.last_accessed,
                    json.dumps(memory.metadata)
                )
        except Exception as e:
            logger.error(f"Failed to save memory to PostgreSQL: {e}")

    async def _get_from_postgres(self, key: str) -> Optional[Memory]:
        """Retrieves memory from PostgreSQL"""
        if not self.db_pool:
            return None

        try:
            async with self.db_pool.acquire() as conn:
                row = await conn.fetchrow(
                    """
                    SELECT key, value, agent_id, created_at, expires_at,
                           access_count, last_accessed, metadata
                    FROM agent_memory
                    WHERE namespace = $1 AND key = $2
                    """,
                    self.namespace,
                    key
                )

                if row:
                    return self._row_to_memory(row)
        except Exception as e:
            logger.error(f"Failed to get memory from PostgreSQL: {e}")

        return None

    async def _update_access(self, memory: Memory) -> None:
        """Updates access count and timestamp"""
        memory.access_count += 1
        memory.last_accessed = datetime.now(timezone.utc)

        # Update in PostgreSQL
        if self.db_pool:
            try:
                async with self.db_pool.acquire() as conn:
                    await conn.execute(
                        """
                        UPDATE agent_memory
                        SET access_count = access_count + 1,
                            last_accessed = NOW()
                        WHERE namespace = $1 AND key = $2
                        """,
                        self.namespace,
                        memory.key
                    )
            except Exception as e:
                logger.warning(f"Failed to update access count: {e}")

    def _row_to_memory(self, row) -> Memory:
        """Converts a database row to Memory"""
        value = row["value"]
        if isinstance(value, str):
            value = json.loads(value)

        metadata = row.get("metadata", {})
        if isinstance(metadata, str):
            metadata = json.loads(metadata)

        return Memory(
            key=row["key"],
            value=value,
            agent_id=row["agent_id"],
            created_at=row["created_at"],
            expires_at=row["expires_at"],
            access_count=row["access_count"] or 0,
            last_accessed=row["last_accessed"],
            metadata=metadata
        )
