use bincode::{self, config::Configuration};
use sled::{self, Batch};
use std::sync::Arc;
use vlogger::*;
use shared::core::{self, Block, ChainData, Hasher, Address};

use crate::{
  db::error::DatabaseError,
  log,
};

use shared::core::print_error_chain;

static BINCODE_CONFIG: Configuration = bincode::config::standard();

const DB_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/database");

const DB_TREE: &str = "blocks";

const BLOCK_PREFIX: &str = "blocks";
const CHAIN_DATA_PREFIX: &str = "chain_data";
const DATA_TO_BLOCK_PREFIX: &str = "data_to_block";
const METADATA_PREFIX: &str = "metadata";
const BALANCE_PREFIX: &str = "balance";

const MEMPOOL_PREFIX: &str = "mempool";

const TIP_KEY: &str = "chain_tip";
const HEIGHT_KEY: &str = "chain_height";
const HEIGHT_TO_HASH_PREFIX: &str = "height_to_hash";

#[derive(Debug)]
pub struct ChainDb {
  db: sled::Tree,
}

fn data_prefix(key: &str) -> String {
  format!("{}:{}", CHAIN_DATA_PREFIX, key)
}

fn data_to_block_prefix(key: &str) -> String {
  format!("{}:{}", DATA_TO_BLOCK_PREFIX, key)
}

fn block_prefix(key: &str) -> String {
  format!("{}:{}", BLOCK_PREFIX, key)
}

fn metadata_prefix(key: &str) -> String {
  format!("{}:{}", METADATA_PREFIX, key)
}

fn height_to_hash_prefix(height: u64) -> String {
  format!("{}:{:020}", HEIGHT_TO_HASH_PREFIX, height)
}

fn balances_prefix(address: &str) -> String {
  format!("{}:{}", BALANCE_PREFIX, address)
}

fn nonce_prefix(address: &str) -> String {
  format!("{}:{}", ACCOUNT_PREFIX, address)
}

fn mempool_prefix(tx_hash: &str) -> String {
  format!("{}:{}", MEMPOOL_PREFIX, tx_hash)
}

fn account_prefix(addr: &str) -> String {
  format!("{}:{}", ACCOUNT_PREFIX, addr)
}

impl ChainDb {
  pub fn new(path: Option<String>) -> Result<ChainDb, DatabaseError> {
    let path = if path.is_some() {
      &path.unwrap()
    } else {
      DB_PATH
    };
    let config = sled::Config::new()
      .cache_capacity(512 * 1024)
      .segment_size(1024 * 1024)
      .path(&path);

    match config.open() {
      Ok(db) => {
        if db.was_recovered() {
          log(msg!(
            INFO,
            "Loaded Database from Previous state at: {}",
            path
          ));
        } else {
          log(msg!(INFO, "Created Database at {}", path));
        }
        let db = db.open_tree(DB_TREE)?;
        Ok(ChainDb { db })
      }
      Err(err) => {
        print_error_chain(&err.clone().into());
        Err(DatabaseError::Init(err.into()))
      }
    }
  }

  pub fn get_account(&self, addr: &Address) -> Result<Option<Account>, DatabaseError> {
    if let Some(bin_acc) = self.db.get(account_prefix(addr))? {
      let (acc, _) = bincode::decode_from_slice::<Account, _>(&bin_acc, BINCODE_CONFIG)?;
      Ok(Some(acc))
    } else {
      Ok(None)
    }
  }

  fn insert_account(&self, acc: &Account) -> Result<(), DatabaseError> {
    let mut batch = Batch::default();
    if let Some(_) = self.get_account(acc.address())? {
      return Err(DatabaseError::AccountExists(acc.address().to_string()));
    }
    let bin_acc = bincode::encode_to_vec(acc, BINCODE_CONFIG)?;
    let bin_nonce = bincode::encode_to_vec(&acc.nonce(), BINCODE_CONFIG)?;
    let bin_balance = bincode::encode_to_vec(&acc.balance(), BINCODE_CONFIG)?;
    batch.insert(account_prefix(acc.address()).as_str(), bin_acc);
    batch.insert(nonce_prefix(acc.address()).as_str(), bin_nonce);
    batch.insert(balances_prefix(acc.address()).as_str(), bin_balance);
    Ok(self.db.apply_batch(batch)?)
  }

