"""Memory provenance and lineage tracking models."""

from datetime import datetime, timezone
from enum import Enum
from typing import Any, Optional
from uuid import uuid4

from pydantic import BaseModel, Field


class MemorySource(str, Enum):
    """Source of memory creation."""

    CONVERSATION = "conversation"  # Extracted from user conversation
    API_DIRECT = "api_direct"  # Created directly via API
    INFERENCE = "inference"  # Inferred from other memories
    MERGE = "merge"  # Merged from multiple memories
    IMPORT = "import"  # Imported from external system
    SYSTEM_GENERATED = "system_generated"  # Generated by system (e.g., summaries)
    USER_INPUT = "user_input"  # Direct user input
    REFINEMENT = "refinement"  # Refined/improved version of existing memory


class TransformationType(str, Enum):
    """Type of transformation applied to memory."""

    CREATED = "created"  # Initial creation
    MERGED = "merged"  # Merged with other memories
    SPLIT = "split"  # Split into multiple memories
    REFINED = "refined"  # Text quality improvement
    ENRICHED = "enriched"  # Added metadata/entities
    SUMMARIZED = "summarized"  # Condensed from longer content
    INFERRED = "inferred"  # Derived from other information
    CONFLICT_RESOLVED = "conflict_resolved"  # Result of conflict resolution
    AUTO_UPDATED = "auto_updated"  # Automatically updated by system


class QualityMetrics(BaseModel):
    """Quality assessment metrics for a memory."""

    specificity: float = Field(default=0.5, ge=0.0, le=1.0)  # How specific vs vague
    verifiability: float = Field(default=0.5, ge=0.0, le=1.0)  # Can it be verified
    completeness: float = Field(default=0.5, ge=0.0, le=1.0)  # Has all necessary info
    clarity: float = Field(default=0.5, ge=0.0, le=1.0)  # Clear and unambiguous
    relevance: float = Field(default=0.5, ge=0.0, le=1.0)  # Relevant to user
    overall_score: float = Field(default=0.5, ge=0.0, le=1.0)  # Composite score

    def calculate_overall(self) -> float:
        """Calculate overall quality score."""
        self.overall_score = (
            self.specificity * 0.25
            + self.verifiability * 0.20
            + self.completeness * 0.25
            + self.clarity * 0.15
            + self.relevance * 0.15
        )
        return self.overall_score


class Citation(BaseModel):
    """Citation linking memory to original source."""

    citation_id: str = Field(default_factory=lambda: str(uuid4()))
    source_type: str  # "message", "document", "url", "memory", "conversation"
    source_id: Optional[str] = None  # ID of source (message_id, doc_id, etc.)
    source_url: Optional[str] = None  # URL if applicable
    source_text: Optional[str] = None  # Excerpt from source
    confidence: float = Field(default=1.0, ge=0.0, le=1.0)  # Confidence in citation
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    metadata: dict[str, Any] = Field(default_factory=dict)


class MemoryTransformation(BaseModel):
    """Record of a transformation applied to memory."""

    transformation_id: str = Field(default_factory=lambda: str(uuid4()))
    transformation_type: TransformationType
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    parent_memory_ids: list[str] = Field(default_factory=list)  # Memories it derived from
    description: str = ""  # Human-readable description
    metadata: dict[str, Any] = Field(default_factory=dict)  # Additional context


