"""
Engram Core

MemoryTree for managing hierarchical memory nodes.
"""

import hashlib
import time
import uuid
from pathlib import Path
from typing import Any, Optional, Union

import numpy as np
from numpy.typing import NDArray

from .types import (
    ContentType,
    DecayTier,
    EngramFile,
    EngramHeader,
    FileMetadata,
    FileStats,
    MemoryNode,
    NodeContent,
    NodeMetadata,
    QualityInfo,
    QualitySource,
    SchemaConfig,
    SecurityConfig,
    TemporalInfo,
)
from .io import read_engram, write_engram


def generate_id() -> str:
    """Generate a unique ID for a node."""
    return uuid.uuid4().hex[:16]


class MemoryTree:
    """
    Hierarchical memory tree for managing engram nodes.
    
    Provides methods for adding, updating, searching, and organizing memories.
    """
    
    def __init__(self, nodes: Optional[list[MemoryNode]] = None):
        """
        Initialize a MemoryTree.
        
        Args:
            nodes: Optional list of initial nodes
        """
        self._nodes: dict[str, MemoryNode] = {}
        self._root_ids: list[str] = []
        self._embeddings: dict[str, NDArray[np.float32]] = {}
        
        if nodes:
            for node in nodes:
                self._nodes[node.id] = node
                if node.parent_id is None:
                    self._root_ids.append(node.id)
                if node.embedding is not None:
                    self._embeddings[node.id] = node.embedding
    
    @classmethod
    def from_file(cls, path: Union[str, Path]) -> "MemoryTree":
        """Load a MemoryTree from an engram file."""
        engram = read_engram(path)
        return cls(engram.nodes)
    
    def save(self, path: Union[str, Path], metadata: Optional[FileMetadata] = None) -> None:
        """Save the MemoryTree to an engram file."""
        now = int(time.time() * 1000)
        
        nodes = list(self._nodes.values())
        
        header = EngramHeader(
            version=(1, 0),
            created=now,
            modified=now,
            security=SecurityConfig(),
            metadata=metadata or FileMetadata(source="engram-py"),
            schema=SchemaConfig(),
            stats=FileStats(
                total_chunks=len(nodes),
                root_nodes=len(self._root_ids),
                max_depth=max((n.depth for n in nodes), default=0),
            ),
        )
        
        engram = EngramFile(header=header, nodes=nodes)
        write_engram(engram, path)
    
    # ============== CRUD ==============
    
    def get(self, node_id: str) -> Optional[MemoryNode]:
        """Get a node by ID."""
        return self._nodes.get(node_id)
    
    def get_all(self) -> list[MemoryNode]:
        """Get all nodes."""
        return list(self._nodes.values())
    
    def get_roots(self) -> list[MemoryNode]:
        """Get all root nodes (no parent)."""
        return [self._nodes[id] for id in self._root_ids if id in self._nodes]
    
    def size(self) -> int:
        """Get total number of nodes."""
        return len(self._nodes)
    
    def add(
        self,
        content: str,
        *,
        content_type: ContentType = ContentType.TEXT,
        tags: Optional[list[str]] = None,
        parent_id: Optional[str] = None,
        embedding: Optional[NDArray[np.float32]] = None,
        custom: Optional[dict[str, Any]] = None,
    ) -> MemoryNode:
        """
        Add a new memory node.
        
        Args:
            content: The content string
            content_type: Type of content
            tags: Optional tags for the node
            parent_id: Optional parent node ID
            embedding: Optional embedding vector
            custom: Optional custom metadata
        
        Returns:
            The created MemoryNode
        """
        node_id = generate_id()
        now = int(time.time() * 1000)
        
        # Determine depth and path
        depth = 0
        path = f"/{node_id}"
        if parent_id and parent_id in self._nodes:
            parent = self._nodes[parent_id]
            depth = parent.depth + 1
            path = f"{parent.path}/{node_id}"
        
        node = MemoryNode(
            id=node_id,
            parent_id=parent_id,
            children=[],
            depth=depth,
            path=path,
            content=NodeContent(
                type=content_type,
                data=content,
                tokens=len(content) // 4,  # Rough estimate
            ),
            embedding=embedding,
            temporal=TemporalInfo(
                created=now,
                modified=now,
                accessed=now,
                decay_tier=DecayTier.HOT,
            ),
            quality=QualityInfo(
                score=1.0,
                confidence=1.0,
                source=QualitySource.DIRECT,
            ),
            metadata=NodeMetadata(
                tags=tags or [],
                custom=custom or {},
            ),
        )
        
        self._nodes[node_id] = node
        
        if parent_id is None:
            self._root_ids.append(node_id)
        elif parent_id in self._nodes:
            self._nodes[parent_id].children.append(node_id)
        
        if embedding is not None:
            self._embeddings[node_id] = embedding
        
        return node
    
    def update(self, node_id: str, **updates: Any) -> Optional[MemoryNode]:
        """
        Update a node's properties.
        
        Args:
            node_id: The node ID to update
            **updates: Properties to update
        
        Returns:
            The updated node, or None if not found
        """
        node = self._nodes.get(node_id)
        if not node:
            return None
        
        now = int(time.time() * 1000)
        
        if "content" in updates and node.content:
            node.content.data = updates["content"]
        
        if "tags" in updates and node.metadata:
            node.metadata.tags = updates["tags"]
        
        if "embedding" in updates:
            node.embedding = updates["embedding"]
            if updates["embedding"] is not None:
                self._embeddings[node_id] = updates["embedding"]
            elif node_id in self._embeddings:
                del self._embeddings[node_id]
        
        if node.temporal:
            node.temporal.modified = now
        
        return node
    
    def delete(self, node_id: str) -> bool:
        """
        Delete a node and optionally its children.
        
        Args:
            node_id: The node ID to delete
        
        Returns:
            True if deleted, False if not found
        """
        if node_id not in self._nodes:
            return False
        
        node = self._nodes[node_id]
        
        # Remove from parent's children
        if node.parent_id and node.parent_id in self._nodes:
            parent = self._nodes[node.parent_id]
            if node_id in parent.children:
                parent.children.remove(node_id)
        
        # Remove from root_ids
        if node_id in self._root_ids:
            self._root_ids.remove(node_id)
        
        # Remove node
        del self._nodes[node_id]
        
        # Remove embedding
        if node_id in self._embeddings:
            del self._embeddings[node_id]
        
        return True
    
    # ============== HIERARCHY ==============
    
    def get_parent(self, node_id: str) -> Optional[MemoryNode]:
        """Get the parent of a node."""
        node = self._nodes.get(node_id)
        if node and node.parent_id:
            return self._nodes.get(node.parent_id)
        return None
    
    def get_children(self, node_id: str) -> list[MemoryNode]:
        """Get all children of a node."""
        node = self._nodes.get(node_id)
        if not node:
            return []
        return [self._nodes[cid] for cid in node.children if cid in self._nodes]
    
    def get_ancestors(self, node_id: str) -> list[MemoryNode]:
        """Get all ancestors of a node (parent, grandparent, etc.)."""
        ancestors = []
        current = self._nodes.get(node_id)
        while current and current.parent_id:
            parent = self._nodes.get(current.parent_id)
            if parent:
                ancestors.append(parent)
                current = parent
            else:
                break
        return ancestors
    
    def get_descendants(self, node_id: str) -> list[MemoryNode]:
        """Get all descendants of a node (children, grandchildren, etc.)."""
        descendants = []
        stack = [node_id]
        while stack:
            current_id = stack.pop()
            node = self._nodes.get(current_id)
            if node:
                for child_id in node.children:
                    if child_id in self._nodes:
                        descendants.append(self._nodes[child_id])
                        stack.append(child_id)
        return descendants
    
    # ============== SEARCH ==============
    
    def find_by_tag(self, tag: str) -> list[MemoryNode]:
        """Find all nodes with a specific tag."""
        return [
            node for node in self._nodes.values()
            if node.metadata and tag in node.metadata.tags
        ]
    
    def find_by_path(self, path_prefix: str) -> list[MemoryNode]:
        """Find all nodes whose path starts with the given prefix."""
        return [
            node for node in self._nodes.values()
            if node.path.startswith(path_prefix)
        ]
    
    def find_by_depth(self, depth: int) -> list[MemoryNode]:
        """Find all nodes at a specific depth."""
        return [node for node in self._nodes.values() if node.depth == depth]
    
    def search_text(self, query: str) -> list[MemoryNode]:
        """Simple text search across all node content."""
        query_lower = query.lower()
        results = []
        for node in self._nodes.values():
            if node.content and isinstance(node.content.data, str):
                if query_lower in node.content.data.lower():
                    results.append(node)
        return results
    
    def search_similar(
        self,
        query_embedding: NDArray[np.float32],
        top_k: int = 10,
        min_score: float = 0.0,
    ) -> list[tuple[MemoryNode, float]]:
        """
        Search for similar nodes using cosine similarity.
        
        Args:
            query_embedding: The query embedding vector
            top_k: Maximum number of results
            min_score: Minimum similarity score (0-1)
        
        Returns:
            List of (node, score) tuples, sorted by similarity
        """
        if not self._embeddings:
            return []
        
        # Normalize query
        query_norm = query_embedding / np.linalg.norm(query_embedding)
        
        scores: list[tuple[str, float]] = []
        for node_id, embedding in self._embeddings.items():
            # Cosine similarity
            emb_norm = embedding / np.linalg.norm(embedding)
            score = float(np.dot(query_norm, emb_norm))
            if score >= min_score:
                scores.append((node_id, score))
        
        # Sort by score descending
        scores.sort(key=lambda x: x[1], reverse=True)
        
        # Return top_k results
        results = []
        for node_id, score in scores[:top_k]:
            if node_id in self._nodes:
                results.append((self._nodes[node_id], score))
        
        return results
    
    # ============== UTILITIES ==============
    
    def touch(self, node_id: str) -> None:
        """Update the accessed timestamp of a node."""
        node = self._nodes.get(node_id)
        if node and node.temporal:
            node.temporal.accessed = int(time.time() * 1000)
    
    def get_stats(self) -> dict[str, Any]:
        """Get statistics about the tree."""
        nodes = list(self._nodes.values())
        return {
            "total_nodes": len(nodes),
            "root_nodes": len(self._root_ids),
            "max_depth": max((n.depth for n in nodes), default=0),
            "nodes_with_embeddings": len(self._embeddings),
            "content_types": list(set(
                n.content.type.value for n in nodes
                if n.content and hasattr(n.content.type, 'value')
            )),
            "total_tags": len(set(
                tag for n in nodes
                if n.metadata
                for tag in n.metadata.tags
            )),
        }
