"""Memory backend implementations."""
from __future__ import annotations

from collections import OrderedDict
import hashlib
import re

from logos_adk.memory import MemoryBackend, MemoryEntry
from logos_adk.memory.backends.postgresql import PostgreSQLBackend
from logos_adk.memory.backends.sqlite import SQLiteBackend


class InMemoryBackend(MemoryBackend):
    """Simple in-process backend for tests and local development."""

    def __init__(self) -> None:
        self._entries: OrderedDict[str, MemoryEntry] = OrderedDict()

    async def store(self, entry: MemoryEntry) -> None:
        self._entries[entry.id] = entry

    async def search(
        self,
        query: str,
        limit: int = 10,
        session_id: str | None = None,
    ) -> list[MemoryEntry]:
        normalized_query = query.strip().lower()
        keywords = [
            token
            for token in re.findall(r"[a-zA-Z0-9_]+", normalized_query)
            if len(token) >= 4
        ]
        matches: list[MemoryEntry] = []

        for entry in reversed(self._entries.values()):
            if session_id is not None and entry.session_id != session_id:
                continue
            content = entry.content.lower()
            if normalized_query and normalized_query not in content:
                if keywords and not any(keyword in content for keyword in keywords):
                    continue
                if not keywords:
                    continue
            matches.append(entry)
            if len(matches) >= limit:
                break

        return matches

    async def delete(self, entry_id: str) -> None:
        self._entries.pop(entry_id, None)


class ShardedMemoryBackend(MemoryBackend):
    """Routes memory operations across multiple backend shards."""

    def __init__(self, shards: list[MemoryBackend]) -> None:
        if not shards:
            raise ValueError("ShardedMemoryBackend requires at least one shard")
        self._shards = shards

    def _shard_index(self, key: str) -> int:
        digest = hashlib.sha256(key.encode("utf-8")).hexdigest()
        return int(digest, 16) % len(self._shards)

    def _select_store_shard(self, entry: MemoryEntry) -> MemoryBackend:
        shard_key = entry.session_id or entry.id
        return self._shards[self._shard_index(shard_key)]

    def _search_shards(self, session_id: str | None) -> list[MemoryBackend]:
        if session_id is not None:
            return [self._shards[self._shard_index(session_id)]]
        return self._shards

    async def store(self, entry: MemoryEntry) -> None:
        await self._select_store_shard(entry).store(entry)

    async def search(
        self,
        query: str,
        limit: int = 10,
        session_id: str | None = None,
    ) -> list[MemoryEntry]:
        matches: list[MemoryEntry] = []
        for shard in self._search_shards(session_id):
            matches.extend(await shard.search(query, limit=limit, session_id=session_id))
        matches.sort(key=lambda entry: entry.created_at, reverse=True)
        return matches[:limit]

    async def delete(self, entry_id: str) -> None:
        for shard in self._shards:
            await shard.delete(entry_id)


__all__ = [
    "InMemoryBackend",
    "ShardedMemoryBackend",
    "PostgreSQLBackend",
    "SQLiteBackend",
]
