from ...entities import (
    EngineMessage,
    MessageContentImage,
    MessageContentText,
    MessageRole,
    TextPartition,
)
from ...memory import MemoryStore as MemoryStoreBase
from ...memory import MessageMemory

from abc import abstractmethod
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime
from enum import StrEnum
from typing import TYPE_CHECKING, Any, Literal, TypeAlias
from uuid import UUID, uuid4

if TYPE_CHECKING:
    from ...model.nlp.sentence import SentenceTransformerModel

    from numpy.typing import NDArray
else:

    class NDArray:  # noqa: D101
        def __class_getitem__(cls, _: Any) -> Any:
            return Any

    SentenceTransformerModel: TypeAlias = Any

Order = Literal["asc", "desc"]


class VectorFunction(StrEnum):
    # cosine distance: 1 – cosine similarity (angle between vectors). smaller
    # means  more similar
    COSINE_DISTANCE = "cosine_distance"
    # negative inner product: returns –(x·y) so that ORDER BY finds top
    # dot-product hits
    INNER_PRODUCT = "inner_product"
    # L1 (Manhattan) distance: sum of absolute differences; often more robust
    # in high-D
    L1_DISTANCE = "l1_distance"
    # L2 (Euclidean) distance: straight-line distance; smaller means more
    # similar
    L2_DISTANCE = "l2_distance"
    # Hamming distance for binary vectors
    VECTOR_DIMS = "vector_dims"
    # Jaccard distance for binary vectors
    VECTOR_NORMS = "vector_norms"


class MemoryType(StrEnum):
    CODE = "code"
    FILE = "file"
    URL = "url"
    RAW = "raw"


@dataclass(frozen=True, kw_only=True, slots=True)
class Session:
    id: UUID
    agent_id: UUID
    participant_id: UUID
    messages: int
    created_at: datetime


@dataclass(frozen=True, kw_only=True, slots=True)
class Memory:
    id: UUID
    model_id: str | None
    type: MemoryType
    participant_id: UUID
    namespace: str
    identifier: str
    data: str | None = None
    partitions: int
    symbols: dict[str, Any] | None
    created_at: datetime
    title: str | None = None
    description: str | None = None


@dataclass(frozen=True, kw_only=True, slots=True)
class PermanentMemoryStore:
    namespace: str
    description: str | None = None


@dataclass(frozen=True, kw_only=True, slots=True)
class PermanentMessage:
    id: UUID
    agent_id: UUID
    model_id: str
    session_id: UUID | None
    author: MessageRole
    data: str
    partitions: int
    created_at: datetime


@dataclass(frozen=True, kw_only=True, slots=True)
class PermanentMessageScored(PermanentMessage):
    score: float


@dataclass(frozen=True, kw_only=True, slots=True)
class PermanentMessagePartition:
    agent_id: UUID
    session_id: UUID | None
    message_id: UUID
    partition: int
    data: str
    embedding: NDArray[Any]
    created_at: datetime


@dataclass(frozen=True, kw_only=True, slots=True)
class PermanentMemoryPartition:
    participant_id: UUID
    memory_id: UUID
    partition: int
    data: str
    embedding: NDArray[Any]
    created_at: datetime


@dataclass(frozen=True, kw_only=True, slots=True)
class Hyperedge:
    id: UUID
    relation: str
    surface_text: str
    embedding: NDArray[Any]
    symbols: dict[str, Any] | None = None
    created_at: datetime


@dataclass(frozen=True, kw_only=True, slots=True)
class HyperedgeMemory:
    hyperedge_id: UUID
    memory_id: UUID
    char_start: int | None = None
    char_end: int | None = None


@dataclass(frozen=True, kw_only=True, slots=True)
class Entity:
    id: UUID
    name: str
    type: str | None = None
    embedding: NDArray[Any]
    symbols: dict[str, Any] | None = None
    participant_id: UUID | None = None
    namespace: str | None = None
    created_at: datetime


@dataclass(frozen=True, kw_only=True, slots=True)
class HyperedgeEntity:
    hyperedge_id: UUID
    entity_id: UUID
    role_idx: int
    role_label: str | None = None


