"""
Engram File I/O

Reading and writing .engram binary files.
"""

import hashlib
import struct
from pathlib import Path
from typing import Any, Optional, Union

import msgpack
import numpy as np

from .types import (
    ContentType,
    DecayTier,
    EngramFile,
    EngramHeader,
    Entity,
    EntityMention,
    EntityRelationship,
    EntityType,
    ExternalRef,
    FileMetadata,
    FileStats,
    LinkCreator,
    LinkType,
    MemoryLink,
    MemoryNode,
    NodeContent,
    NodeMetadata,
    QualityInfo,
    QualitySource,
    SchemaConfig,
    SecurityConfig,
    TemporalInfo,
)

# ============== CONSTANTS ==============

MAGIC = b"ENGRAM"
VERSION_MAJOR = 1
VERSION_MINOR = 0
HEADER_OFFSET = 12  # magic(6) + version(2) + header_len(4)


# ============== UTILITIES ==============

def ensure_engram_extension(filename: str) -> str:
    """Ensure filename has .engram extension."""
    if not filename.lower().endswith(".engram"):
        return filename + ".engram"
    return filename


def _hash_payload(payload: bytes) -> bytes:
    """Compute SHA-256 hash of payload."""
    return hashlib.sha256(payload).digest()


# ============== SERIALIZATION ==============

def _serialize_node(node: MemoryNode) -> dict[str, Any]:
    """Convert MemoryNode to dict for msgpack serialization."""
    result: dict[str, Any] = {
        "id": node.id,
        "parentId": node.parent_id,
        "children": node.children,
        "depth": node.depth,
        "path": node.path,
    }
    
    if node.content:
        content_dict: dict[str, Any] = {
            "type": node.content.type.value if isinstance(node.content.type, ContentType) else node.content.type,
            "data": node.content.data,
        }
        if node.content.mime_type:
            content_dict["mimeType"] = node.content.mime_type
        if node.content.language:
            content_dict["language"] = node.content.language
        if node.content.tokens:
            content_dict["tokens"] = node.content.tokens
        result["content"] = content_dict
    
    if node.embedding is not None:
        # Store as list of floats for msgpack compatibility
        result["embedding"] = node.embedding.tolist()
        if node.embedding_model:
            result["embeddingModel"] = node.embedding_model
    
    if node.temporal:
        result["temporal"] = {
            "created": node.temporal.created,
            "modified": node.temporal.modified,
            "accessed": node.temporal.accessed,
            "decayTier": node.temporal.decay_tier.value if isinstance(node.temporal.decay_tier, DecayTier) else node.temporal.decay_tier,
        }
        if node.temporal.expires:
            result["temporal"]["expires"] = node.temporal.expires
    
    if node.quality:
        result["quality"] = {
            "score": node.quality.score,
            "confidence": node.quality.confidence,
            "source": node.quality.source.value if isinstance(node.quality.source, QualitySource) else node.quality.source,
        }
        if node.quality.verified is not None:
            result["quality"]["verified"] = node.quality.verified
    
    if node.metadata:
        meta: dict[str, Any] = {}
        if node.metadata.source_file:
            meta["sourceFile"] = node.metadata.source_file
        if node.metadata.source_line:
            meta["sourceLine"] = node.metadata.source_line
        if node.metadata.tags:
            meta["tags"] = node.metadata.tags
        if node.metadata.custom:
            meta["custom"] = node.metadata.custom
        if meta:
            result["metadata"] = meta
    
    return result


