from __future__ import annotations

from typing import Any, Optional

from pydantic import BaseModel, ConfigDict, Field, model_validator

from goodmem.models.memory import Memory
from goodmem.models.chunk_reference import ChunkReference


class RetrievedItem(BaseModel):
    """A retrieved result that can be either a Memory or MemoryChunk"""

    model_config = ConfigDict(populate_by_name=True, protected_namespaces=())

    memory: Optional[Memory] = Field(default=None, description='Complete memory object (if retrieved)')
    chunk: Optional[ChunkReference] = Field(default=None, description='Reference to a memory chunk (if retrieved)')

    @model_validator(mode='after')
    def _check_exactly_one_variant(self) -> 'RetrievedItem':
        _variant_fields = ['memory', 'chunk']
        _set = [f for f in _variant_fields if getattr(self, f) is not None]
        if len(_set) != 1:
            raise ValueError(
                f"Exactly one of memory, chunk must be set, "
                f"got {len(_set)}: {_set or 'none'}"
            )
        return self

