//! Block storage for L2 consensus using RocksDB.
//!
//! This crate provides block storage for L2 nodes and validators.
//! All block types are defined here for consistency across applications.

use rocksdb::{ColumnFamilyDescriptor, IteratorMode, Options, DB};
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use std::path::Path;
use std::sync::Arc;

const CF_BLOCKS: &str = "blocks";
const CF_CANONICAL: &str = "canonical";
const CF_STATE: &str = "state";
const CF_INDEX: &str = "index";

/// Storage error.
#[derive(Debug, thiserror::Error)]
pub enum StorageError {
    #[error("Database error: {0}")]
    Database(#[from] rocksdb::Error),

    #[error("Serialization error: {0}")]
    Serialization(String),

    #[error("Invalid data: {0}")]
    Invalid(String),
}

pub type StorageResult<T> = Result<T, StorageError>;

// ─────────────────────────────────────────────────────────────────────────────────
// Block Types
// ─────────────────────────────────────────────────────────────────────────────────

/// Block header.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BlockHeader {
    pub epoch: u64,
    pub height: u64,
    pub prev_block_root: [u8; 32],
    pub timestamp: u64,
    pub proposer_pubkey: Vec<u8>,
    pub extra_data: [u8; 32],
}

/// Emission entry.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Emission {
    pub pubkey: Vec<u8>,
    pub weight: u64,
}

/// Stored block with all data for proof generation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoredBlock {
    pub header: BlockHeader,
    pub data_hashes: Vec<[u8; 32]>,
    pub emissions: Vec<Emission>,
    pub data_root: [u8; 32],
    pub emissions_root: [u8; 32],
}

impl StoredBlock {
    /// Deserialize from raw bytes.
    pub fn from_bytes(bytes: &[u8]) -> StorageResult<Self> {
        deserialize(bytes)
    }

    /// Serialize to raw bytes.
    pub fn to_bytes(&self) -> StorageResult<Vec<u8>> {
        serialize(self)
    }
}

/// Data location in a block (for indexing).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DataLocation {
    pub epoch: u64,
    pub block_root: [u8; 32],
    pub index: usize,
}

// ─────────────────────────────────────────────────────────────────────────────────
// Block Storage
// ─────────────────────────────────────────────────────────────────────────────────

/// Block storage using RocksDB.
#[derive(Clone)]
pub struct BlockStorage {
    db: Arc<DB>,
}

impl BlockStorage {
    /// Open or create storage at the given path.
    pub fn open(path: &Path) -> StorageResult<Self> {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.create_missing_column_families(true);

        let cfs = vec![
            ColumnFamilyDescriptor::new(CF_BLOCKS, Options::default()),
            ColumnFamilyDescriptor::new(CF_CANONICAL, Options::default()),
            ColumnFamilyDescriptor::new(CF_STATE, Options::default()),
            ColumnFamilyDescriptor::new(CF_INDEX, Options::default()),
        ];

        let db = DB::open_cf_descriptors(&opts, path, cfs)?;
        Ok(Self { db: Arc::new(db) })
    }

    // ─────────────────────────────────────────────────────────────────────────────
    // Block Operations
    // ─────────────────────────────────────────────────────────────────────────────

    /// Store a block.
    pub fn put_block(&self, epoch: u64, root: &[u8; 32], block: &StoredBlock) -> StorageResult<()> {
        let cf = self.db.cf_handle(CF_BLOCKS).expect("blocks CF missing");
        let key = block_key(epoch, root);
        let value = block.to_bytes()?;
        self.db.put_cf(&cf, &key, &value)?;
        Ok(())
    }

    /// Get a block.
    pub fn get_block(&self, epoch: u64, root: &[u8; 32]) -> StorageResult<Option<StoredBlock>> {
        let cf = self.db.cf_handle(CF_BLOCKS).expect("blocks CF missing");
        let key = block_key(epoch, root);
        match self.db.get_cf(&cf, &key)? {
            Some(bytes) => Ok(Some(StoredBlock::from_bytes(&bytes)?)),
            None => Ok(None),
        }
    }

    /// Store a block as raw bytes.
    pub fn store_block(&self, epoch: u64, root: &[u8; 32], data: &[u8]) -> StorageResult<()> {
        let cf = self.db.cf_handle(CF_BLOCKS).expect("blocks CF missing");
        let key = block_key(epoch, root);
        self.db.put_cf(&cf, &key, data)?;
        Ok(())
    }

    /// Get a block as raw bytes.
    pub fn get_block_raw(&self, epoch: u64, root: &[u8; 32]) -> StorageResult<Option<Vec<u8>>> {
        let cf = self.db.cf_handle(CF_BLOCKS).expect("blocks CF missing");
        let key = block_key(epoch, root);
        Ok(self.db.get_cf(&cf, &key)?)
    }

    /// Check if a block exists.
    pub fn has_block(&self, epoch: u64, root: &[u8; 32]) -> StorageResult<bool> {
        let cf = self.db.cf_handle(CF_BLOCKS).expect("blocks CF missing");
        let key = block_key(epoch, root);
        Ok(self.db.get_cf(&cf, &key)?.is_some())
    }

