from pydantic import BaseModel, Field
from datetime import datetime
from uuid import uuid4
from typing import Any


# =============================================
# BASE MODEL
# =============================================

class MemoryBase(BaseModel):
    id: str = Field(default_factory=lambda: str(uuid4()))
    user_id: str
    org_id: str
    content: str
    content_type: str = "text"  # text / structured / embedding
    topic_tags: list[str] = []
    metadata: dict[str, Any] = {}
    created_at: datetime = Field(default_factory=datetime.utcnow)


# =============================================
# WORKING MEMORY (Redis, ephemeral)
# =============================================

class WorkingMemory(MemoryBase):
    session_id: str
    ttl_seconds: int = 1800  # 30 min default
    importance: float = 0.5  # 0-1, determines TTL
    tier: str = "working"

    def calculate_ttl(self) -> int:
        """TTL from importance: 900s (15min) to 3600s (60min)."""
        return int(900 + (self.importance * 2700))


# =============================================
# EPISODIC MEMORY (Memgraph + PostgreSQL)
# =============================================

class EpisodicMemory(MemoryBase):
    session_id: str
    encrypted_content: bytes | None = None
    embedding: list[float] | None = None
    pii_tags: list[str] = []
    classification: str = "general"  # personal / financial / medical / general
    confidence_score: float = 1.0
    relevance_score: float = 1.0
    access_count: int = 0
    last_accessed: datetime | None = None
    decay_rate: float = 0.1   # λ in Ebbinghaus formula
    stability: float = 1.0    # S in Ebbinghaus formula
    is_canary: bool = False
    is_decayed: bool = False
    source_type: str = "chat"  # chat / api / system / promotion
    tier: str = "episodic"
    memgraph_node_id: str | None = None


# =============================================
# SEMANTIC MEMORY (Memgraph + PostgreSQL)
# =============================================

class SemanticMemory(MemoryBase):
    embedding: list[float] | None = None
    pattern_description: str = ""
    confidence_score: float = 1.0
    source_episodic_count: int = 0
    source_episodic_ids: list[str] = []
    pii_verified_clean: bool = False
    tier: str = "semantic"
    memgraph_node_id: str | None = None


# =============================================
# QUERY RESULTS
# =============================================

class MemoryQueryResult(BaseModel):
    working: list[WorkingMemory] = []
    episodic: list[EpisodicMemory] = []
    semantic: list[SemanticMemory] = []
    arbitration: dict[str, Any] | None = None
    receipt_id: str | None = None
    primed: bool = False
    query_duration_ms: float = 0.0


# =============================================
# API REQUEST / RESPONSE MODELS
# =============================================

class MemoryStoreRequest(BaseModel):
    user_id: str
    org_id: str
    session_id: str
    content: str
    content_type: str = "text"
    topic_tags: list[str] = []
    importance: float = 0.5
    tier: str = "working"  # working / episodic
    metadata: dict[str, Any] = {}


class MemoryStoreResponse(BaseModel):
    memory_id: str
    tier: str
    ttl_seconds: int | None = None
    encrypted: bool = False
    pii_detected: list[str] = []
    receipt_id: str | None = None


class MemoryQueryRequest(BaseModel):
    user_id: str
    session_id: str
    query_text: str
    task_context: str = "general"  # customer_service / strategic_analysis / compliance_audit
    max_results: int = 20
    tiers: list[str] = ["working", "episodic", "semantic"]


class WorkingMemoryListResponse(BaseModel):
    session_id: str
    memories: list[WorkingMemory]
    count: int