class MemoryLineage(BaseModel):
    """Complete lineage information for a memory."""

    memory_id: str
    source: MemorySource = MemorySource.CONVERSATION
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    created_by: Optional[str] = None  # User/agent/system that created it

    # Ancestry
    parent_memory_ids: list[str] = Field(default_factory=list)  # Direct parents
    ancestor_memory_ids: list[str] = Field(default_factory=list)  # All ancestors
    child_memory_ids: list[str] = Field(default_factory=list)  # Derived memories

    # Transformations
    transformations: list[MemoryTransformation] = Field(default_factory=list)

    # Citations
    citations: list[Citation] = Field(default_factory=list)

    # Quality
    quality_metrics: Optional[QualityMetrics] = None

    # Access tracking (for provenance)
    access_count: int = 0
    last_accessed_at: Optional[datetime] = None
    accessed_by: list[str] = Field(default_factory=list)  # List of user/agent IDs

    def add_transformation(
        self,
        transformation_type: TransformationType,
        parent_ids: Optional[list[str]] = None,
        description: str = "",
        metadata: Optional[dict[str, Any]] = None,
    ) -> None:
        """Add a transformation record."""
        transformation = MemoryTransformation(
            transformation_type=transformation_type,
            parent_memory_ids=parent_ids or [],
            description=description,
            metadata=metadata or {},
        )
        self.transformations.append(transformation)

        # Update parent_memory_ids if this is a derived memory
        if parent_ids:
            self.parent_memory_ids = list(set(self.parent_memory_ids + parent_ids))

    def add_citation(
        self,
        source_type: str,
        source_id: Optional[str] = None,
        source_url: Optional[str] = None,
        source_text: Optional[str] = None,
        confidence: float = 1.0,
        metadata: Optional[dict[str, Any]] = None,
    ) -> None:
        """Add a citation."""
        citation = Citation(
            source_type=source_type,
            source_id=source_id,
            source_url=source_url,
            source_text=source_text,
            confidence=confidence,
            metadata=metadata or {},
        )
        self.citations.append(citation)

    def record_access(self, accessor_id: str) -> None:
        """Record an access to this memory."""
        self.access_count += 1
        self.last_accessed_at = datetime.now(timezone.utc)
        if accessor_id not in self.accessed_by:
            self.accessed_by.append(accessor_id)

    def get_generation(self) -> int:
        """Get generation level (0 = original, 1 = first derivative, etc.)."""
        if not self.parent_memory_ids:
            return 0
        # This is simplified - in reality you'd traverse the full tree
        return len(self.transformations)

    def is_derived(self) -> bool:
        """Check if this memory is derived from others."""
        return len(self.parent_memory_ids) > 0

    def has_citations(self) -> bool:
        """Check if memory has citations."""
        return len(self.citations) > 0


class ProvenanceChain(BaseModel):
    """Complete provenance chain for a memory."""

    memory_id: str
    chain: list[MemoryLineage] = Field(default_factory=list)  # Ordered from root to current
    total_generations: int = 0
    root_memory_ids: list[str] = Field(default_factory=list)  # Original source memories

    def add_link(self, lineage: MemoryLineage) -> None:
        """Add a lineage link to the chain."""
        self.chain.append(lineage)
        self.total_generations = len(self.chain)

    def get_root_sources(self) -> list[MemorySource]:
        """Get the original sources of this memory."""
        if not self.chain:
            return []
        return [link.source for link in self.chain if not link.parent_memory_ids]

    def get_all_transformations(self) -> list[MemoryTransformation]:
        """Get all transformations in the chain."""
        transformations = []
        for link in self.chain:
            transformations.extend(link.transformations)
        return transformations

    def get_all_citations(self) -> list[Citation]:
        """Get all citations in the chain."""
        citations = []
        for link in self.chain:
            citations.extend(link.citations)
        return citations


# Extended Memory metadata schema for provenance
PROVENANCE_METADATA_SCHEMA = {
    # Source tracking
    "source": str,  # MemorySource enum value
    "source_id": Optional[str],  # Original source identifier
    "source_metadata": Optional[dict[str, Any]],  # Additional source context
    # Lineage
    "parent_memory_ids": list[str],  # Direct parents
    "derived_from": list[str],  # Alias for parent_memory_ids
    "supersedes": list[str],  # Memories this replaces
    "related_memories": list[str],  # Related but not parent/child
    # Transformation history
    "transformations": list[dict[str, Any]],  # List of transformations
    "last_transformation": Optional[dict[str, Any]],  # Most recent transformation
    # Citations
    "citations": list[dict[str, Any]],  # List of citations
    "primary_citation": Optional[dict[str, Any]],  # Main citation
    # Quality
    "quality_metrics": Optional[dict[str, Any]],  # Quality scores
    "quality_assessed_at": Optional[str],  # When quality was last assessed
    # Provenance flags
    "is_derived": bool,  # Whether memory is derived
    "is_inferred": bool,  # Whether memory is inferred
    "is_verified": bool,  # Whether memory has been verified
    "verification_source": Optional[str],  # How it was verified
}
