"""MemoryManager — tiered memory backend routing."""
from __future__ import annotations

import logging
import threading
from enum import Enum

from langchain_core.chat_history import BaseChatMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory

logger = logging.getLogger(__name__)


class MemoryType(str, Enum):
    WITHIN_CLAIM = "within_claim"
    CROSS_CLAIM = "cross_claim"
    MULTI_TURN = "multi_turn"


class MemoryManager:
    """Thread-safe memory manager with tiered backends."""

    def __init__(self, redis_url: str | None = None, postgres_url: str | None = None) -> None:
        self._in_memory_store: dict[str, ChatMessageHistory] = {}
        self._lock = threading.Lock()
        self._redis_url = redis_url
        self._postgres_url = postgres_url

    def get_history(self, session_id: str, memory_type: MemoryType = MemoryType.WITHIN_CLAIM) -> BaseChatMessageHistory:
        match memory_type:
            case MemoryType.WITHIN_CLAIM:
                return self._get_in_memory(session_id)
            case MemoryType.CROSS_CLAIM:
                return self._get_redis(session_id)
            case MemoryType.MULTI_TURN:
                return self._get_postgres(session_id)

    def clear_claim_memory(self, claim_id: str) -> None:
        with self._lock:
            self._in_memory_store.pop(claim_id, None)

    def _get_in_memory(self, session_id: str) -> ChatMessageHistory:
        with self._lock:
            if session_id not in self._in_memory_store:
                self._in_memory_store[session_id] = ChatMessageHistory()
            return self._in_memory_store[session_id]

    def _get_redis(self, session_id: str) -> BaseChatMessageHistory:
        if not self._redis_url:
            logger.warning("Redis URL not configured, falling back to in-memory")
            return self._get_in_memory(f"redis-fallback-{session_id}")
        try:
            from langchain_community.chat_message_histories import RedisChatMessageHistory
            return RedisChatMessageHistory(session_id=session_id, url=self._redis_url, ttl=86400 * 7)
        except ImportError:
            logger.warning("redis not installed, falling back to in-memory")
            return self._get_in_memory(f"redis-fallback-{session_id}")

    def _get_postgres(self, session_id: str) -> BaseChatMessageHistory:
        if not self._postgres_url:
            logger.warning("Postgres URL not configured, falling back to in-memory")
            return self._get_in_memory(f"pg-fallback-{session_id}")
        try:
            from langchain_community.chat_message_histories import PostgresChatMessageHistory
            return PostgresChatMessageHistory(session_id=session_id, connection_string=self._postgres_url)
        except ImportError:
            logger.warning("psycopg not installed, falling back to in-memory")
            return self._get_in_memory(f"pg-fallback-{session_id}")