  pub fn recover_mempool(&self) -> Result<Vec<ChainData>, DatabaseError> {
    let mem_tree = self.db.scan_prefix(MEMPOOL_PREFIX);
    let mut mem_vec = vec![];
    for mem in mem_tree {
      if let Ok((_, value)) = mem {
        let (value, _) = bincode::decode_from_slice::<ChainData, _>(&value, BINCODE_CONFIG)?;
        mem_vec.push(value);
      } else {
        return Err(DatabaseError::Corruption("Failed to recover Mempool".to_string()))
      }
    }
    Ok(mem_vec)
  }

  pub fn dump_mempool(&self, mempool: &[ChainData]) -> Result<(), DatabaseError> {
    let mut batch = Batch::default();
    for mem in mempool {
      let data_hash = Hasher::hash_chain_data(&mem);
      let bin_data = bincode::encode_to_vec(&mem, BINCODE_CONFIG)?;
      let key = mempool_prefix(&data_hash);
      batch.insert(key.as_str(), bin_data);
    }
    self.db.apply_batch(batch)?;
    Ok(())
  }

  pub async fn shutdown(&self) -> Result<usize, DatabaseError> {
    Ok(self.db.flush_async().await?)
  }

  pub fn add_balance(&self, address: &Address, new_balance: u64) -> Result<(), DatabaseError> {
    let new_balance = if let Some(old_balance) = self.get_balance(address)? {
      old_balance + new_balance
    } else {
      new_balance
    };
    Ok(self.set_balance(address, new_balance)?)
  }

  pub fn set_balance(&self, address: &Address, new_balance: u64) -> Result<(), DatabaseError> {
    let mut batch = Batch::default();
    let account_nonce = self.get_nonce(address)?;
    let account_nonce = account_nonce + 1;
    let bin_balance = bincode::encode_to_vec(new_balance, BINCODE_CONFIG)?;
    let bin_nonce = bincode::encode_to_vec(account_nonce, BINCODE_CONFIG)?;
    batch.insert(balances_prefix(address).as_str(), bin_balance);
    batch.insert(nonce_prefix(address).as_str(), bin_nonce);
    Ok(self.db.apply_batch(batch)?)
  }

  fn set_nonce(&self, address: &Address, nonce: u64) -> Result<(), DatabaseError> {
    let bin_nonce = bincode::encode_to_vec(&nonce, BINCODE_CONFIG)?;
    let address = nonce_prefix(address);
    self.db.insert(address, bin_nonce)?;
    Ok(())
  }

  fn get_nonce(&self, address: &Address) -> Result<u64, DatabaseError> {
    match self.db.get(nonce_prefix(address))? {
      Some(n) => {
        let (nonce, _) = bincode::decode_from_slice::<u64, _>(&n, BINCODE_CONFIG)?;
        Ok(nonce)
      },
      None => Err(DatabaseError::AccountNotFound(address.to_string()))
    }
  }

  pub fn set_balance_batch(&self, changes: Vec<(&Address, u64)>) -> Result<(), DatabaseError> {
    let mut batch = Batch::default();

    for (address, amount) in changes {
      let account_nonce = self.get_nonce(address)?;
      let account_nonce = account_nonce + 1;
      let balance_key = balances_prefix(&address);
      let bin_amount = bincode::encode_to_vec(amount, BINCODE_CONFIG)?;
      let bin_nonce = bincode::encode_to_vec(&account_nonce, BINCODE_CONFIG)?;
      batch.insert(nonce_prefix(address).as_str(), bin_nonce);
      batch.insert(balance_key.as_str(), bin_amount);
    }
    Ok(self.db.apply_batch(batch)?)
  }

