"""In-memory working memory store with TTL support.

Replaces Redis for working memory (Tier 1) operations.
Working memory is session-scoped and ephemeral — no persistence needed.
"""

import asyncio
import json
import logging
import time
from typing import Any

logger = logging.getLogger(__name__)

# Key: (session_id, memory_id) -> (memory_dict, expires_at_timestamp)
_store: dict[tuple[str, str], tuple[dict, float]] = {}

# Index: session_id -> set of memory_ids
_index: dict[str, set[str]] = {}

# Background cleanup task reference
_cleanup_task: asyncio.Task | None = None

CLEANUP_INTERVAL = 30  # seconds


# ─── Store Operations ────────────────────────────────────────────────────────

async def store(session_id: str, memory_id: str, memory_data: dict, ttl_seconds: int) -> None:
    """Store a working memory entry with TTL."""
    expires_at = time.time() + ttl_seconds
    _store[(session_id, memory_id)] = (memory_data, expires_at)

    if session_id not in _index:
        _index[session_id] = set()
    _index[session_id].add(memory_id)

    logger.debug(f"Working memory stored: {session_id}/{memory_id} (TTL={ttl_seconds}s)")


async def get(session_id: str, memory_id: str) -> dict | None:
    """Get a single working memory entry."""
    key = (session_id, memory_id)
    entry = _store.get(key)
    if entry is None:
        return None

    memory_data, expires_at = entry
    if time.time() > expires_at:
        # Expired — clean up
        del _store[key]
        _index.get(session_id, set()).discard(memory_id)
        return None

    return memory_data


async def get_all(session_id: str) -> list[dict]:
    """Get all working memories for a session (non-expired)."""
    memory_ids = _index.get(session_id, set()).copy()
    result = []
    now = time.time()

    for mid in memory_ids:
        key = (session_id, mid)
        entry = _store.get(key)
        if entry is None:
            continue
        memory_data, expires_at = entry
        if now > expires_at:
            del _store[key]
            _index[session_id].discard(mid)
            continue
        result.append(memory_data)

    return result


async def query(session_id: str, query_text: str, limit: int = 10) -> list[dict]:
    """Keyword search in working memory for a session."""
    all_memories = await get_all(session_id)
    query_lower = query_text.lower()
    query_terms = query_lower.split()

    scored = []
    for mem in all_memories:
        content = mem.get("content", "").lower()
        tags = " ".join(mem.get("topic_tags", [])).lower()
        searchable = content + " " + tags

        score = sum(1 for term in query_terms if term in searchable)
        if score > 0:
            scored.append((score, mem))

    scored.sort(key=lambda x: x[0], reverse=True)
    return [mem for _, mem in scored[:limit]]


async def get_ttl(session_id: str, memory_id: str) -> int:
    """Get remaining TTL in seconds for a working memory entry."""
    key = (session_id, memory_id)
    entry = _store.get(key)
    if entry is None:
        return -1
    _, expires_at = entry
    remaining = int(expires_at - time.time())
    return max(remaining, 0)


async def count(session_id: str) -> int:
    """Count active (non-expired) working memories for a session."""
    memories = await get_all(session_id)  # This cleans expired ones
    return len(memories)


async def flush(session_id: str) -> int:
    """Delete all working memories for a session. Returns count deleted."""
    memory_ids = _index.pop(session_id, set())
    deleted = 0
    for mid in memory_ids:
        key = (session_id, mid)
        if key in _store:
            del _store[key]
            deleted += 1
    logger.info(f"Flushed {deleted} working memories for session {session_id}")
    return deleted


async def flush_all() -> int:
    """Delete all working memories across all sessions."""
    total = len(_store)
    _store.clear()
    _index.clear()
    return total


# ─── Background Cleanup ─────────────────────────────────────────────────────

async def _cleanup_loop():
    """Periodically remove expired entries."""
    while True:
        try:
            await asyncio.sleep(CLEANUP_INTERVAL)
            now = time.time()
            expired_keys = [
                key for key, (_, expires_at) in _store.items()
                if now > expires_at
            ]
            for key in expired_keys:
                session_id, memory_id = key
                del _store[key]
                _index.get(session_id, set()).discard(memory_id)

            if expired_keys:
                logger.debug(f"Cleaned up {len(expired_keys)} expired working memories")
        except asyncio.CancelledError:
            break
        except Exception as e:
            logger.error(f"Working memory cleanup error: {e}")


def start_cleanup():
    """Start the background cleanup task."""
    global _cleanup_task
    if _cleanup_task is None or _cleanup_task.done():
        _cleanup_task = asyncio.create_task(_cleanup_loop())
        logger.info("Working memory cleanup task started")


def stop_cleanup():
    """Stop the background cleanup task."""
    global _cleanup_task
    if _cleanup_task and not _cleanup_task.done():
        _cleanup_task.cancel()
        _cleanup_task = None
