from __future__ import annotations

"""Vector-backed memory backend."""

import math
from dataclasses import dataclass, field
from typing import Protocol
from uuid import uuid4

from agentrspy.core import Memory, Message

from .in_memory import InMemoryMemory


@dataclass(slots=True)
class VectorSearchResult:
    """Search result returned by a vector store."""

    score: float
    payload: Message


class Embedder(Protocol):
    """Computes embeddings for messages."""

    async def embed(self, text: str) -> list[float]:
        """Generate an embedding vector."""


class VectorStore(Protocol):
    """Persists and searches embedding vectors."""

    async def upsert(self, id: str, vector: list[float], payload: Message) -> None:
        """Upsert a vector and payload."""

    async def search(self, query: list[float], limit: int) -> list[VectorSearchResult]:
        """Search the store."""

    async def clear(self) -> None:
        """Clear all stored vectors."""


class SimpleEmbedder:
    """Small deterministic embedder useful for tests and local demos."""

    async def embed(self, text: str) -> list[float]:
        buckets = [0.0] * 16
        for index, byte in enumerate(text.encode("utf-8")):
            buckets[index % 16] += byte / 255.0
        return buckets


@dataclass(slots=True)
class InMemoryVectorStore:
    """In-memory vector store with cosine similarity search."""

    items: list[tuple[str, list[float], Message]] = field(default_factory=list)

    @classmethod
    def new(cls) -> "InMemoryVectorStore":
        return cls()

    async def upsert(self, id: str, vector: list[float], payload: Message) -> None:
        for index, item in enumerate(self.items):
            if item[0] == id:
                self.items[index] = (id, vector, payload)
                return
        self.items.append((id, vector, payload))

    async def search(self, query: list[float], limit: int) -> list[VectorSearchResult]:
        scored = [
            VectorSearchResult(score=_cosine_similarity(vector, query), payload=payload)
            for _, vector, payload in self.items
        ]
        scored.sort(key=lambda item: item.score, reverse=True)
        return scored[:limit]

    async def clear(self) -> None:
        self.items.clear()


@dataclass(slots=True)
class VectorMemory(Memory):
    """Memory backend that combines recent history with semantic retrieval."""

    embedder: Embedder = field(default_factory=SimpleEmbedder)
    store_backend: VectorStore = field(default_factory=InMemoryVectorStore)
    recent: InMemoryMemory = field(default_factory=InMemoryMemory.new)

    @classmethod
    def new(cls) -> "VectorMemory":
        return cls()

    @classmethod
    def with_components(cls, embedder: Embedder, store: VectorStore) -> "VectorMemory":
        return cls(embedder=embedder, store_backend=store)

    async def store(self, key: str, value: Message) -> None:
        vector = await self.embedder.embed(value.text_content())
        await self.store_backend.upsert(f"{key}-{uuid4()}", vector, value)
        await self.recent.store(key, value)

    async def retrieve(self, query: str, limit: int) -> list[Message]:
        vector = await self.embedder.embed(query)
        return [
            result.payload for result in await self.store_backend.search(vector, limit)
        ]

    async def history(self) -> list[Message]:
        return await self.recent.history()

    async def clear(self) -> None:
        await self.store_backend.clear()
        await self.recent.clear()

    async def token_count(self) -> int:
        history = await self.recent.history()
        return sum(len(message.text_content()) // 4 for message in history)


def _cosine_similarity(left: list[float], right: list[float]) -> float:
    if len(left) != len(right) or not left:
        return 0.0
    dot = sum(a * b for a, b in zip(left, right))
    left_norm = math.sqrt(sum(value * value for value in left))
    right_norm = math.sqrt(sum(value * value for value in right))
    if left_norm == 0.0 or right_norm == 0.0:
        return 0.0
    return dot / (left_norm * right_norm)
