from enum import Enum
from typing import List, Optional
from pydantic import BaseModel, Field, model_validator

class Role(str, Enum):
    USER = "user"
    ASSISTANT = "model"
    SYSTEM = "system"

class MemoryType(str, Enum):
    # Add actual memory types based on your Go definitions
    CORE = "core"
    GENERAL = "general" 

class Message(BaseModel):
    role: Role
    content: str

class UserResponse(BaseModel):
    user_id: str = Field(alias="userId")

class InsertMemoryRequest(BaseModel):
    user_id: str = Field(alias="userId")
    messages: List[Message]

class MemoryInsertionResponse(BaseModel):
    req_id: str = Field(alias="ReqId")
    msg: str = Field(alias="Msg")

class MemoryRetrievalRequest(BaseModel):
    user_id: str = Field(alias="userId")
    messages: Optional[List[Message]] = None
    query: Optional[str] = None
    threshold: Optional[float] = None

    @model_validator(mode='after')
    def check_query_or_messages(self):
        # Enforces XOR logic: exactly one must be provided, but not both.
        has_messages = bool(self.messages)
        has_query = bool(self.query)
        
        if has_messages == has_query:
            raise ValueError("Exactly one of 'messages' or 'query' must be provided. Cannot provide both or neither.")
        return self

class Memory(BaseModel):
    memory_text: str = Field(alias="Memory_text")
    type: str = Field(alias="Type")  # Or MemoryType if enum is strictly enforced
    memory_id: str = Field(alias="Memory_Id")
    user_id: str = Field(alias="UserId")

class FullReqStatus(BaseModel):
    status: str  # Assuming types.ReqStatus resolves to a string. Use Enum if strict.
    error: Optional[str] = None
    created_at: str = Field(alias="createdAt")