    /// Get all block roots for an epoch.
    pub fn get_epoch_blocks(&self, epoch: u64) -> StorageResult<Vec<[u8; 32]>> {
        let cf = self.db.cf_handle(CF_BLOCKS).expect("blocks CF missing");
        let prefix = format!("{}:", epoch);
        let mut roots = Vec::new();

        let iter = self.db.prefix_iterator_cf(&cf, prefix.as_bytes());
        for item in iter {
            let (key, _) = item?;
            let key_str = String::from_utf8_lossy(&key);
            if let Some(hex_part) = key_str.strip_prefix(&prefix) {
                if let Ok(bytes) = hex::decode(hex_part) {
                    if bytes.len() == 32 {
                        let mut root = [0u8; 32];
                        root.copy_from_slice(&bytes);
                        roots.push(root);
                    }
                }
            }
        }
        Ok(roots)
    }

    // ─────────────────────────────────────────────────────────────────────────────
    // Canonical Chain
    // ─────────────────────────────────────────────────────────────────────────────

    /// Set the canonical block root for an epoch.
    pub fn set_canonical(&self, epoch: u64, root: &[u8; 32]) -> StorageResult<()> {
        let cf = self.db.cf_handle(CF_CANONICAL).expect("canonical CF missing");
        self.db.put_cf(&cf, epoch.to_be_bytes(), root)?;
        Ok(())
    }

    /// Get the canonical block root for an epoch.
    pub fn get_canonical(&self, epoch: u64) -> StorageResult<Option<[u8; 32]>> {
        let cf = self.db.cf_handle(CF_CANONICAL).expect("canonical CF missing");
        match self.db.get_cf(&cf, epoch.to_be_bytes())? {
            Some(bytes) if bytes.len() == 32 => {
                let mut root = [0u8; 32];
                root.copy_from_slice(&bytes);
                Ok(Some(root))
            }
            Some(_) => Err(StorageError::Invalid("Invalid canonical root length".into())),
            None => Ok(None),
        }
    }

    // ─────────────────────────────────────────────────────────────────────────────
    // State (Key-Value)
    // ─────────────────────────────────────────────────────────────────────────────

    /// Store arbitrary state.
    pub fn put_state(&self, key: &str, value: &[u8]) -> StorageResult<()> {
        let cf = self.db.cf_handle(CF_STATE).expect("state CF missing");
        self.db.put_cf(&cf, key.as_bytes(), value)?;
        Ok(())
    }

    /// Get state.
    pub fn get_state(&self, key: &str) -> StorageResult<Option<Vec<u8>>> {
        let cf = self.db.cf_handle(CF_STATE).expect("state CF missing");
        Ok(self.db.get_cf(&cf, key.as_bytes())?)
    }

    /// Store last processed L1 height.
    pub fn set_last_height(&self, height: u64) -> StorageResult<()> {
        self.put_state("last_height", &height.to_le_bytes())
    }

    /// Get last processed L1 height.
    pub fn get_last_height(&self) -> StorageResult<Option<u64>> {
        match self.get_state("last_height")? {
            Some(bytes) if bytes.len() == 8 => {
                Ok(Some(u64::from_le_bytes(bytes.try_into().unwrap())))
            }
            _ => Ok(None),
        }
    }

    // ─────────────────────────────────────────────────────────────────────────────
    // Data Index
    // ─────────────────────────────────────────────────────────────────────────────

    /// Index a data item by its hash.
    pub fn index_data(&self, hash: &[u8; 32], location: &DataLocation) -> StorageResult<()> {
        let cf = self.db.cf_handle(CF_INDEX).expect("index CF missing");
        let value = serialize(location)?;
        self.db.put_cf(&cf, hash, &value)?;
        Ok(())
    }

    /// Find data location by hash.
    pub fn find_data(&self, hash: &[u8; 32]) -> StorageResult<Option<DataLocation>> {
        let cf = self.db.cf_handle(CF_INDEX).expect("index CF missing");
        match self.db.get_cf(&cf, hash)? {
            Some(bytes) => Ok(Some(deserialize(&bytes)?)),
            None => Ok(None),
        }
    }

    // ─────────────────────────────────────────────────────────────────────────────
    // Maintenance
    // ─────────────────────────────────────────────────────────────────────────────

    /// Get all epochs that have canonical blocks.
    pub fn all_canonical_epochs(&self) -> StorageResult<Vec<u64>> {
        let cf = self.db.cf_handle(CF_CANONICAL).expect("canonical CF missing");
        let mut epochs = Vec::new();

        let iter = self.db.iterator_cf(&cf, IteratorMode::Start);
        for item in iter {
            let (key, _) = item?;
            if key.len() == 8 {
                let epoch = u64::from_be_bytes(key.as_ref().try_into().unwrap());
                epochs.push(epoch);
            }
        }

        epochs.sort();
        Ok(epochs)
    }

