"""
Engram Core Types

Neural memory format for AI systems — hierarchical, temporal, multi-modal.
"""

from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Optional
import numpy as np
from numpy.typing import NDArray


# ============== ENUMS ==============

class Modality(str, Enum):
    TEXT = "text"
    IMAGE = "image"
    AUDIO = "audio"
    CODE = "code"


class DecayTier(str, Enum):
    HOT = "hot"
    WARM = "warm"
    COLD = "cold"
    ARCHIVE = "archive"


class QualitySource(str, Enum):
    DIRECT = "direct"
    INFERRED = "inferred"
    SUMMARIZED = "summarized"


class ContentType(str, Enum):
    TEXT = "text"
    IMAGE = "image"
    AUDIO = "audio"
    CODE = "code"
    SUMMARY = "summary"


class EntityType(str, Enum):
    PERSON = "person"
    PLACE = "place"
    ORGANIZATION = "organization"
    CONCEPT = "concept"
    EVENT = "event"
    DOCUMENT = "document"


class LinkType(str, Enum):
    RELATED = "related"
    REFERENCES = "references"
    CONTRADICTS = "contradicts"
    SUPERSEDES = "supersedes"
    ELABORATES = "elaborates"
    SUMMARIZES = "summarizes"
    CAUSES = "causes"
    FOLLOWS = "follows"


class LinkCreator(str, Enum):
    USER = "user"
    AGENT = "agent"
    SYSTEM = "system"


# ============== CONTENT ==============

@dataclass
class ExternalRef:
    """Reference to external file or URL."""
    type: str  # "file" | "url"
    path: str
    hash: Optional[str] = None


@dataclass
class NodeContent:
    """Content payload for a memory node."""
    type: ContentType
    data: str | bytes
    mime_type: Optional[str] = None
    language: Optional[str] = None
    tokens: Optional[int] = None
    original_length: Optional[int] = None
    original_hash: Optional[str] = None
    ref: Optional[ExternalRef] = None


# ============== TEMPORAL ==============

@dataclass
class TemporalInfo:
    """Temporal metadata for a memory node."""
    created: int  # Unix timestamp ms
    modified: int
    accessed: int
    decay_tier: DecayTier = DecayTier.HOT
    expires: Optional[int] = None


# ============== QUALITY ==============

@dataclass
class QualityInfo:
    """Quality metrics for a memory node."""
    score: float
    confidence: float
    source: QualitySource = QualitySource.DIRECT
    verified: Optional[bool] = None


# ============== METADATA ==============

@dataclass
class NodeMetadata:
    """Metadata for a memory node."""
    source_file: Optional[str] = None
    source_line: Optional[int] = None
    tags: list[str] = field(default_factory=list)
    custom: dict[str, Any] = field(default_factory=dict)


# ============== MEMORY NODE ==============

@dataclass
class MemoryNode:
    """A single memory node in the engram."""
    id: str
    
    # Hierarchy
    parent_id: Optional[str] = None
    children: list[str] = field(default_factory=list)
    depth: int = 0
    path: str = "/"
    
    # Content
    content: Optional[NodeContent] = None
    
    # Embedding
    embedding: Optional[NDArray[np.float32]] = None
    embedding_model: Optional[str] = None
    
    # Temporal
    temporal: Optional[TemporalInfo] = None
    
    # Quality
    quality: Optional[QualityInfo] = None
    
    # Metadata
    metadata: Optional[NodeMetadata] = None


# ============== ENTITIES ==============

@dataclass
class EntityMention:
    """A mention of an entity in a node."""
    node_id: str
    span: tuple[int, int]
    confidence: float


@dataclass
class EntityRelationship:
    """A relationship between entities."""
    target_id: str
    type: str
    properties: dict[str, Any] = field(default_factory=dict)


@dataclass
class Entity:
    """A named entity extracted from memories."""
    id: str
    type: EntityType
    name: str
    aliases: list[str] = field(default_factory=list)
    properties: dict[str, Any] = field(default_factory=dict)
    mentions: list[EntityMention] = field(default_factory=list)
    relationships: list[EntityRelationship] = field(default_factory=list)


# ============== LINKS ==============

@dataclass
class MemoryLink:
    """A link between memory nodes."""
    id: str
    source_id: str
    target_id: str
    type: LinkType
    confidence: float
    bidirectional: bool = False
    created: int = 0
    created_by: LinkCreator = LinkCreator.SYSTEM
    metadata: dict[str, Any] = field(default_factory=dict)


# ============== SECURITY ==============

@dataclass
class SecurityConfig:
    """Security configuration for an engram file."""
    encrypted: bool = False
    algorithm: str = "none"  # "aes-256-gcm" | "none"
    kdf: str = "none"  # "argon2id" | "pbkdf2" | "none"
    salt: Optional[bytes] = None
    nonce: Optional[bytes] = None
    integrity: bytes = b""
    signature: Optional[bytes] = None


# ============== FILE METADATA ==============

@dataclass
class FileMetadata:
    """Metadata for the engram file."""
    source: str
    description: Optional[str] = None
    tags: list[str] = field(default_factory=list)
    author: Optional[str] = None
    persona: Optional[str] = None


@dataclass
class SchemaConfig:
    """Schema configuration for the engram file."""
    embedding_model: str = "none"
    embedding_dims: int = 0
    chunk_strategy: str = "paragraph"  # "paragraph" | "sentence" | "fixed" | "semantic"
    modalities: list[str] = field(default_factory=lambda: ["text"])


@dataclass
class FileStats:
    """Statistics for the engram file."""
    total_chunks: int = 0
    total_tokens: int = 0
    root_nodes: int = 0
    max_depth: int = 0
    entity_count: int = 0
    link_count: int = 0


# ============== HEADER ==============

@dataclass
class EngramHeader:
    """Header for an engram file."""
    version: tuple[int, int] = (1, 0)
    created: int = 0
    modified: int = 0
    security: SecurityConfig = field(default_factory=SecurityConfig)
    metadata: FileMetadata = field(default_factory=lambda: FileMetadata(source="unknown"))
    schema: SchemaConfig = field(default_factory=SchemaConfig)
    stats: FileStats = field(default_factory=FileStats)


# ============== FILE STRUCTURE ==============

@dataclass
class EngramFile:
    """Complete engram file structure."""
    header: EngramHeader
    nodes: list[MemoryNode] = field(default_factory=list)
    entities: list[Entity] = field(default_factory=list)
    links: list[MemoryLink] = field(default_factory=list)