def _deserialize_node(data: dict[str, Any]) -> MemoryNode:
    """Convert dict to MemoryNode."""
    content = None
    if "content" in data and data["content"]:
        c = data["content"]
        content = NodeContent(
            type=ContentType(c.get("type", "text")),
            data=c.get("data", ""),
            mime_type=c.get("mimeType"),
            language=c.get("language"),
            tokens=c.get("tokens"),
        )
    
    embedding = None
    if "embedding" in data and data["embedding"] is not None:
        emb = data["embedding"]
        if isinstance(emb, list):
            embedding = np.array(emb, dtype=np.float32)
        elif isinstance(emb, (bytes, bytearray)):
            # Handle raw bytes (4 bytes per float32)
            embedding = np.frombuffer(emb, dtype=np.float32)
    
    temporal = None
    if "temporal" in data and data["temporal"]:
        t = data["temporal"]
        temporal = TemporalInfo(
            created=t.get("created", 0),
            modified=t.get("modified", 0),
            accessed=t.get("accessed", 0),
            decay_tier=DecayTier(t.get("decayTier", "hot")),
            expires=t.get("expires"),
        )
    
    quality = None
    if "quality" in data and data["quality"]:
        q = data["quality"]
        quality = QualityInfo(
            score=q.get("score", 0.0),
            confidence=q.get("confidence", 0.0),
            source=QualitySource(q.get("source", "direct")),
            verified=q.get("verified"),
        )
    
    metadata = None
    if "metadata" in data and data["metadata"]:
        m = data["metadata"]
        metadata = NodeMetadata(
            source_file=m.get("sourceFile"),
            source_line=m.get("sourceLine"),
            tags=m.get("tags", []),
            custom=m.get("custom", {}),
        )
    
    return MemoryNode(
        id=data.get("id", ""),
        parent_id=data.get("parentId"),
        children=data.get("children", []),
        depth=data.get("depth", 0),
        path=data.get("path", "/"),
        content=content,
        embedding=embedding,
        embedding_model=data.get("embeddingModel"),
        temporal=temporal,
        quality=quality,
        metadata=metadata,
    )


def _serialize_header(header: EngramHeader) -> dict[str, Any]:
    """Convert EngramHeader to dict for msgpack serialization."""
    return {
        "version": list(header.version),
        "created": header.created,
        "modified": header.modified,
        "security": {
            "encrypted": header.security.encrypted,
            "algorithm": header.security.algorithm,
            "kdf": header.security.kdf,
            "integrity": header.security.integrity,
        },
        "metadata": {
            "source": header.metadata.source,
            "description": header.metadata.description,
            "tags": header.metadata.tags,
            "author": header.metadata.author,
        },
        "schema": {
            "embeddingModel": header.schema.embedding_model,
            "embeddingDims": header.schema.embedding_dims,
            "chunkStrategy": header.schema.chunk_strategy,
            "modalities": header.schema.modalities,
        },
        "stats": {
            "totalChunks": header.stats.total_chunks,
            "totalTokens": header.stats.total_tokens,
            "rootNodes": header.stats.root_nodes,
            "maxDepth": header.stats.max_depth,
            "entityCount": header.stats.entity_count,
            "linkCount": header.stats.link_count,
        },
    }


def _deserialize_header(data: dict[str, Any]) -> EngramHeader:
    """Convert dict to EngramHeader."""
    security = SecurityConfig()
    if "security" in data:
        s = data["security"]
        integrity = s.get("integrity", b"")
        if isinstance(integrity, (list, tuple)):
            integrity = bytes(integrity)
        security = SecurityConfig(
            encrypted=s.get("encrypted", False),
            algorithm=s.get("algorithm", "none"),
            kdf=s.get("kdf", "none"),
            integrity=integrity,
        )
    
    metadata = FileMetadata(source="unknown")
    if "metadata" in data:
        m = data["metadata"]
        metadata = FileMetadata(
            source=m.get("source", "unknown"),
            description=m.get("description"),
            tags=m.get("tags", []),
            author=m.get("author"),
        )
    
    schema = SchemaConfig()
    if "schema" in data:
        sc = data["schema"]
        schema = SchemaConfig(
            embedding_model=sc.get("embeddingModel", "none"),
            embedding_dims=sc.get("embeddingDims", 0),
            chunk_strategy=sc.get("chunkStrategy", "paragraph"),
            modalities=sc.get("modalities", ["text"]),
        )
    
    stats = FileStats()
    if "stats" in data:
        st = data["stats"]
        stats = FileStats(
            total_chunks=st.get("totalChunks", 0),
            total_tokens=st.get("totalTokens", 0),
            root_nodes=st.get("rootNodes", 0),
            max_depth=st.get("maxDepth", 0),
            entity_count=st.get("entityCount", 0),
            link_count=st.get("linkCount", 0),
        )
    
    version = data.get("version", [1, 0])
    if isinstance(version, list):
        version = tuple(version)
    
    return EngramHeader(
        version=version,
        created=data.get("created", 0),
        modified=data.get("modified", 0),
        security=security,
        metadata=metadata,
        schema=schema,
        stats=stats,
    )


