package com.cube.index;

import com.cube.storage.StorageEngine;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

/**
 * Cubic Indexed Storage Engine - Combines LSM storage with cubic indexing.
 * 
 * Features:
 * - Cubic index tree (N³×6) for fast lookups
 * - 6-sided distribution for load balancing  
 * - Auto-expanding levels
 * - Prefix and range queries
 */
public class CubicIndexedStorage implements StorageEngine {
    
    private static final Logger logger = LoggerFactory.getLogger(CubicIndexedStorage.class);
    
    private final StorageEngine backingStorage;
    private final CubicIndexTree index;
    private final boolean indexEnabled;
    
    public CubicIndexedStorage(StorageEngine backingStorage) {
        this(backingStorage, true, 5, 20);
    }
    
    public CubicIndexedStorage(
            StorageEngine backingStorage,
            boolean indexEnabled,
            int initialLevels,
            int maxLevels) {
        
        this.backingStorage = backingStorage;
        this.indexEnabled = indexEnabled;
        this.index = indexEnabled ? new CubicIndexTree(initialLevels, maxLevels, true) : null;
        
        logger.info("Cubic indexed storage initialized (indexing: {})", indexEnabled);
    }
    
    @Override
    public void put(String key, byte[] value) throws IOException {
        // Write to backing storage
        backingStorage.put(key, value);
        
        // Update index
        if (indexEnabled) {
            index.put(key, value);
        }
    }
    
    @Override
    public byte[] get(String key) throws IOException {
        // Try index first for fast lookup
        if (indexEnabled) {
            byte[] value = index.get(key);
            if (value != null) {
                return value;
            }
        }
        
        // Fallback to backing storage
        return backingStorage.get(key);
    }
    
    @Override
    public boolean delete(String key) throws IOException {
        // Remove from index
        if (indexEnabled) {
            index.remove(key);
        }
        
        // Delete from backing storage
        return backingStorage.delete(key);
    }
    
    @Override
    public Iterator<String> scan(String prefix) throws IOException {
        if (indexEnabled) {
            // Use index for fast prefix search
            List<String> results = index.searchPrefix(prefix);
            return results.iterator();
        }
        
        // Fallback to backing storage
        return backingStorage.scan(prefix);
    }
    
    @Override
    public Iterator<Map.Entry<String, byte[]>> scanEntries(String prefix) throws IOException {
        if (indexEnabled) {
            // Use index for fast prefix search
            List<String> keys = index.searchPrefix(prefix);
            List<Map.Entry<String, byte[]>> entries = new ArrayList<>();
            
            for (String key : keys) {
                byte[] value = index.get(key);
                if (value != null) {
                    entries.add(new AbstractMap.SimpleEntry<>(key, value));
                }
            }
            
            return entries.iterator();
        }
        
        // Fallback to backing storage
        return backingStorage.scanEntries(prefix);
    }
    
    /**
     * Range search using cubic index
     */
    public List<String> rangeSearch(String startKey, String endKey) {
        if (!indexEnabled) {
            throw new UnsupportedOperationException("Index is disabled");
        }
        
        return index.searchRange(startKey, endKey);
    }
    
    /**
     * Get keys at a specific cubic level
     */
    public Set<String> getKeysAtLevel(int level) {
        if (!indexEnabled) {
            throw new UnsupportedOperationException("Index is disabled");
        }
        
        CubicIndexNode node = index.getLevel(level);
        return node != null ? node.getAllKeys() : Collections.emptySet();
    }
    
    /**
     * Get keys on a specific side of a level
     */
    public Set<String> getKeysOnSide(int level, CubicIndexNode.Side side) {
        if (!indexEnabled) {
            throw new UnsupportedOperationException("Index is disabled");
        }
        
        CubicIndexNode node = index.getLevel(level);
        return node != null ? node.getSide(side).keys() : Collections.emptySet();
    }
    
    /**
     * Rebalance the cubic index
     */
    public void rebalanceIndex() {
        if (!indexEnabled) {
            return;
        }
        
        index.rebalance();
        logger.info("Cubic index rebalanced");
    }
    
    /**
     * Rebuild index from backing storage
     */
    public void rebuildIndex() throws IOException {
        if (!indexEnabled) {
            return;
        }
        
        logger.info("Rebuilding cubic index from storage...");
        
        index.clear();
        
        // Scan all keys from backing storage
        Iterator<Map.Entry<String, byte[]>> entries = backingStorage.scanEntries("");
        int count = 0;
        
        while (entries.hasNext()) {
            Map.Entry<String, byte[]> entry = entries.next();
            index.put(entry.getKey(), entry.getValue());
            count++;
        }
        
        logger.info("Rebuilt cubic index with {} keys", count);
    }
    
    @Override
    public void flush() throws IOException {
        backingStorage.flush();
    }
    
    @Override
    public void compact() throws IOException {
        backingStorage.compact();
    }
    
    @Override
    public StorageStats getStats() {
        StorageStats backingStats = backingStorage.getStats();
        
        if (!indexEnabled) {
            return backingStats;
        }
        
        // Combine backing storage stats with index stats
        Map<String, Object> indexStats = index.getStats();
        
        return new StorageStats(
            backingStats.getTotalKeys(),
            backingStats.getTotalSize(),
            backingStats.getMemtableSize(),
            backingStats.getSstableCount()
        );
    }
    
    /**
     * Get cubic index statistics
     */
    public Map<String, Object> getIndexStats() {
        if (!indexEnabled) {
            return Collections.emptyMap();
        }
        
        return index.getStats();
    }
    
    /**
     * Print cubic index structure
     */
    public void printIndexStructure() {
        if (indexEnabled) {
            index.printStructure();
        }
    }
    
    /**
     * Get the cubic index tree (for advanced operations)
     */
    public CubicIndexTree getIndex() {
        return index;
    }
    
    public boolean isIndexEnabled() {
        return indexEnabled;
    }
    
    @Override
    public void close() throws IOException {
        if (indexEnabled) {
            index.clear();
        }
        backingStorage.close();
    }
}