    /// Get all epochs that have any blocks (canonical or not).
    pub fn all_epochs_with_blocks(&self) -> StorageResult<Vec<u64>> {
        let cf = self.db.cf_handle(CF_BLOCKS).expect("blocks CF missing");
        let mut epochs = std::collections::HashSet::new();

        let iter = self.db.iterator_cf(&cf, IteratorMode::Start);
        for item in iter {
            let (key, _) = item?;
            // Keys are "epoch:hex_root" format
            let key_str = String::from_utf8_lossy(&key);
            if let Some(epoch_str) = key_str.split(':').next() {
                if let Ok(epoch) = epoch_str.parse::<u64>() {
                    epochs.insert(epoch);
                }
            }
        }

        let mut epochs: Vec<_> = epochs.into_iter().collect();
        epochs.sort();
        Ok(epochs)
    }

    /// Delete all data before a given epoch.
    pub fn prune_before(&self, epoch: u64) -> StorageResult<usize> {
        let blocks_cf = self.db.cf_handle(CF_BLOCKS).expect("blocks CF missing");
        let canonical_cf = self.db.cf_handle(CF_CANONICAL).expect("canonical CF missing");
        let mut deleted = 0;

        for e in 0..epoch {
            for root in self.get_epoch_blocks(e)? {
                let key = block_key(e, &root);
                self.db.delete_cf(&blocks_cf, &key)?;
                deleted += 1;
            }
            self.db.delete_cf(&canonical_cf, e.to_be_bytes())?;
        }

        Ok(deleted)
    }

    /// Flush all pending writes.
    pub fn flush(&self) -> StorageResult<()> {
        self.db.flush()?;
        Ok(())
    }

    /// Get total block count.
    pub fn block_count(&self) -> StorageResult<usize> {
        let cf = self.db.cf_handle(CF_BLOCKS).expect("blocks CF missing");
        let mut count = 0;
        let iter = self.db.iterator_cf(&cf, IteratorMode::Start);
        for item in iter {
            item?;
            count += 1;
        }
        Ok(count)
    }
}

fn block_key(epoch: u64, root: &[u8; 32]) -> Vec<u8> {
    format!("{}:{}", epoch, hex::encode(root)).into_bytes()
}

fn serialize<T: Serialize>(value: &T) -> StorageResult<Vec<u8>> {
    bincode::serialize(value).map_err(|e| StorageError::Serialization(e.to_string()))
}

fn deserialize<T: DeserializeOwned>(bytes: &[u8]) -> StorageResult<T> {
    bincode::deserialize(bytes).map_err(|e| StorageError::Serialization(e.to_string()))
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;

    fn test_block() -> StoredBlock {
        StoredBlock {
            header: BlockHeader {
                epoch: 1,
                height: 100,
                prev_block_root: [0u8; 32],
                timestamp: 1234567890,
                proposer_pubkey: vec![1, 2, 3],
                extra_data: [0u8; 32],
            },
            data_hashes: vec![[1u8; 32], [2u8; 32]],
            emissions: vec![Emission { pubkey: vec![4, 5, 6], weight: 100 }],
            data_root: [3u8; 32],
            emissions_root: [4u8; 32],
        }
    }

    #[test]
    fn test_block_storage() {
        let tmp = TempDir::new().unwrap();
        let storage = BlockStorage::open(tmp.path()).unwrap();

        let root = [42u8; 32];
        let block = test_block();

        storage.put_block(1, &root, &block).unwrap();
        assert!(storage.has_block(1, &root).unwrap());

        let retrieved = storage.get_block(1, &root).unwrap().unwrap();
        assert_eq!(retrieved.header.epoch, block.header.epoch);
    }

    #[test]
    fn test_raw_block_storage() {
        let tmp = TempDir::new().unwrap();
        let storage = BlockStorage::open(tmp.path()).unwrap();

        let root = [42u8; 32];
        let data = b"test block data";

        storage.store_block(1, &root, data).unwrap();
        assert!(storage.has_block(1, &root).unwrap());

        let retrieved = storage.get_block_raw(1, &root).unwrap().unwrap();
        assert_eq!(retrieved, data);
    }

    #[test]
    fn test_canonical() {
        let tmp = TempDir::new().unwrap();
        let storage = BlockStorage::open(tmp.path()).unwrap();

        let root = [42u8; 32];
        storage.set_canonical(1, &root).unwrap();

        let retrieved = storage.get_canonical(1).unwrap().unwrap();
        assert_eq!(retrieved, root);
    }

    #[test]
    fn test_data_index() {
        let tmp = TempDir::new().unwrap();
        let storage = BlockStorage::open(tmp.path()).unwrap();

        let hash = [1u8; 32];
        let location = DataLocation {
            epoch: 5,
            block_root: [2u8; 32],
            index: 3,
        };

        storage.index_data(&hash, &location).unwrap();
        let found = storage.find_data(&hash).unwrap().unwrap();
        assert_eq!(found.epoch, 5);
        assert_eq!(found.index, 3);
    }

    #[test]
    fn test_state() {
        let tmp = TempDir::new().unwrap();
        let storage = BlockStorage::open(tmp.path()).unwrap();

        storage.set_last_height(12345).unwrap();
        assert_eq!(storage.get_last_height().unwrap(), Some(12345));
    }
}
