"""Small bounded pool for periodic memory extraction; no extra service required."""
from concurrent.futures import ThreadPoolExecutor
from copy import deepcopy
import logging
import threading

from engine.memory import MAX_RECENT, update_memory

logger = logging.getLogger(__name__)


class MemoryScheduler:
    def __init__(self, workers=2):
        self._executor = ThreadPoolExecutor(max_workers=workers, thread_name_prefix='memory')
        self._capacity = workers
        self._pending = set()
        self._lock = threading.Lock()

    def schedule(self, session_id, character_id, conversation, client, model, persona=None):
        if len(conversation['recent_messages']) - conversation['summarized_count'] <= MAX_RECENT:
            return None
        key = (session_id, character_id)
        with self._lock:
            # No unbounded queue and at most one job for each conversation.
            # Skipped batches remain intact and are retried after a later reply.
            if key in self._pending or len(self._pending) >= self._capacity:
                return None
            self._pending.add(key)
        try:
            snapshot = deepcopy(conversation)
            future = self._executor.submit(update_memory, session_id, character_id, snapshot, client, model, deepcopy(persona))
        except Exception:
            self._release(key)
            logger.warning('Could not schedule memory update; messages retained.', exc_info=True)
            return None
        future.add_done_callback(lambda _: self._release(key))
        return future

    def _release(self, key):
        with self._lock:
            self._pending.discard(key)

    def shutdown(self):
        self._executor.shutdown(wait=True, cancel_futures=True)


_scheduler = MemoryScheduler()
schedule_memory_update = _scheduler.schedule