  pub fn get_balance(&self, address: &Address) -> Result<Option<u64>, DatabaseError> {
    match self.db.get(balances_prefix(address))? {
      Some(b) => {
        let (balance, _) = bincode::decode_from_slice::<u64, _>(&b, BINCODE_CONFIG)?;
        Ok(Some(balance))
      },
      None => Ok(None),
    }
  }

  pub fn get_block_by_key(&self, block_hash: &str) -> Result<Option<Arc<core::Block>>, DatabaseError> {
    let block_hash = block_prefix(block_hash);
    if let Some(bin_block) = self.db.get(block_hash)? {
      let (block, _size) = bincode::decode_from_slice::<core::Block, _>(&bin_block, BINCODE_CONFIG)
        .map_err(|e| DatabaseError::Decode(e))?;
      Ok(Some(block.into()))
    } else {
      Ok(None)
    }
  }

  pub fn get_block_by_height(
    &self,
    height: u64,
  ) -> Result<Option<Arc<core::Block>>, DatabaseError> {
    if let Some(hash) = self.db.get(height_to_hash_prefix(height))? {
      let (hash_str, _) = bincode::decode_from_slice::<String, _>(&hash, BINCODE_CONFIG)?;
      Ok(self.get_block_by_key(&hash_str)?)
    } else {
      Ok(None)
    }
  }

  pub fn get_block_data(&self, block: &Block) -> Result<Arc<Vec<ChainData>>, DatabaseError> {
    let mut chain_data = Vec::new();

    for data_hash in &block.data {
      let data_hash = data_prefix(data_hash);
      if let Some(bin_data) = self.db.get(&data_hash)? {
        let (data, _) = bincode::decode_from_slice::<ChainData, _>(&bin_data, BINCODE_CONFIG)?;
        chain_data.push(data);
      } else {
        return Err(DatabaseError::MissingData(data_hash.clone()));
      }
    }

    Ok(Arc::new(chain_data))
  }

  pub fn get_all_blocks(&self) -> Result<Vec<Arc<core::Block>>, DatabaseError> {
    self
      .db
      .scan_prefix(BLOCK_PREFIX)
      .map(|res| -> Result<Arc<core::Block>, DatabaseError> {
        let (_key, value) = res?;
        let (block, _size) = bincode::decode_from_slice::<core::Block, _>(&value, BINCODE_CONFIG)
          .map_err(|e| DatabaseError::Decode(e))?;
        Ok(Arc::new(block))
      })
      .collect()
  }

  pub fn add_data(&self, data: &core::ChainData) -> Result<(), DatabaseError> {
    let bin_data = bincode::encode_to_vec(data, BINCODE_CONFIG)?;
    let data_hash = data_prefix(&Hasher::hash_chain_data(data));
    self.db.insert(data_hash, bin_data)?;
    Ok(())
  }

  pub fn add_block(&self, block: &core::Block) -> Result<(), DatabaseError> {
    let mut db_batch = Batch::default();
    let bin_block = bincode::encode_to_vec(block, BINCODE_CONFIG)?;
    db_batch.insert(block_prefix(block.head().block_hash()).as_str(), bin_block);
    db_batch.insert(
      height_to_hash_prefix(block.head().height).as_str(),
      block.head().block_hash(),
    );
    for data in block.data() {
      db_batch.insert(
        data_to_block_prefix(data.as_str()).as_str(),
        block.head().block_hash(),
      );
    }
    db_batch.insert(metadata_prefix(TIP_KEY).as_str(), block.head().block_hash());
    let bin_head = bincode::encode_to_vec(&block.head().height, BINCODE_CONFIG)?;
    db_batch.insert(
      metadata_prefix(HEIGHT_KEY).as_str(),
      bin_head
    );
    self.db.apply_batch(db_batch)?;
    Ok(())
  }
}