class PermanentMessageMemory(MessageMemory):
    _session_id: UUID | None = None
    _sentence_model: SentenceTransformerModel | None

    def __init__(
        self,
        sentence_model: SentenceTransformerModel | None,
        **kwargs: Any,
    ) -> None:
        self._sentence_model = sentence_model
        super().__init__(**kwargs)

    @property
    def has_session(self) -> bool:
        return bool(self._session_id)

    @property
    def session_id(self) -> UUID | None:
        return self._session_id

    async def reset(self) -> None:
        raise NotImplementedError()

    async def reset_session(
        self, agent_id: UUID, participant_id: UUID
    ) -> None:
        self._session_id = await self.create_session(
            agent_id=agent_id, participant_id=participant_id
        )

    async def continue_session(
        self,
        agent_id: UUID,
        participant_id: UUID,
        session_id: UUID,
    ) -> None:
        self._session_id = await self.continue_session_and_get_id(
            agent_id=agent_id,
            participant_id=participant_id,
            session_id=session_id,
        )

    async def append(self, agent_id: UUID, data: EngineMessage) -> None:
        del agent_id
        del data
        raise NotImplementedError()

    @abstractmethod
    async def create_session(
        self, agent_id: UUID, participant_id: UUID
    ) -> UUID:
        raise NotImplementedError()

    @abstractmethod
    async def continue_session_and_get_id(
        self,
        agent_id: UUID,
        participant_id: UUID,
        session_id: UUID,
    ) -> UUID:
        raise NotImplementedError()

    @abstractmethod
    async def append_with_partitions(
        self,
        engine_message: EngineMessage,
        partitions: list[TextPartition],
    ) -> None:
        raise NotImplementedError()

    @abstractmethod
    async def get_recent_messages(
        self,
        session_id: UUID,
        participant_id: UUID,
        limit: int | None = None,
    ) -> list[EngineMessage]:
        raise NotImplementedError()

    @abstractmethod
    async def search_messages(
        self,
        agent_id: UUID,
        function: VectorFunction,
        participant_id: UUID,
        search_partitions: list[TextPartition],
        search_user_messages: bool,
        session_id: UUID | None,
        exclude_session_id: UUID | None,
        limit: int | None = None,
    ) -> list[EngineMessage]:
        raise NotImplementedError()

    @staticmethod
    def _build_message_with_partitions(
        engine_message: EngineMessage,
        session_id: UUID | None,
        partitions: list[TextPartition],
        *,
        created_at: datetime,
        message_id: UUID | None = None,
    ) -> tuple[PermanentMessage, list[PermanentMessagePartition]]:
        if message_id is None:
            message_id = uuid4()
        content = engine_message.message.content
        data = (
            content.text
            if isinstance(content, MessageContentText)
            else (
                str(content.image_url)
                if isinstance(content, MessageContentImage)
                else str(content)
            )
        )

        message = PermanentMessage(
            id=message_id,
            agent_id=engine_message.agent_id,
            model_id=engine_message.model_id,
            session_id=session_id,
            author=engine_message.message.role,
            data=data,
            partitions=len(partitions),
            created_at=created_at,
        )
        message_partitions = [
            PermanentMessagePartition(
                agent_id=message.agent_id,
                session_id=message.session_id,
                message_id=message.id,
                partition=i + 1,
                data=p.data,
                embedding=p.embeddings,
                created_at=created_at,
            )
            for i, p in enumerate(partitions)
        ]
        return message, message_partitions

    @staticmethod
    def _build_session(
        agent_id: UUID,
        participant_id: UUID,
        *,
        created_at: datetime,
        session_id: UUID | None = None,
    ) -> Session:
        if session_id is None:
            session_id = uuid4()
        return Session(
            id=session_id,
            agent_id=agent_id,
            participant_id=participant_id,
            messages=0,
            created_at=created_at,
        )


class PermanentMemory(MemoryStoreBase[Memory]):
    _sentence_model: SentenceTransformerModel | None

    def __init__(
        self,
        sentence_model: SentenceTransformerModel | None,
        **kwargs: Any,
    ) -> None:
        self._sentence_model = sentence_model
        super().__init__(**kwargs)

    async def append(self, agent_id: UUID, data: Memory) -> None:
        del agent_id
        del data
        raise NotImplementedError()

    async def reset(self) -> None:
        raise NotImplementedError()

    @abstractmethod
    async def append_with_partitions(
        self,
        namespace: str,
        participant_id: UUID,
        memory_type: MemoryType,
        data: str,
        identifier: str,
        partitions: list[TextPartition],
        symbols: Mapping[str, Any] | None = None,
        model_id: str | None = None,
        title: str | None = None,
        description: str | None = None,
    ) -> None:
        raise NotImplementedError()

    @abstractmethod
    async def search_memories(
        self,
        search_partitions: list[TextPartition],
        participant_id: UUID,
        namespace: str,
        function: VectorFunction,
        limit: int | None = None,
    ) -> list[PermanentMemoryPartition]:
        raise NotImplementedError()

    @abstractmethod
    async def list_memories(
        self,
        *,
        participant_id: UUID,
        namespace: str,
    ) -> list[Memory]:
        raise NotImplementedError()

    @staticmethod
    def _build_memory_with_partitions(
        namespace: str,
        participant_id: UUID,
        memory_type: MemoryType,
        data: str,
        identifier: str,
        partitions: list[TextPartition],
        *,
        created_at: datetime,
        symbols: Mapping[str, Any] | None = None,
        model_id: str | None = None,
        memory_id: UUID | None = None,
        title: str | None = None,
        description: str | None = None,
    ) -> tuple[Memory, list[PermanentMemoryPartition]]:
        if memory_id is None:
            memory_id = uuid4()
        entry = Memory(
            id=memory_id,
            model_id=model_id,
            type=memory_type,
            participant_id=participant_id,
            namespace=namespace,
            identifier=identifier,
            data=data,
            partitions=len(partitions),
            symbols=dict(symbols) if symbols is not None else None,
            created_at=created_at,
            title=title,
            description=description,
        )
        partition_rows = [
            PermanentMemoryPartition(
                participant_id=participant_id,
                memory_id=memory_id,
                partition=i + 1,
                data=p.data,
                embedding=p.embeddings,
                created_at=created_at,
            )
            for i, p in enumerate(partitions)
        ]
        return entry, partition_rows


class RecordNotFoundException(Exception):
    def __init__(self) -> None:
        super(RecordNotFoundException, self).__init__("record_not_found")


class RecordNotSavedException(Exception):
    def __init__(self) -> None:
        super(RecordNotSavedException, self).__init__("record_not_saved")
