package com.cube.storage;

import org.junit.jupiter.api.*;
import java.io.IOException;
import java.nio.file.*;
import java.util.*;

import static org.junit.jupiter.api.Assertions.*;

/**
 * Comprehensive tests for Cube LSM storage engine
 */
public class CubeStorageEngineTest {
    
    private Path testDir;
    private LSMStorageEngine storage;
    
    @BeforeEach
    public void setUp() throws IOException {
        testDir = Files.createTempDirectory("cube-test-");
        storage = new LSMStorageEngine(testDir.toString());
    }
    
    @AfterEach
    public void tearDown() throws IOException {
        if (storage != null) {
            storage.close();
        }
        deleteDirectory(testDir);
    }
    
    @Test
    public void testBasicPutAndGet() throws IOException {
        storage.put("user:1", "Alice".getBytes());
        
        byte[] result = storage.get("user:1");
        assertNotNull(result);
        assertEquals("Alice", new String(result));
    }
    
    @Test
    public void testMultipleOperations() throws IOException {
        // Insert 100 records
        for (int i = 0; i < 100; i++) {
            storage.put("key:" + i, ("value:" + i).getBytes());
        }
        
        // Verify all records
        for (int i = 0; i < 100; i++) {
            byte[] result = storage.get("key:" + i);
            assertNotNull(result);
            assertEquals("value:" + i, new String(result));
        }
    }
    
    @Test
    public void testUpdate() throws IOException {
        storage.put("key1", "value1".getBytes());
        assertEquals("value1", new String(storage.get("key1")));
        
        storage.put("key1", "value2".getBytes());
        assertEquals("value2", new String(storage.get("key1")));
    }
    
    @Test
    public void testDelete() throws IOException {
        storage.put("key1", "value1".getBytes());
        assertNotNull(storage.get("key1"));
        
        storage.delete("key1");
        
        // After delete, tombstone or null
        byte[] result = storage.get("key1");
        if (result != null) {
            assertArrayEquals(WriteAheadLog.TOMBSTONE, result);
        }
    }
    
    @Test
    public void testScanWithPrefix() throws IOException {
        storage.put("user:1:name", "Alice".getBytes());
        storage.put("user:1:email", "alice@example.com".getBytes());
        storage.put("user:2:name", "Bob".getBytes());
        storage.put("product:1:name", "Laptop".getBytes());
        
        Iterator<String> keys = storage.scan("user:1");
        Set<String> results = new HashSet<>();
        keys.forEachRemaining(results::add);
        
        assertEquals(2, results.size());
        assertTrue(results.contains("user:1:name"));
        assertTrue(results.contains("user:1:email"));
    }
    
    @Test
    public void testScanEntries() throws IOException {
        storage.put("user:1:name", "Alice".getBytes());
        storage.put("user:1:email", "alice@example.com".getBytes());
        
        Iterator<Map.Entry<String, byte[]>> entries = storage.scanEntries("user:1");
        Map<String, String> results = new HashMap<>();
        
        entries.forEachRemaining(e -> 
            results.put(e.getKey(), new String(e.getValue())));
        
        assertEquals(2, results.size());
        assertEquals("Alice", results.get("user:1:name"));
        assertEquals("alice@example.com", results.get("user:1:email"));
    }
    
    @Test
    public void testFlush() throws IOException, InterruptedException {
        storage.put("key1", "value1".getBytes());
        storage.put("key2", "value2".getBytes());
        
        storage.flush();
        
        // Wait a bit for async flush to complete
        Thread.sleep(100);
        
        byte[] value1 = storage.get("key1");
        byte[] value2 = storage.get("key2");
        
        assertNotNull(value1, "key1 should not be null after flush");
        assertNotNull(value2, "key2 should not be null after flush");
        
        assertEquals("value1", new String(value1));
        assertEquals("value2", new String(value2));
    }
    
    @Test
    public void testRecovery() throws IOException, InterruptedException {
        storage.put("key1", "value1".getBytes());
        storage.put("key2", "value2".getBytes());
        
        storage.close();
        
        // Wait a moment before reopening
        Thread.sleep(100);
        
        // Reopen and verify recovery
        storage = new LSMStorageEngine(testDir.toString());
        
        byte[] value1 = storage.get("key1");
        byte[] value2 = storage.get("key2");
        
        assertNotNull(value1, "key1 should be recovered");
        assertNotNull(value2, "key2 should be recovered");
        
        assertEquals("value1", new String(value1));
        assertEquals("value2", new String(value2));
    }
    
    @Test
    public void testStats() throws IOException {
        for (int i = 0; i < 50; i++) {
            storage.put("key:" + i, ("value:" + i).getBytes());
        }
        
        StorageEngine.StorageStats stats = storage.getStats();
        
        assertNotNull(stats);
        assertTrue(stats.getMemtableSize() > 0);
        System.out.println("Stats: " + stats);
    }
    
    @Test
    public void testNonExistentKey() throws IOException {
        byte[] result = storage.get("nonexistent");
        assertNull(result);
    }
    
    private void deleteDirectory(Path dir) throws IOException {
        if (Files.exists(dir)) {
            Files.walk(dir)
                .sorted(Comparator.reverseOrder())
                .forEach(path -> {
                    try {
                        Files.delete(path);
                    } catch (IOException e) {
                        // Ignore
                    }
                });
        }
    }
}