# ============== READER ==============

def read_engram(
    path: Union[str, Path],
    verify_integrity: bool = True,
    password: Optional[str] = None,
) -> EngramFile:
    """
    Read an engram file from disk.
    
    Args:
        path: Path to the .engram file
        verify_integrity: Whether to verify the integrity hash
        password: Password for encrypted files (not yet implemented)
    
    Returns:
        EngramFile containing header, nodes, entities, and links
    
    Raises:
        ValueError: If file is invalid or corrupted
        FileNotFoundError: If file doesn't exist
    """
    path = Path(path)
    data = path.read_bytes()
    
    # Verify magic bytes
    if data[:6] != MAGIC:
        raise ValueError(f"Invalid engram file: bad magic bytes")
    
    # Read version
    version_major = data[6]
    version_minor = data[7]
    
    # Read header length
    header_len = struct.unpack("<I", data[8:12])[0]
    
    # Decode header
    header_bytes = data[HEADER_OFFSET : HEADER_OFFSET + header_len]
    header_data = msgpack.unpackb(header_bytes, raw=False)
    header = _deserialize_header(header_data)
    
    # Decode payload
    payload_bytes = data[HEADER_OFFSET + header_len :]
    
    # Verify integrity
    if verify_integrity:
        expected_hash = header.security.integrity
        actual_hash = _hash_payload(payload_bytes)
        if expected_hash and expected_hash != actual_hash:
            raise ValueError("Integrity check failed: file may be corrupted or tampered")
    
    # Check encryption
    if header.security.encrypted:
        raise NotImplementedError("Encrypted engram files not yet supported in Python SDK")
    
    payload = msgpack.unpackb(payload_bytes, raw=False)
    
    # Deserialize nodes
    nodes = [_deserialize_node(n) for n in payload.get("nodes", [])]
    
    # TODO: Deserialize entities and links
    entities: list[Entity] = []
    links: list[MemoryLink] = []
    
    return EngramFile(
        header=header,
        nodes=nodes,
        entities=entities,
        links=links,
    )


# ============== WRITER ==============

def write_engram(
    engram: EngramFile,
    path: Union[str, Path],
    encrypt: bool = False,
    password: Optional[str] = None,
) -> None:
    """
    Write an engram file to disk.
    
    Args:
        engram: The EngramFile to write
        path: Output path
        encrypt: Whether to encrypt (not yet implemented)
        password: Password for encryption
    
    Raises:
        NotImplementedError: If encryption is requested
    """
    if encrypt:
        raise NotImplementedError("Encryption not yet supported in Python SDK")
    
    path = Path(path)
    
    # Serialize payload
    payload = {
        "nodes": [_serialize_node(n) for n in engram.nodes],
        "entities": [],  # TODO
        "links": [],  # TODO
    }
    payload_bytes = msgpack.packb(payload, use_bin_type=True)
    
    # Compute integrity hash
    integrity = _hash_payload(payload_bytes)
    engram.header.security.integrity = integrity
    
    # Serialize header
    header_dict = _serialize_header(engram.header)
    header_bytes = msgpack.packb(header_dict, use_bin_type=True)
    
    # Build file
    result = bytearray()
    result.extend(MAGIC)
    result.append(VERSION_MAJOR)
    result.append(VERSION_MINOR)
    result.extend(struct.pack("<I", len(header_bytes)))
    result.extend(header_bytes)
    result.extend(payload_bytes)
    
    # Write to disk
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(bytes(result))


# ============== CONVENIENCE ==============

def load(path: Union[str, Path]) -> EngramFile:
    """Shorthand for read_engram with default options."""
    return read_engram(path)


def save(engram: EngramFile, path: Union[str, Path]) -> None:
    """Shorthand for write_engram with default options."""
    write_engram(engram, path)
