"""
Basic tests for Engram Python SDK
"""

import os
import tempfile
from pathlib import Path

import numpy as np
import pytest

from engram import (
    ContentType,
    EngramFile,
    EngramHeader,
    FileMetadata,
    MemoryNode,
    MemoryTree,
    load,
    read_engram,
    save,
    write_engram,
)


class TestMemoryTree:
    """Tests for MemoryTree class."""
    
    def test_create_empty_tree(self):
        tree = MemoryTree()
        assert tree.size() == 0
        assert tree.get_roots() == []
    
    def test_add_memory(self):
        tree = MemoryTree()
        node = tree.add("Hello, world!", tags=["greeting"])
        
        assert tree.size() == 1
        assert node.content.data == "Hello, world!"
        assert "greeting" in node.metadata.tags
    
    def test_add_with_parent(self):
        tree = MemoryTree()
        parent = tree.add("Parent node")
        child = tree.add("Child node", parent_id=parent.id)
        
        assert child.parent_id == parent.id
        assert child.depth == 1
        assert child.id in parent.children
    
    def test_delete_memory(self):
        tree = MemoryTree()
        node = tree.add("To be deleted")
        node_id = node.id
        
        assert tree.delete(node_id)
        assert tree.get(node_id) is None
        assert tree.size() == 0
    
    def test_find_by_tag(self):
        tree = MemoryTree()
        tree.add("Node 1", tags=["a", "b"])
        tree.add("Node 2", tags=["b", "c"])
        tree.add("Node 3", tags=["c", "d"])
        
        results = tree.find_by_tag("b")
        assert len(results) == 2
    
    def test_search_text(self):
        tree = MemoryTree()
        tree.add("The quick brown fox")
        tree.add("The lazy dog")
        tree.add("Something else")
        
        results = tree.search_text("fox")
        assert len(results) == 1
        assert "fox" in results[0].content.data


class TestIO:
    """Tests for reading and writing engram files."""
    
    def test_write_and_read(self):
        tree = MemoryTree()
        tree.add("Test memory 1", tags=["test"])
        tree.add("Test memory 2", tags=["test"])
        
        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "test.engram"
            tree.save(path)
            
            # Verify file exists
            assert path.exists()
            
            # Read it back
            tree2 = MemoryTree.from_file(path)
            assert tree2.size() == 2
    
    def test_read_openstax(self):
        """Test reading the OpenStax US History engram if it exists."""
        openstax_path = Path.home() / ".allo" / "brains" / "openstax-us-history.engram"
        if not openstax_path.exists():
            pytest.skip("OpenStax engram not found")
        
        tree = MemoryTree.from_file(openstax_path)
        assert tree.size() > 0
        
        # Check it has chapter tags
        ch1_nodes = tree.find_by_tag("chapter-1")
        assert len(ch1_nodes) > 0
    
    def test_integrity_verification(self):
        """Test that corrupted files are rejected."""
        tree = MemoryTree()
        tree.add("Test memory")
        
        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "test.engram"
            tree.save(path)
            
            # Corrupt the file by modifying bytes
            data = bytearray(path.read_bytes())
            data[-10] ^= 0xFF  # Flip some bits in payload
            path.write_bytes(bytes(data))
            
            # Should raise on read
            with pytest.raises(ValueError, match="Integrity"):
                read_engram(path, verify_integrity=True)


class TestEmbeddings:
    """Tests for embedding functionality."""
    
    def test_add_with_embedding(self):
        tree = MemoryTree()
        embedding = np.random.rand(384).astype(np.float32)
        node = tree.add("With embedding", embedding=embedding)
        
        assert node.embedding is not None
        assert len(node.embedding) == 384
    
    def test_search_similar(self):
        tree = MemoryTree()
        
        # Add nodes with embeddings
        for i in range(10):
            embedding = np.random.rand(384).astype(np.float32)
            tree.add(f"Node {i}", embedding=embedding)
        
        # Search
        query = np.random.rand(384).astype(np.float32)
        results = tree.search_similar(query, top_k=5)
        
        assert len(results) == 5
        assert all(isinstance(r[1], float) for r in results)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
