"""
Memory models for agent cognition and recall.

This module provides memory-related models that enable agents to store,
organize, and recall different types of information.
"""

from typing import Any, Dict, List, Literal, Optional
from pydantic import Field, field_validator
from datetime import datetime, timezone

from .base_resource import BaseResource


class MemoryBlock(BaseResource):
    """
    Represents a block of agent memory.

    Memory blocks are categorized into three types based on cognitive science:
    - Episodic: Specific events and experiences
    - Procedural: Skills, habits, and how-to knowledge
    - Executive: High-level goals, plans, and decision-making context
    """

    resource_type: Literal["MemoryBlock"] = Field(
        default="MemoryBlock", description="Resource type identifier"
    )

    memory_type: Literal["episodic", "procedural", "executive"] = Field(
        description="Type of memory this block represents",
        examples=["episodic", "procedural", "executive"],
    )

    content: str = Field(
        description="The actual memory content or information",
        examples=[
            "Patient John Doe expressed concern about chest pain during consultation",
            "To calculate BMI: divide weight in kg by height in meters squared",
            "Current goal: Complete patient triage within 15 minutes",
        ],
    )

    metadata: Dict[str, Any] = Field(
        default_factory=dict,
        description="Additional structured metadata about this memory",
        examples=[
            {
                "patient_id": "pat-001",
                "encounter_id": "enc-001",
                "confidence": 0.95,
                "tags": ["chest_pain", "cardiology"],
            }
        ],
    )

    related_memories: List[str] = Field(
        default_factory=list,
        description="IDs of other memory blocks related to this one",
        examples=[["mem-001", "mem-002"]],
    )

    importance_score: float = Field(
        default=0.5,
        ge=0.0,
        le=1.0,
        description="Importance score from 0.0 to 1.0 for memory prioritization",
        examples=[0.8, 0.3, 0.95],
    )

    access_count: int = Field(
        default=0,
        ge=0,
        description="Number of times this memory has been accessed",
        examples=[0, 5, 23],
    )

    # Additional compatibility fields
    vector_id: Optional[str] = Field(
        default=None,
        description="Vector embedding ID for semantic search",
        examples=["vec_abc123", None],
    )

    last_accessed: Optional[datetime] = Field(
        default=None,
        description="Timestamp of last access",
        examples=["2024-01-15T10:30:00Z", None],
    )

    @field_validator("content")
    @classmethod
    def validate_content(cls, v: str) -> str:
        """Ensure content is not empty."""
        if not v.strip():
            raise ValueError("Memory content cannot be empty")
        return v.strip()

    @field_validator("metadata")
    @classmethod
    def validate_metadata(cls, v: Dict[str, Any]) -> Dict[str, Any]:
        """Ensure metadata is a valid dictionary."""
        if not isinstance(v, dict):
            raise ValueError("Metadata must be a dictionary")
        return v

    @property
    def linked_memories(self) -> List[str]:
        """Alias for related_memories for backward compatibility."""
        return self.related_memories

    @linked_memories.setter
    def linked_memories(self, value: List[str]) -> None:
        """Setter for linked_memories."""
        self.related_memories = value

    def add_related_memory(self, memory_id: str) -> None:
        """
        Add a related memory ID if not already present.

        Args:
            memory_id: ID of the related memory block
        """
        if memory_id not in self.related_memories:
            self.related_memories.append(memory_id)
            self.update_timestamp()

    def remove_related_memory(self, memory_id: str) -> bool:
        """
        Remove a related memory ID.

        Args:
            memory_id: ID of the memory block to remove

        Returns:
            True if the memory was removed, False if it wasn't found
        """
        if memory_id in self.related_memories:
            self.related_memories.remove(memory_id)
            self.update_timestamp()
            return True
        return False

    def increment_access(self) -> None:
        """Increment the access count and update timestamp."""
        self.access_count += 1
        self.last_accessed = datetime.now(timezone.utc)
        self.update_timestamp()

    def set_importance(self, score: float) -> None:
        """
        Set the importance score.

        Args:
            score: Importance score between 0.0 and 1.0
        """
        if not 0.0 <= score <= 1.0:
            raise ValueError("Importance score must be between 0.0 and 1.0")
        self.importance_score = score
        self.update_timestamp()

    def add_metadata(self, key: str, value: Any) -> None:
        """
        Add or update a metadata field.

        Args:
            key: Metadata key
            value: Metadata value
        """
        self.metadata[key] = value
        self.update_timestamp()

    def get_metadata(self, key: str, default: Any = None) -> Any:
        """
        Get a metadata value.

        Args:
            key: Metadata key
            default: Default value if key not found

        Returns:
            Metadata value or default
        """
        return self.metadata.get(key, default)

    def __repr__(self) -> str:
        """Enhanced representation including memory type."""
        return f"MemoryBlock(id='{self.id}', type='{self.memory_type}', importance={self.importance_score})"
