// Copyright The Hush Developers  2019-2022
// Released under the GPLv3
use crate::lightwallet::LightWallet;

use std::sync::{Arc, RwLock, Mutex, mpsc::channel};
use std::sync::atomic::{AtomicBool, AtomicI32, AtomicUsize, Ordering};
use std::path::{Path, PathBuf};
use std::fs::File;
use std::collections::{HashSet, HashMap};
use std::cmp::{max, min};
use std::io;
use std::io::prelude::*;
use std::io::{BufReader, BufWriter, Error, ErrorKind};
use tokio::runtime::Runtime;
use tokio::time::{Duration};

use crate::grpc_client::RawTransaction;

use protobuf::parse_from_bytes;

use threadpool::ThreadPool;

use json::{object, array, JsonValue};
use zcash_primitives::transaction::{TxId, Transaction};
use zcash_client_backend::{constants::testnet, constants::mainnet, constants::regtest,};

use log::{info, warn, error, LevelFilter};
use log4rs::append::rolling_file::RollingFileAppender;
use log4rs::encode::pattern::PatternEncoder;
use log4rs::config::{Appender, Config, Root};
use log4rs::filter::threshold::ThresholdFilter;
use log4rs::append::rolling_file::policy::compound::{
    CompoundPolicy,
    trigger::size::SizeTrigger,
    roll::fixed_window::FixedWindowRoller,
};

use crate::grpcconnector::{self, *};


use crate::ANCHOR_OFFSET;

mod checkpoints;

pub const DEFAULT_SERVER: &str = "https://lite.dragonx.is";
pub const WALLET_NAME: &str    = "silentdragonxlite-wallet.dat";
pub const LOGFILE_NAME: &str   = "silentdragonxlite-wallet.debug.log";

#[derive(Clone, Debug)]
pub struct WalletStatus {
    pub is_syncing: bool,
    pub total_blocks: u64,
    pub synced_blocks: u64,
}

impl WalletStatus {
    pub fn new() -> Self {
        WalletStatus {
            is_syncing: false,
            total_blocks: 0,
            synced_blocks: 0
        }
    }
}

#[derive(Clone, Debug)]
pub struct LightClientConfig {
    pub server                      : http::Uri,
    pub chain_name                  : String,
    pub sapling_activation_height   : u64,
    pub consensus_branch_id         : String,
    pub anchor_offset               : u32,
    pub no_cert_verification        : bool,
    pub data_dir                    : Option<String>
}

impl LightClientConfig {

    // Create an unconnected (to any server) config to test for local wallet etc...
    pub fn create_unconnected(chain_name: String, dir: Option<String>) -> LightClientConfig {
        LightClientConfig {
            server                      : http::Uri::default(),
            chain_name                  : chain_name,
            sapling_activation_height   : 0,
            consensus_branch_id         : "".to_string(),
            anchor_offset               : ANCHOR_OFFSET,
            no_cert_verification        : false,
            data_dir                    : dir,
        }
    }

    pub fn create(server: http::Uri, dangerous: bool) -> io::Result<(LightClientConfig, u64)> {
        use std::net::ToSocketAddrs;
        // Test for a connection first
        format!("{}:{}", server.host().unwrap(), server.port().unwrap())
            .to_socket_addrs()?
            .next()
            .ok_or(std::io::Error::new(ErrorKind::ConnectionRefused, "Couldn't resolve server!"))?;

        // Do a getinfo first, before opening the wallet
        let info = grpcconnector::get_info(&server, dangerous)
            .map_err(|e| std::io::Error::new(ErrorKind::ConnectionRefused, e))?;

        // Create a Light Client Config
        let config = LightClientConfig {
            server,
            chain_name                  : info.chain_name,
            sapling_activation_height   : info.sapling_activation_height,
            consensus_branch_id         : info.consensus_branch_id,
            anchor_offset               : ANCHOR_OFFSET,
            no_cert_verification        : dangerous,
            data_dir                    : None,
        };

        Ok((config, info.block_height))
    }


    /// Build the Logging config
    pub fn get_log_config(&self) -> io::Result<Config> {
        let window_size = 3; // log0, log1, log2
        let fixed_window_roller =
            FixedWindowRoller::builder().build("zecwallet-light-wallet-log{}",window_size).unwrap();
        let size_limit = 5 * 1024 * 1024; // 5MB as max log file size to roll
        let size_trigger = SizeTrigger::new(size_limit);
        let compound_policy = CompoundPolicy::new(Box::new(size_trigger),Box::new(fixed_window_roller));

        Config::builder()
            .appender(
                Appender::builder()
                    .filter(Box::new(ThresholdFilter::new(LevelFilter::Info)))
                    .build(
                        "logfile",
                        Box::new(
                            RollingFileAppender::builder()
                                .encoder(Box::new(PatternEncoder::new("{d} {l}::{m}{n}")))
                                .build(self.get_log_path(), Box::new(compound_policy))?,
                        ),
                    ),
            )
            .build(
                Root::builder()
                    .appender("logfile")
                    .build(LevelFilter::Debug),
            )
            .map_err(|e|Error::new(ErrorKind::Other, format!("{}", e)))
    }

    pub fn get_zcash_data_path(&self) -> Box<Path> {
        let mut zcash_data_location; 
        if self.data_dir.is_some() {
            zcash_data_location = PathBuf::from(&self.data_dir.as_ref().unwrap());
        } else {
            if cfg!(target_os="macos") || cfg!(target_os="windows") {
                zcash_data_location = dirs::data_dir().expect("Couldn't determine app data directory!");
                zcash_data_location.push("silentdragonxlite");
            } else {
                zcash_data_location = dirs::home_dir().expect("Couldn't determine home directory!");
                zcash_data_location.push(".silentdragonxlite");
            };

            match &self.chain_name[..] {
                "main" => {},
                "test" => zcash_data_location.push("testnet3"),
                "regtest" => zcash_data_location.push("regtest"),
                c         => panic!("Unknown chain {}", c),
            };
        }
        
       // Create directory if it doesn't exist
        match std::fs::create_dir_all(zcash_data_location.clone()) {
            Ok(_) => zcash_data_location.into_boxed_path(),
            Err(e) => {
                eprintln!("Couldn't create zcash directory!\n{}", e);
                panic!("Couldn't create zcash directory!");
            }
        }
    }

    pub fn get_wallet_path(&self) -> Box<Path> {
        let mut wallet_location = self.get_zcash_data_path().into_path_buf();
        wallet_location.push(WALLET_NAME);
        
        wallet_location.into_boxed_path()
    }

    pub fn wallet_exists(&self) -> bool {
        return self.get_wallet_path().exists()
    }

    pub fn get_zcash_params_path(&self) -> io::Result<Box<Path>> {
        let mut zcash_params = self.get_zcash_data_path().into_path_buf();
        zcash_params.push("..");
        if cfg!(target_os="macos") || cfg!(target_os="windows") {
            zcash_params.push("ZcashParams");
        } else {
            zcash_params.push(".zcash-params");
        }

        match std::fs::create_dir_all(zcash_params.clone()) {
            Ok(_) => Ok(zcash_params.into_boxed_path()),
            Err(e) => {
                eprintln!("Couldn't create zcash params directory\n{}", e);
                Err(e)
            }
        }
    }

    pub fn get_log_path(&self) -> Box<Path> {
        let mut log_path = self.get_zcash_data_path().into_path_buf();
        log_path.push(LOGFILE_NAME);

        log_path.into_boxed_path()
    }

    pub fn get_initial_state(&self, height: u64) -> Option<(u64, &str, &str)> {
        checkpoints::get_closest_checkpoint(&self.chain_name, height)
    }

    pub fn get_server_or_default(server: Option<String>) -> http::Uri {
        match server {
            Some(s) => {
                let mut s = if s.starts_with("http") {s} else { "http://".to_string() + &s};
                let uri: http::Uri = s.parse().unwrap();
                if uri.port().is_none() {
                    s = s + ":443";
                }
                s
            }
            None    => DEFAULT_SERVER.to_string()
        }.parse().unwrap()
    }

    pub fn get_coin_type(&self) -> u32 {
        match &self.chain_name[..] {
            "main"    => mainnet::COIN_TYPE,
            "test"    => testnet::COIN_TYPE,
            "regtest" => regtest::COIN_TYPE,
            c         => panic!("Unknown chain {}", c)
        }
    }

    pub fn hrp_sapling_address(&self) -> &str {
        match &self.chain_name[..] {
            "main"    => mainnet::HRP_SAPLING_PAYMENT_ADDRESS,
            "test"    => testnet::HRP_SAPLING_PAYMENT_ADDRESS,
            "regtest" => regtest::HRP_SAPLING_PAYMENT_ADDRESS,
            c         => panic!("Unknown chain {}", c)
        }
    }

    pub fn hrp_sapling_private_key(&self) -> &str {
        match &self.chain_name[..] {
            "main"    => mainnet::HRP_SAPLING_EXTENDED_SPENDING_KEY,
            "test"    => testnet::HRP_SAPLING_EXTENDED_SPENDING_KEY,
            "regtest" => regtest::HRP_SAPLING_EXTENDED_SPENDING_KEY,
            c         => panic!("Unknown chain {}", c)
        }
    }

    pub fn hrp_sapling_viewing_key(&self) -> &str {
        match &self.chain_name[..] {
            "main"    => mainnet::HRP_SAPLING_EXTENDED_FULL_VIEWING_KEY,
            "test"    => testnet::HRP_SAPLING_EXTENDED_FULL_VIEWING_KEY,
            "regtest" => regtest::HRP_SAPLING_EXTENDED_FULL_VIEWING_KEY,
            c         => panic!("Unknown chain {}", c)
        }
    }

    pub fn base58_pubkey_address(&self) -> [u8; 1] {
        match &self.chain_name[..] {
            "main"    => mainnet::B58_PUBKEY_ADDRESS_PREFIX,
        

            c         => panic!("Unknown chain {}", c)
        }
    }


    pub fn base58_script_address(&self) -> [u8; 1] {
        match &self.chain_name[..] {
            "main"    => mainnet::B58_SCRIPT_ADDRESS_PREFIX,
            

            c         => panic!("Unknown chain {}", c)
        }
    }

    pub fn base58_secretkey_prefix(&self) -> [u8; 1] {
        match &self.chain_name[..] {
            "main"    => [0xBC],
            "test"    => [0xEF],
            "regtest" => [0xEF],
            c         => panic!("Unknown chain {}", c)
        }
    }
}

pub struct LightClient {
    pub wallet          : Arc<RwLock<LightWallet>>,

    pub config          : LightClientConfig,

    // zcash-params
    pub sapling_output  : Vec<u8>,
    pub sapling_spend   : Vec<u8>,

    sync_lock           : Mutex<()>,
    sync_status         : Arc<RwLock<WalletStatus>>, // The current syncing status of the Wallet.
    pub shutdown_flag   : Arc<AtomicBool>,           // Signal mempool threads to stop
}

impl LightClient {
    
    pub fn set_wallet_initial_state(&self, height: u64) {
        use std::convert::TryInto;

        let state = self.config.get_initial_state(height);

        match state {
            Some((height, hash, tree)) => self.wallet.read().unwrap().set_initial_block(height.try_into().unwrap(), hash, tree),
            _ => true,
        };
    }

    fn write_file_if_not_exists(dir: &Box<Path>, name: &str, bytes: &[u8]) -> io::Result<()> {
        let mut file_path = dir.to_path_buf();
        file_path.push(name);
        if !file_path.exists() {
            let mut file = File::create(&file_path)?;
            file.write_all(bytes)?;
        }

        Ok(())
    }

    #[cfg(feature = "embed_params")]
    fn read_sapling_params(&mut self) {
        // Read Sapling Params
        use crate::SaplingParams;
        self.sapling_output.extend_from_slice(SaplingParams::get("sapling-output.params").unwrap().as_ref());
        self.sapling_spend.extend_from_slice(SaplingParams::get("sapling-spend.params").unwrap().as_ref());
    }
    pub fn set_sapling_params(&mut self, sapling_output: &[u8], sapling_spend: &[u8]) -> Result<(), String> {
        use sha2::{Sha256, Digest};
        // The hashes of the params need to match
        const SAPLING_OUTPUT_HASH: &str = "2f0ebbcbb9bb0bcffe95a397e7eba89c29eb4dde6191c339db88570e3f3fb0e4";
        const SAPLING_SPEND_HASH: &str = "8e48ffd23abb3a5fd9c5589204f32d9c31285a04b78096ba40a79b75677efc13";
        if SAPLING_OUTPUT_HASH.to_string() != hex::encode(Sha256::digest(&sapling_output)) {
            return Err(format!("sapling-output hash didn't match. expected {}, found {}", SAPLING_OUTPUT_HASH, hex::encode(Sha256::digest(&sapling_output)) ))
        }
        if SAPLING_SPEND_HASH.to_string() != hex::encode(Sha256::digest(&sapling_spend)) {
            return Err(format!("sapling-spend hash didn't match. expected {}, found {}", SAPLING_SPEND_HASH, hex::encode(Sha256::digest(&sapling_spend)) ))
        }
        // Will not overwrite previous params
        if self.sapling_output.is_empty() {
            self.sapling_output.extend_from_slice(sapling_output);
        }
        if self.sapling_spend.is_empty() {
            self.sapling_spend.extend_from_slice(sapling_spend);
        }

        // Ensure that the sapling params are stored on disk properly as well. 
        match self.config.get_zcash_params_path() {
            Ok(zcash_params_dir) => {
                // Create the sapling output and spend params files
                match LightClient::write_file_if_not_exists(&zcash_params_dir, "sapling-output.params", &self.sapling_output) {
                    Ok(_) => {},
                    Err(e) => eprintln!("Warning: Couldn't write the output params!\n{}", e)
                };

                match LightClient::write_file_if_not_exists(&zcash_params_dir, "sapling-spend.params", &self.sapling_spend) {
                    Ok(_) => {},
                    Err(e) => eprintln!("Warning: Couldn't write the output params!\n{}", e)
                }
            },
            Err(e) => {
                eprintln!("{}", e);
            }
        };

        Ok(())
    }

    /// Method to create a test-only version of the LightClient
    #[allow(dead_code)]
    pub fn unconnected(seed_phrase: String, dir: Option<String>) -> io::Result<Self> {
        let config = LightClientConfig::create_unconnected("test".to_string(), dir);
        let mut l = LightClient {
                wallet          : Arc::new(RwLock::new(LightWallet::new(Some(seed_phrase), &config, 0,0 )?)),
                config          : config.clone(),
                sapling_output  : vec![], 
                sapling_spend   : vec![],
                sync_lock       : Mutex::new(()),
                sync_status     : Arc::new(RwLock::new(WalletStatus::new())),
                shutdown_flag   : Arc::new(AtomicBool::new(false)),
            };

        l.set_wallet_initial_state(0);
        
        #[cfg(feature = "embed_params")]
        l.read_sapling_params();

        info!("Created new wallet!");
        info!("Created LightClient to {}", &config.server);

        Ok(l)
    }

    /// Create a brand new wallet with a new seed phrase. Will fail if a wallet file 
    /// already exists on disk
    pub fn new(config: &LightClientConfig, latest_block: u64) -> io::Result<Self> {
        if config.wallet_exists() {
            return Err(Error::new(ErrorKind::AlreadyExists,
                    "Cannot create a new wallet from seed, because a wallet already exists"));
        }

        let mut l = LightClient {
                wallet          : Arc::new(RwLock::new(LightWallet::new(None, config, latest_block, 0)?)),
                config          : config.clone(),
                sapling_output  : vec![], 
                sapling_spend   : vec![],
                sync_lock       : Mutex::new(()),
                sync_status     : Arc::new(RwLock::new(WalletStatus::new())),
                shutdown_flag   : Arc::new(AtomicBool::new(false)),
            };

        l.set_wallet_initial_state(latest_block);
        
        #[cfg(feature = "embed_params")]
        l.read_sapling_params();

        info!("Created new wallet with a new seed!");
        info!("Created LightClient to {}", &config.server);
        l.do_save().map_err(|s| io::Error::new(ErrorKind::PermissionDenied, s))?;

        Ok(l)
    }

    pub fn new_from_phrase(seed_phrase: String, config: &LightClientConfig, birthday: u64,number: u64, overwrite: bool) -> io::Result<Self> {
        if !overwrite && config.wallet_exists() {
            return Err(Error::new(ErrorKind::AlreadyExists,
                    "Cannot create a new wallet from seed, because a wallet already exists"));
        }

        let mut l = LightClient {
                wallet          : Arc::new(RwLock::new(LightWallet::new(Some(seed_phrase), config, birthday, number)?)),
                config          : config.clone(),
                sapling_output  : vec![], 
                sapling_spend   : vec![],
                sync_lock       : Mutex::new(()),
                sync_status     : Arc::new(RwLock::new(WalletStatus::new())),
                shutdown_flag   : Arc::new(AtomicBool::new(false)),
            };

        // println!("Setting birthday to {}", birthday);
        l.set_wallet_initial_state(birthday);
        #[cfg(feature = "embed_params")]
        l.read_sapling_params();


        info!("Created new wallet!");
        info!("Created LightClient to {}", &config.server);
        l.do_save().map_err(|s| io::Error::new(ErrorKind::PermissionDenied, s))?;

        Ok(l)
    }

    pub fn read_from_disk(config: &LightClientConfig) -> io::Result<Self> {
        if !config.wallet_exists() {
            // Try to recover from backup
            let bak_path = config.get_wallet_path().with_extension("dat.bak");
            if bak_path.exists() {
                warn!("Wallet file missing but backup found, attempting recovery from {:?}", bak_path);
                std::fs::copy(&bak_path, config.get_wallet_path())?;
                info!("Wallet recovered from backup");
            } else {
                return Err(Error::new(ErrorKind::AlreadyExists,
                        format!("Cannot read wallet. No file at {}", config.get_wallet_path().display())));
            }
        }

        // Try to open the wallet file; if it fails (corrupted/truncated), try the backup
        let wallet = match File::open(config.get_wallet_path())
            .and_then(|f| {
                let mut file_buffer = BufReader::new(f);
                LightWallet::read(&mut file_buffer, config)
            }) {
            Ok(w) => w,
            Err(e) => {
                warn!("Failed to read wallet file: {}, trying backup", e);
                let bak_path = config.get_wallet_path().with_extension("dat.bak");
                if bak_path.exists() {
                    std::fs::copy(&bak_path, config.get_wallet_path())?;
                    let mut file_buffer = BufReader::new(File::open(config.get_wallet_path())?);
                    LightWallet::read(&mut file_buffer, config)?
                } else {
                    return Err(e);
                }
            }
        };

        let  mut lc = LightClient {
            wallet          : Arc::new(RwLock::new(wallet)),
            config          : config.clone(),
            sapling_output  : vec![], 
            sapling_spend   : vec![],
            sync_lock       : Mutex::new(()),
            sync_status     : Arc::new(RwLock::new(WalletStatus::new())),
            shutdown_flag   : Arc::new(AtomicBool::new(false)),
        };

        #[cfg(feature = "embed_params")]
        lc.read_sapling_params();

        info!("Read wallet with birthday {}", lc.wallet.read().unwrap().get_first_tx_block());
        info!("Created LightClient to {}", &config.server);

        Ok(lc)
    }

    pub fn init_logging(&self) -> io::Result<()> {
        // Configure logging first.
        let log_config = self.config.get_log_config()?;
        log4rs::init_config(log_config).map_err(|e| {
            std::io::Error::new(ErrorKind::Other, e)
        })?;

        Ok(())
    }

        pub fn attempt_recover_seed(config: &LightClientConfig, password: Option<String>) -> Result<String, String> {
        use std::io::prelude::*;
        use byteorder::{LittleEndian, ReadBytesExt};
        use libflate::gzip::Decoder;
        use bip39::{Mnemonic, Language};
        use zcash_primitives::serialize::Vector;

        let mut inp = BufReader::new(File::open(config.get_wallet_path()).unwrap());
        let version = inp.read_u64::<LittleEndian>().unwrap();
        println!("Reading wallet version {}", version);

        // At version 5, we're writing the rest of the file as a compressed stream (gzip)
        let mut reader: Box<dyn Read> = if version != 5 {
            Box::new(inp)
        } else {
            Box::new(Decoder::new(inp).unwrap())
        };

        let encrypted = if version >= 4 {
            reader.read_u8().unwrap() > 0
        } else {
            false
        };

        if encrypted && password.is_none() {
            return Err("The wallet is encrypted and a password was not specified. Please specify the password with '--password'!".to_string());
        }

        let mut enc_seed = [0u8; 48];
        if version >= 4 {
            reader.read_exact(&mut enc_seed).unwrap();
        }

        let nonce = if version >= 4 {
            Vector::read(&mut reader, |r| r.read_u8()).unwrap()
        } else {
            vec![]
        };

        let phrase = if encrypted {
            use sodiumoxide::crypto::secretbox;
            use crate::lightwallet::double_sha256;

         // Get the doublesha256 of the password, which is the right length
         let key = secretbox::Key::from_slice(&double_sha256(password.unwrap().as_bytes())).unwrap();
         let nonce = secretbox::Nonce::from_slice(&nonce).unwrap();

         let seed = match secretbox::open(&enc_seed, &nonce, &key) {
            Ok(s) => s,
            Err(_) => return Err("Decryption failed. Is your password correct?".to_string())
        };

        Mnemonic::from_entropy(&seed, Language::English)
    } else {
        // Seed
        let mut seed_bytes = [0u8; 32];
        reader.read_exact(&mut seed_bytes).unwrap();

        Mnemonic::from_entropy(&seed_bytes, Language::English) 
    }.map_err(|e| format!("Failed to read seed. {:?}", e));

    phrase.map(|m| m.phrase().to_string())
    }


    pub fn last_scanned_height(&self) -> u64 {
        self.wallet.read().unwrap().last_scanned_height() as u64
    }

    // Export private keys
    pub fn do_export(&self, addr: Option<String>) -> Result<JsonValue, &str> {
        if !self.wallet.read().unwrap().is_unlocked_for_spending() {
            error!("Wallet is locked");
            return Err("Wallet is locked");
        }

        // Clone address so it can be moved into the closure
        let address = addr.clone();
        let wallet = self.wallet.read().unwrap();
        // Go over all z addresses
        let z_keys = wallet.get_z_private_keys().iter()
        .filter( move |(addr, _, _)| address.is_none() || address.as_ref() == Some(addr))
        .map( |(addr, pk, vk)|
                object!{
                    "address"     => addr.clone(),
                    "private_key" => pk.clone(),
                    "viewing_key" => vk.clone(),
                }
            ).collect::<Vec<JsonValue>>();

        // Clone address so it can be moved into the closure
        let address = addr.clone();

        // Go over all t addresses
        let t_keys = wallet.get_t_secret_keys().iter()
            .filter( move |(addr, _)| address.is_none() || address.as_ref() == Some(addr))
            .map( |(addr, sk)|
                object!{
                    "address"     => addr.clone(),
                    "private_key" => sk.clone(),
                }
            ).collect::<Vec<JsonValue>>();

        let mut all_keys = vec![];
        all_keys.extend_from_slice(&z_keys);
        all_keys.extend_from_slice(&t_keys);

        Ok(all_keys.into())
    }

    pub fn do_address(&self) -> JsonValue {
        let wallet = self.wallet.read().unwrap();

        // Collect z addresses
        let z_addresses = wallet.get_all_zaddresses();

        // Collect t addresses
        let t_addresses = wallet.get_all_taddresses().iter().map( |a| a.clone() )
                            .collect::<Vec<String>>();

        object!{
            "z_addresses" => z_addresses,
            "t_addresses" => t_addresses,
        }
    }

    pub fn do_balance(&self) -> JsonValue {
        let wallet = self.wallet.read().unwrap();

        // Collect z addresses
        let z_addresses = wallet.get_all_zaddresses().iter().map(|zaddress| {
            object!{
                "address" => zaddress.clone(),
                "zbalance" => wallet.zbalance(Some(zaddress.clone())),
                "unconfirmed" => wallet.unconfirmed_zbalance(Some(zaddress.clone())),
                "verified_zbalance" => wallet.verified_zbalance(Some(zaddress.clone())),
                "spendable_zbalance" => wallet.spendable_zbalance(Some(zaddress.clone()))
            }
        }).collect::<Vec<JsonValue>>();

        // Collect t addresses
        let t_addresses =  wallet.get_all_taddresses().iter().map( |address| {
            // Get the balance for this address
            let balance = wallet.tbalance(Some(address.clone())) ;
            
            object!{
                "address" => address.clone(),
                "balance" => balance,
            }
        }).collect::<Vec<JsonValue>>();

        object!{
            "zbalance"           => wallet.zbalance(None),
            "unconfirmed"        => wallet.unconfirmed_zbalance(None),
            "verified_zbalance"  => wallet.verified_zbalance(None),
            "spendable_zbalance" => wallet.spendable_zbalance(None),
            "tbalance"           => wallet.tbalance(None), 
            "z_addresses"        => z_addresses,
            "t_addresses"        => t_addresses,
        }
    }

    pub fn do_save(&self) -> Result<(), String> {        
        // If the wallet is encrypted but unlocked, lock it again.
        {
            let mut wallet = self.wallet.write().unwrap();
            if wallet.is_encrypted() && wallet.is_unlocked_for_spending() {
                match wallet.lock() {
                    Ok(_) => {},
                    Err(e) => {
                        let err = format!("ERR: {}", e);
                        error!("{}", err);
                        return Err(e.to_string());
                    }
                }
            }
        }        

        let r;
        {
            // Prevent any overlapping syncs during save, and don't save in the middle of a sync
            let _lock = self.sync_lock.lock().unwrap();

            let wallet_path = self.config.get_wallet_path();
            let tmp_path = wallet_path.with_extension("dat.tmp");
            let bak_path = wallet_path.with_extension("dat.bak");

            // Write to a temporary file first (atomic save)
            let wallet = self.wallet.write().unwrap();
            let mut file_buffer = BufWriter::with_capacity(
                1_000_000, // 1 MB write buffer
                File::create(&tmp_path).map_err(|e| format!("Failed to create temp file: {}", e))?);

                r = match wallet.write(&mut file_buffer) {
                    Ok(_) => Ok(()),
                    Err(e) => {
                        let err = format!("ERR: {}", e);
                        error!("{}", err);
                        // Clean up temp file on write failure
                        let _ = std::fs::remove_file(&tmp_path);
                        Err(e.to_string())
                    }
                };

                file_buffer.flush().map_err(|e| {
                    let _ = std::fs::remove_file(&tmp_path);
                    format!("{}", e)
                })?;

            // Only proceed with rename if the write succeeded
            if r.is_ok() {
                // Create backup of existing wallet (if it exists)
                if wallet_path.exists() {
                    if let Err(e) = std::fs::copy(&wallet_path, &bak_path) {
                        warn!("Failed to create wallet backup: {}", e);
                        // Non-fatal: continue with the save
                    }
                }

                // Atomically replace the wallet file with the new one
                if let Err(e) = std::fs::rename(&tmp_path, &wallet_path) {
                    error!("Failed to rename temp wallet file: {}", e);
                    // Try direct copy as fallback (rename can fail across filesystems)
                    if let Err(e2) = std::fs::copy(&tmp_path, &wallet_path) {
                        let _ = std::fs::remove_file(&tmp_path);
                        return Err(format!("Failed to save wallet: {} / {}", e, e2));
                    }
                    let _ = std::fs::remove_file(&tmp_path);
                }
            }
        }

        r
    }

    pub fn get_server_uri(&self) -> http::Uri {
        self.config.server.clone()
    }

    // DEBUG: dump the wallet's latest Sapling commitment tree so it can be diffed against the
    // node's `getblockmerkletree <height>` to localize a tree/anchor divergence.
    pub fn do_sapling_tree(&self) -> JsonValue {
        match self.wallet.read().unwrap().get_sapling_tree() {
            Ok((height, hash, tree)) => object!{
                "height" => height,
                "hash"   => hash,
                "tree"   => tree,
                "anchor_offset" => self.config.anchor_offset,
            },
            Err(e) => object!{ "error" => e },
        }
    }

    pub fn do_info(&self) -> String {
        match get_info(&self.get_server_uri(), self.config.no_cert_verification) {
            Ok(i) => {
                let o = object!{
                    "version" => i.version,
                    "vendor" => i.vendor,
                    "taddr_support" => i.taddr_support,
                    "chain_name" => i.chain_name,
                    "sapling_activation_height" => i.sapling_activation_height,
                    "consensus_branch_id" => i.consensus_branch_id,
                    "latest_block_height" => i.block_height,
                    "difficulty" => i.difficulty,
                    "longestchain" => i.longestchain,
                    "notarized" => i.notarized,
                };
                o.pretty(2)
            },
            Err(e) => e
        }
    }

    pub fn do_coinsupply(&self) -> String {
        match get_coinsupply(self.get_server_uri(), self.config.no_cert_verification) {
            Ok(i) => {
                let o = object!{
                    "result" => i.result,
                    "coin" => i.coin,
                    "height" => i.height,
                    "supply" => i.supply,
                    "zfunds" => i.zfunds,
                    "total" => i.total,
                   
                };
                o.pretty(2)
            },
            Err(e) => e
        }
    }

    pub fn do_seed_phrase(&self) -> Result<JsonValue, &str> {
        if !self.wallet.read().unwrap().is_unlocked_for_spending() {
            error!("Wallet is locked");
            return Err("Wallet is locked");
        }

        let wallet = self.wallet.read().unwrap();
        Ok(object!{
            "seed"     => wallet.get_seed_phrase(),
            "birthday" => wallet.get_birthday()
        })
    }

    // Return a list of all notes, spent and unspent
    pub fn do_list_notes(&self, all_notes: bool) -> JsonValue {
        let mut unspent_notes: Vec<JsonValue> = vec![];
        let mut spent_notes  : Vec<JsonValue> = vec![];
        let mut pending_notes: Vec<JsonValue> = vec![];

        let anchor_height: i32 = self.wallet.read().unwrap().get_anchor_height() as i32;

        {
           
            let wallet = self.wallet.read().unwrap();

            // First, collect all extfvk's that are spendable (i.e., we have the private key)
            let spendable_address: HashSet<String> = wallet.get_all_zaddresses().iter()
                .filter(|address| wallet.have_spending_key_for_zaddress(address))
                .map(|address| address.clone())
                .collect();

            // Collect Sapling notes
            wallet.txs.read().unwrap().iter()
                .flat_map( |(txid, wtx)| {
                    let spendable_address = spendable_address.clone();
                    wtx.notes.iter().filter_map(move |nd| 
                        if !all_notes && nd.spent.is_some() {
                            None
                        } else {

                            let address = LightWallet::note_address(self.config.hrp_sapling_address(), nd);
                            let spendable = address.is_some() && 
                                                    spendable_address.contains(&address.clone().unwrap()) && 
                                                    wtx.block <= anchor_height && nd.spent.is_none() && nd.unconfirmed_spent.is_none();
                            Some(object!{
                                "created_in_block"   => wtx.block,
                                "datetime"           => wtx.datetime,
                                "created_in_txid"    => format!("{}", txid),
                                "value"              => nd.note.value,
                                "is_change"          => nd.is_change,
                                "address"            => address,
                                "spendable"          => spendable,
                                "spent"              => nd.spent.map(|spent_txid| format!("{}", spent_txid)),
                                "unconfirmed_spent"  => nd.unconfirmed_spent.map(|spent_txid| format!("{}", spent_txid)),
                            })
                        }
                    )
                })
                .for_each( |note| {
                    if note["spent"].is_null() && note["unconfirmed_spent"].is_null() {
                        unspent_notes.push(note);
                    } else if !note["spent"].is_null() {
                        spent_notes.push(note);
                    } else {
                        pending_notes.push(note);
                    }
                });
        }
        
        let mut unspent_utxos: Vec<JsonValue> = vec![];
        let mut spent_utxos  : Vec<JsonValue> = vec![];
        let mut pending_utxos: Vec<JsonValue> = vec![];
        
        {
            let wallet = self.wallet.read().unwrap();
            wallet.txs.read().unwrap().iter()
                .flat_map( |(txid, wtx)| {
                    wtx.utxos.iter().filter_map(move |utxo| 
                        if !all_notes && utxo.spent.is_some() {
                            None
                        } else {
                            Some(object!{
                                "created_in_block"   => wtx.block,
                                "datetime"           => wtx.datetime,
                                "created_in_txid"    => format!("{}", txid),
                                "value"              => utxo.value,
                                "scriptkey"          => hex::encode(utxo.script.clone()),
                                "is_change"          => false, // TODO: Identify notes as change if we send change to taddrs
                                "address"            => utxo.address.clone(),
                                "spent"              => utxo.spent.map(|spent_txid| format!("{}", spent_txid)),
                                "unconfirmed_spent"  => utxo.unconfirmed_spent.map(|spent_txid| format!("{}", spent_txid)),
                            })
                        }
                    )
                })
                .for_each( |utxo| {
                    if utxo["spent"].is_null() && utxo["unconfirmed_spent"].is_null() {
                        unspent_utxos.push(utxo);
                    } else if !utxo["spent"].is_null() {
                        spent_utxos.push(utxo);
                    } else {
                        pending_utxos.push(utxo);
                    }
                });
        }

        let mut res = object!{
            "unspent_notes" => unspent_notes,
            "pending_notes" => pending_notes,
            "utxos"         => unspent_utxos,
            "pending_utxos" => pending_utxos,
        };

        if all_notes {
            res["spent_notes"] = JsonValue::Array(spent_notes);
            res["spent_utxos"] = JsonValue::Array(spent_utxos);
        }

        res
    }

    pub fn do_encryption_status(&self) -> JsonValue {
        let wallet = self.wallet.read().unwrap();
        object!{
            "encrypted" => wallet.is_encrypted(),
            "locked"    => !wallet.is_unlocked_for_spending()
        }
    }

    pub fn do_list_transactions(&self) -> JsonValue {
        let wallet = self.wallet.read().unwrap();

        // Create a list of TransactionItems from wallet txns
        let mut tx_list = wallet.txs.read().unwrap().iter()
            .flat_map(| (_k, v) | {
                let mut txns: Vec<JsonValue> = vec![];

                if v.total_shielded_value_spent + v.total_transparent_value_spent > 0 {
                    // If money was spent, create a transaction. For this, we'll subtract
                    // all the change notes. TODO: Add transparent change here to subtract it also
                    let total_change: u64 = v.notes.iter()
                        .filter( |nd| nd.is_change )
                        .map( |nd| nd.note.value )
                        .sum();

                    // TODO: What happens if change is > than sent ?

                    // Collect outgoing metadata
                    let outgoing_json = v.outgoing_metadata.iter()
                        .map(|om| 
                            object!{
                                "address" => om.address.clone(),
                                "value"   => om.value,
                                "memo"    => LightWallet::memo_str(&Some(om.memo.clone())),
                        })
                        .collect::<Vec<JsonValue>>();                    

                    txns.push(object! {
                        "block_height" => v.block,
                        "datetime"     => v.datetime,
                        "txid"         => format!("{}", v.txid),
                        "amount"       => total_change as i64 
                                            - v.total_shielded_value_spent as i64 
                                            - v.total_transparent_value_spent as i64,
                        "outgoing_metadata" => outgoing_json,
                    });
                } 

                // For each sapling note that is not a change, add a Tx.
                txns.extend(v.notes.iter()
                    .filter( |nd| !nd.is_change )
                    .enumerate()
                    .map ( |(i, nd)| 
                        object! {
                            "block_height" => v.block,
                            "datetime"     => v.datetime,
                            "position"     => i,
                            "txid"         => format!("{}", v.txid),
                            "amount"       => nd.note.value as i64,
                            "address"      => LightWallet::note_address(self.config.hrp_sapling_address(), nd),
                            "memo"         => LightWallet::memo_str(&nd.memo),
                    })
                );

                // Get the total transparent received
                let total_transparent_received = v.utxos.iter().map(|u| u.value).sum::<u64>();
                if total_transparent_received > v.total_transparent_value_spent {
                    // Create an input transaction for the transparent value as well.
                    txns.push(object!{
                        "block_height" => v.block,
                        "datetime"     => v.datetime,
                        "txid"         => format!("{}", v.txid),
                        "amount"       => total_transparent_received as i64 - v.total_transparent_value_spent as i64,
                        "address"      => v.utxos.iter().map(|u| u.address.clone()).collect::<Vec<String>>().join(","),
                        "memo"         => None::<String>
                    })
                }

                txns
            })
            .collect::<Vec<JsonValue>>();
    
            // Add the incoming Mempool - incoming_mempool flag is atm useless, but we can use that in future maybe
            tx_list.extend(
                wallet.incoming_mempool_txs.read().unwrap().iter().flat_map(|(_, wtxs)| {
                    wtxs.iter().flat_map(|wtx| {
                        wtx.incoming_metadata.iter()
                            .enumerate()
                            .map(move |(i, om)| 
                                object! {
                                    "block_height" => wtx.block.clone(),
                                    "datetime"     => wtx.datetime.clone(),
                                    "position"     => i,
                                    "txid"         => format!("{}", wtx.txid),
                                    "amount"       => om.value as i64,
                                    "address"      => om.address.clone(),
                                    "memo"         => LightWallet::memo_str(&Some(om.memo.clone())),
                                    "unconfirmed"  => true,
                                    "incoming_mempool" => true,
                                }
                            )
                    })
                })
            );
  
        // Add in all mempool txns
        tx_list.extend(wallet.mempool_txs.read().unwrap().iter().map( |(_, wtx)| {
            use zcash_primitives::transaction::components::amount::DEFAULT_FEE;
            use std::convert::TryInto;
            
            let amount: u64 = wtx.outgoing_metadata.iter().map(|om| om.value).sum::<u64>();
            let fee: u64 = DEFAULT_FEE.try_into().unwrap();

            // Collect outgoing metadata
            let outgoing_json = wtx.outgoing_metadata.iter()
                .map(|om| 
                    object!{
                        "address" => om.address.clone(),
                        "value"   => om.value,
                        "memo"    => LightWallet::memo_str(&Some(om.memo.clone())),
                        
                }).collect::<Vec<JsonValue>>();                    

            object! {
                "block_height" => wtx.block,
                "datetime"     => wtx.datetime,
                "txid"         => format!("{}", wtx.txid),
                "amount"       => -1 * (fee + amount) as i64,
                "unconfirmed"  => true,
                "outgoing_metadata" => outgoing_json,
            }
        }));

        tx_list.sort_by( |a, b| if a["block_height"] == b["block_height"] {
                                    a["txid"].as_str().cmp(&b["txid"].as_str())
                                } else {
                                    a["block_height"].as_i32().cmp(&b["block_height"].as_i32())
                                }
        );

        JsonValue::Array(tx_list)
    }

    /// Create a new address, deriving it from the seed.
    pub fn do_new_address(&self, addr_type: &str) -> Result<JsonValue, String> {
        if !self.wallet.read().unwrap().is_unlocked_for_spending() {
            error!("Wallet is locked");
            return Err("Wallet is locked".to_string());
        }

        let new_address = {
            let wallet = self.wallet.write().unwrap();

            let addr = match addr_type {
                "zs" => wallet.add_zaddr(),
                "R" => wallet.add_taddr(),
                _   => {
                    let e = format!("Unrecognized address type: {}", addr_type);
                    error!("{}", e);
                    return Err(e);
                }
            };

            if addr.starts_with("Error") {
                let e = format!("Error creating new address: {}", addr);
                    error!("{}", e);
                    return Err(e);
            }

            addr
        };

        self.do_save()?;

        Ok(array![new_address])
    }

    /// Import a new private key
    pub fn do_import_tk(&self, sk: String,  birthday: u64) -> Result<JsonValue, String> {
        if !self.wallet.read().unwrap().is_unlocked_for_spending() {
            error!("Wallet is locked");
            return Err("Wallet is locked".to_string());
        }

        let new_address = {
            let mut  wallet = self.wallet.write().unwrap();

            let addr = wallet.import_taddr(sk, birthday);
            if addr.starts_with("Error") {
                let e = format!("Error creating new address{}", addr);
                    error!("{}", e);
                    return Err(e);
            }

            addr
        };

        self.do_save()?;

        Ok(array![new_address])
    }

    // Start Mempool-Monitor
pub fn start_mempool_monitor(lc: Arc<LightClient>) -> Result<(), String> {
    let config = lc.config.clone();
    let uri = config.server.clone();
    let shutdown = lc.shutdown_flag.clone();

    let (incoming_mempool_tx, incoming_mempool_rx) = std::sync::mpsc::channel::<RawTransaction>();

    // Thread for receive transactions
    let shutdown_rx = lc.shutdown_flag.clone();
    std::thread::spawn(move || {
            while let Ok(rtx) = incoming_mempool_rx.recv() {
                if shutdown_rx.load(Ordering::Relaxed) {
                    break;
                }
                if let Ok(tx) = Transaction::read(
                &rtx.data[..])
           { 
           let light_wallet_clone = lc.wallet.clone();
           light_wallet_clone.read().unwrap().scan_full_mempool_tx(&tx, rtx.height as i32, 0, true);
            }
    }
    });

    // Thread mempool monitor
    std::thread::spawn(move || {
        let mut rt = Runtime::new().unwrap();
        rt.block_on(async {
            loop {
                if shutdown.load(Ordering::Relaxed) {
                    info!("Mempool monitor shutting down");
                    break;
                }

                let incoming_mempool_tx_clone = incoming_mempool_tx.clone();
                let send_closure = move |rtx: RawTransaction| {
                    incoming_mempool_tx_clone.send(rtx).map_err(|e| Box::new(e) as Box<dyn std::error::Error>)
                };

                match grpcconnector::monitor_mempool(&uri.clone(), true, send_closure).await {
                    Ok(_) => info!("Mempool monitor loop successful"),
                    Err(e) => {
                        // If the server doesn't implement GetMempoolStream, stop the monitor instead
                        // of reconnecting forever against a method that will never succeed (log spam
                        // + needless churn). Other errors fall through to the periodic retry below.
                        if e.to_string().contains("Unimplemented") {
                            info!("Server does not implement GetMempoolStream; stopping mempool monitor");
                            return;
                        }
                        warn!("Mempool monitor returned {:?}, will restart listening", e);
                    }
                }

                // Sleep in 1-second increments so we can check shutdown flag
                for _ in 0..10 {
                    if shutdown.load(Ordering::Relaxed) {
                        info!("Mempool monitor shutting down during sleep");
                        return;
                    }
                    std::thread::sleep(Duration::from_secs(1));
                }
            }
        });
    });

    Ok(())
}

    /// Signal all background threads to stop
    pub fn shutdown(&self) {
        self.shutdown_flag.store(true, Ordering::Relaxed);
        info!("Shutdown flag set");
    }
    /// Convinence function to determine what type of key this is and import it
    pub fn do_import_key(&self, key: String, birthday: u64) -> Result<JsonValue, String> {
        if key.starts_with(self.config.hrp_sapling_private_key()) {
            self.do_import_sk(key, birthday)
        } else if key.starts_with(self.config.hrp_sapling_viewing_key()) {
            self.do_import_vk(key, birthday)
        } else {
            Err(format!("'{}' was not recognized as either a spending key or a viewing key because it didn't start with either '{}' or '{}'", 
                key, self.config.hrp_sapling_private_key(), self.config.hrp_sapling_viewing_key()))
        }
    }

    /// Import a new private key
    pub fn do_import_sk(&self, sk: String, birthday: u64) -> Result<JsonValue, String> {
        if !self.wallet.read().unwrap().is_unlocked_for_spending() {
            error!("Wallet is locked");
            return Err("Wallet is locked".to_string());
        }

        let new_address = {
            let mut wallet = self.wallet.write().unwrap();

            let addr = wallet.add_imported_sk(sk, birthday);
            if addr.starts_with("Error") {
                let e = format!("Error creating new address{}", addr);
                    error!("{}", e);
                    return Err(e);
            }

            addr
        };

        self.do_save()?;

        Ok(array![new_address])
    }

    /// Import a new viewing key
    pub fn do_import_vk(&self, vk: String, birthday: u64) -> Result<JsonValue, String> {
        if !self.wallet.read().unwrap().is_unlocked_for_spending() {
            error!("Wallet is locked");
            return Err("Wallet is locked".to_string());
        }

        let new_address = {
            let mut wallet = self.wallet.write().unwrap();

            let addr = wallet.add_imported_vk(vk, birthday);
            if addr.starts_with("Error") {
                let e = format!("Error creating new address{}", addr);
                    error!("{}", e);
                    return Err(e);
            }
            addr
        };

        self.do_save()?;

        Ok(array![new_address])
    }

    pub fn do_new_sietchaddress(&self, addr_type: &str) -> Result<JsonValue, String> {
       
        let zdust_address = {
            let wallet = self.wallet.write().unwrap();

            match addr_type {
                "zs" => wallet.add_zaddrdust(),

                _   => {
                    let e = format!("Unrecognized address type: {}", addr_type);
                    error!("{}", e);
                    return Err(e);
                }
            }
        };

        Ok(array![zdust_address])
    }

    pub fn clear_state(&self) {
        self.clear_state_from(self.wallet.read().unwrap().get_birthday());
    }

    pub fn clear_state_from(&self, height: u64) {
        // First, clear the state from the wallet
        self.wallet.read().unwrap().clear_blocks();

        // Then set the initial block
        self.set_wallet_initial_state(height);
        info!("Cleared wallet state to height {}", height);
    }

    pub fn do_rescan(&self) -> Result<JsonValue, String> {
        if !self.wallet.read().unwrap().is_unlocked_for_spending() {
            warn!("Wallet is locked, new HD addresses won't be added!");
        }
        
        info!("Rescan starting");
        
        self.clear_state();

        // Then, do a sync, which will force a full rescan from the initial state
        let response = self.do_sync(true);

        self.do_save()?;
        info!("Rescan finished");

        response
    }

    /// Return the syncing status of the wallet
    pub fn do_scan_status(&self) -> WalletStatus {
        self.sync_status.read().unwrap().clone()
    }

    pub fn do_sync(&self, print_updates: bool) -> Result<JsonValue, String> {

        let mut retry_count = 0;
        loop {
            match self.do_sync_internal(print_updates, retry_count) {
                Ok(j) => return Ok(j),
                Err(e) => {
                    retry_count += 1;
                    if retry_count > 5 {
                        return Err(e);
                    }
                    // Sleep exponentially backing off
                    std::thread::sleep(std::time::Duration::from_secs((2 as u64).pow(retry_count)));

                }
            }
        }
    }

    pub fn do_shield(&self, address: Option<String>) -> Result<String, String> {
        use zcash_primitives::transaction::components::amount::DEFAULT_FEE;
        use std::convert::TryInto;

        let fee = DEFAULT_FEE.try_into().unwrap();
        let tbal = self.wallet.read().unwrap().tbalance(None);

        // Make sure there is a balance, and it is greated than the amount
        if tbal <= fee {
            return Err(format!("Not enough transparent balance to shield. Have {} puposhis, need more than {} puposhis to cover tx fee", tbal, fee));
        }

        let addr = address.or(self.wallet.read().unwrap().get_all_zaddresses().get(0).map(|s| s.clone())).unwrap();
        let fee: u64 = DEFAULT_FEE.try_into().unwrap();

        let result = {
            let _lock = self.sync_lock.lock().unwrap();
            self.wallet.read().unwrap().send_to_address(
                u32::from_str_radix(&self.config.consensus_branch_id, 16).unwrap(), 
                &self.sapling_spend, &self.sapling_output, 
                true, 
                vec![(&addr, tbal - fee, None)],
                &fee,
                |txbytes| broadcast_raw_tx(&self.get_server_uri(),self.config.no_cert_verification, txbytes)
            )
        };

        result.map(|(txid, _)| txid)
    }

    fn do_sync_internal(&self, print_updates: bool, retry_count: u32) -> Result<JsonValue, String> {
        // We can only do one sync at a time because we sync blocks in serial order
        // If we allow multiple syncs, they'll all get jumbled up.
        let _lock = self.sync_lock.lock().unwrap();

        // Sync is 3 parts
        // 1. Get the latest block
        // 2. Get all the blocks that we don't have
        // 3. Find all new Txns that don't have the full Tx, and get them as full transactions 
        //    and scan them, mainly to get the memos
        let mut last_scanned_height = self.wallet.read().unwrap().last_scanned_height() as u64;

        // This will hold the latest block fetched from the RPC

        let latest_block = fetch_latest_block(&self.get_server_uri(), self.config.no_cert_verification)?.height;
       
        if latest_block < last_scanned_height {
            let w = format!("Server's latest block({}) is behind ours({})", latest_block, last_scanned_height);
            warn!("{}", w);
            return Err(w);
        }

        info!("Latest block is {}", latest_block);

        // Get the end height to scan to.
        // get_block_range streams an ENTIRE range over ONE gRPC connection, so the batch size is
        // the dominant factor in how many TCP connections a sync opens. The old 1000 meant ~1300
        // connections for a full sync, which on Windows exhausted socket buffers (WSAENOBUFS / os
        // error 10055). A large batch (25k) cuts that ~50x (~52 connections for a full sync) with
        // no extra memory — blocks are still streamed + scanned incrementally, not buffered.
        let scan_batch_size = 25000;
        let mut end_height = std::cmp::min(last_scanned_height + scan_batch_size, latest_block);

        // If there's nothing to scan, just return
        if last_scanned_height == latest_block {
            info!("Nothing to sync, returning");
            return Ok(object!{ "result" => "success" })
        }

        {
            let mut status = self.sync_status.write().unwrap();
            status.is_syncing = true;
            status.synced_blocks = last_scanned_height;
            status.total_blocks = latest_block;
        }

        // Count how many bytes we've downloaded
        let bytes_downloaded = Arc::new(AtomicUsize::new(0));

        let mut total_reorg = 0;

        // Collect all txns in blocks that we have a tx in. We'll fetch all these
        // txs along with our own, so that the server doesn't learn which ones
        // belong to us.
        let all_new_txs = Arc::new(RwLock::new(vec![]));

        // Create a new threadpool (upto 8, atleast 2 threads) to scan with
        let pool = ThreadPool::new(max(2, min(8, num_cpus::get())));

        // Create a single Tokio runtime for all block fetches to avoid per-batch overhead
        let mut rt = tokio::runtime::Runtime::new()
            .map_err(|e| format!("Error creating runtime: {:?}", e))?;

        // Fetch CompactBlocks in increments
        let mut pass = 0;
        loop {
            pass +=1 ;
            // Collect all block times, because we'll need to update transparent tx
            // datetime via the block height timestamp
            let block_times = Arc::new(RwLock::new(HashMap::new()));

            let local_light_wallet = self.wallet.clone();
            let local_bytes_downloaded = bytes_downloaded.clone();

            let start_height = last_scanned_height + 1;
            info!("Start height is {}", start_height);

            // Show updates only if we're syncing a lot of blocks
            if print_updates && (latest_block - start_height) > 100 {
                print!("Syncing {}/{}\r", start_height, latest_block);
                io::stdout().flush().ok().expect("Could not flush stdout");
            }

            {
                let mut status = self.sync_status.write().unwrap();
                status.is_syncing = true;
                status.synced_blocks = start_height;
                status.total_blocks = latest_block;
            }

            // Fetch compact blocks
            info!("Fetching blocks {}-{}", start_height, end_height);

            let all_txs = all_new_txs.clone();
            let block_times_inner = block_times.clone();

            let last_invalid_height = Arc::new(AtomicI32::new(0));
            let last_invalid_height_inner = last_invalid_height.clone();
            let tpool = pool.clone();
            fetch_blocks_with_runtime(&mut rt, &self.get_server_uri(), start_height, end_height, self.config.no_cert_verification, pool.clone(),
                move |encoded_block: &[u8], height: u64| {
                    // Process the block only if there were no previous errors
                    if last_invalid_height_inner.load(Ordering::SeqCst) > 0 {
                        return;
                    }

                    // Parse the block and save it's time. We'll use this timestamp for 
                    // transactions in this block that might belong to us.
                    let block: Result<zcash_client_backend::proto::compact_formats::CompactBlock, _>
                                        = parse_from_bytes(encoded_block);
                    match block {
                        Ok(b) => {
                            block_times_inner.write().unwrap().insert(b.height, b.time);
                        },
                        Err(_) => {}
                    }

                    match local_light_wallet.read().unwrap().scan_block_with_pool(encoded_block, &tpool) {
                        Ok(block_txns) => {
                            // Add to global tx list
                            all_txs.write().unwrap().extend_from_slice(&block_txns.iter().map(|txid| (txid.clone(), height as i32)).collect::<Vec<_>>()[..]);
                        },
                        Err(invalid_height) => {
                            // Block at this height seems to be invalid, so invalidate up till that point
                            last_invalid_height_inner.store(invalid_height, Ordering::SeqCst);
                        }
                    };

                    local_bytes_downloaded.fetch_add(encoded_block.len(), Ordering::SeqCst);
            })?;


            {
               

                let t = self.wallet.read().unwrap();
                let mut d = t.total_scan_duration.write().unwrap();
                d.clear();
                d.push(std::time::Duration::new(0, 0));
            }

            // Check if there was any invalid block, which means we might have to do a reorg
            let invalid_height = last_invalid_height.load(Ordering::SeqCst);
            if invalid_height > 0 {
                total_reorg += self.wallet.read().unwrap().invalidate_block(invalid_height);

                warn!("Invalidated block at height {}. Total reorg is now {}", invalid_height, total_reorg);
            }

            // Make sure we're not re-orging too much!
            if total_reorg > (crate::lightwallet::MAX_REORG - 1) as u64 {
                error!("Reorg has now exceeded {} blocks!", crate::lightwallet::MAX_REORG);
                return Err(format!("Reorg has exceeded {} blocks. Aborting.", crate::lightwallet::MAX_REORG));
            } 
            
            if invalid_height > 0 {
                // Reset the scanning heights
                last_scanned_height = (invalid_height - 1) as u64;
                end_height = std::cmp::min(last_scanned_height + scan_batch_size, latest_block);

                warn!("Reorg: reset scanning from {} to {}", last_scanned_height, end_height);

                continue;
            }

            // If it got here, that means the blocks are scanning properly now. 
            // So, reset the total_reorg
            total_reorg = 0;

            // If this is the first pass after a retry, fetch older t address txids too, becuse
            // they might have been missed last time.
            let transparent_start_height = if pass == 1 && retry_count > 0 {
                start_height - scan_batch_size
            } else {
                start_height
            };

            let no_cert = true;

                      // We'll also fetch all the txids that our transparent addresses are involved with
                      self.scan_taddress_txids(&pool, block_times, transparent_start_height, end_height, no_cert)?;          
            
            // Do block height accounting
            last_scanned_height = end_height;
            end_height = last_scanned_height + scan_batch_size;

            if last_scanned_height >= latest_block {
                break;
            } else if end_height > latest_block {
                end_height = latest_block;
            }
        }

        if print_updates{
            println!(""); // New line to finish up the updates
        }
        
        info!("Synced to {}, Downloaded {} kB", latest_block, bytes_downloaded.load(Ordering::SeqCst) / 1024);
        {
            let mut status = self.sync_status.write().unwrap();
            status.is_syncing = false;
            status.synced_blocks = latest_block;
            status.total_blocks = latest_block;
        }

        // Get the Raw transaction for all the wallet transactions

        {
            let decoy_txids = all_new_txs.read().unwrap();
            match self.scan_fill_fulltxs(&pool, decoy_txids.to_vec()) {
                Ok(_) => Ok(object!{
                    "result" => "success",
                    "latest_block" => latest_block,
                    "downloaded_bytes" => bytes_downloaded.load(Ordering::SeqCst)
                }),
                Err(e) => Err(format!("Error fetching all txns for memos: {}", e))
            }        
        }
    }
    fn scan_taddress_txids(
        &self, 
        pool: &ThreadPool, 
        block_times: Arc<RwLock<HashMap<u64, u32>>>, 
        start_height: u64, 
        end_height: u64, 
        no_cert: bool
    ) -> Result<Vec<()>, String> {
        let addresses = self.wallet.read()
            .map_err(|e| format!("Failed to read wallet: {:?}", e))?
            .get_all_taddresses().iter()
            .cloned()
            .collect::<Vec<String>>();
    
        let (ctx, crx) = channel();
        let num_addresses = addresses.len();
    
        for address in addresses {
            let address_clone = address.clone();
            let wallet = self.wallet.clone();
            let pool = pool.clone();
            let server_uri = self.get_server_uri();
            let ctx = ctx.clone();
            let block_times = block_times.clone();
    
            pool.execute(move || {
              
                let r = fetch_transparent_txids(
                    &server_uri, 
                    address, 
                    start_height, 
                    end_height, 
                    no_cert,
                    move |tx_bytes: &[u8], height: u64| {
                        let tx_result = Transaction::read(tx_bytes)
                            .map_err(|e| format!("Failed to read transaction: {:?}", e));
    
                        match tx_result {
                            Ok(tx) => {
                                let datetime_result = block_times.read()
                                    .map_err(|e| format!("Failed to read block times: {:?}", e))
                                    .and_then(|bt| bt.get(&height).cloned().ok_or_else(|| format!("No datetime for height: {}", height)));
    
                                match datetime_result {
                                    Ok(datetime) => {
                                        match wallet.read().map_err(|e| format!("Failed to read wallet: {:?}", e)) {
                                            Ok(w) => {
                                                w.scan_full_tx(&tx, height as i32, datetime as u64);
                                            },
                                            Err(e) => {
                                                error!("Error reading wallet: {}", e);
                                            },
                                        }
                                    },
                                    Err(e) => {
                                        error!("Error processing transaction: {}", e);
                                    },
                                }
                            },
                            Err(e) => {
                                error!("Error reading transaction: {}", e);
                            },
                        }
                    }
                );
    
                match ctx.send(r) {
                    Ok(_) => info!("Successfully sent data for address: {}", address_clone),
                    Err(e) => error!("Failed to send data for address: {}: {:?}", address_clone, e),
                }
            });
        }
    
        crx.iter().take(num_addresses).collect()
    }
    
 
    fn scan_fill_fulltxs(&self, pool: &ThreadPool, decoy_txids: Vec<(TxId, i32)>) -> Result<Vec<()>, String> {

        // We need to first copy over the Txids from the wallet struct, because
        // we need to free the read lock from here (Because we'll self.wallet.txs later)
        let mut txids_to_fetch: Vec<(TxId, i32)> = self.wallet.read().unwrap().txs.read().unwrap().values()
                                                        .filter(|wtx| wtx.full_tx_scanned == false)
                                                        .map(|wtx| (wtx.txid.clone(), wtx.block))
                                                        .collect::<Vec<(TxId, i32)>>();

        info!("Fetching {} new txids, total {} with decoy", txids_to_fetch.len(), decoy_txids.len());
        txids_to_fetch.extend_from_slice(&decoy_txids[..]);
        txids_to_fetch.sort();
        txids_to_fetch.dedup();

        let num_fetches = txids_to_fetch.len();
        let (ctx, crx) = channel();

        // And go and fetch the txids, getting the full transaction, so we can 
        // read the memos
        for (txid, height) in txids_to_fetch {
            let light_wallet_clone = self.wallet.clone();

                let pool = pool.clone();
                let server_uri = self.get_server_uri();
                let ctx = ctx.clone();
                let no_cert =  self.config.no_cert_verification;
    
                pool.execute(move || {
                    //info!("Fetching full Tx: {}", txid);

                    match fetch_full_tx(&server_uri, txid, no_cert) {
                        Ok(tx_bytes) => {
                            let tx = Transaction::read(&tx_bytes[..]).unwrap();
    
                            light_wallet_clone.read().unwrap().scan_full_tx(&tx, height, 0);
                            ctx.send(Ok(())).unwrap();
                        },
                        Err(e) => ctx.send(Err(e)).unwrap()
                    };   
            });
        };

            // Wait for all the fetches to finish.
            crx.iter().take(num_fetches).collect::<Result<Vec<()>, String>>()  
    }

    pub fn do_send(&self, addrs: Vec<(&str, u64, Option<String>)>, fee: &u64) -> Result<String, String> {
        if !self.wallet.read().unwrap().is_unlocked_for_spending() {
            error!("Wallet is locked");
            return Err("Wallet is locked".to_string());
        }

        info!("Creating transaction");

        let result = {
            let _lock = self.sync_lock.lock().unwrap();

            self.wallet.write().unwrap().send_to_address(
                u32::from_str_radix(&self.config.consensus_branch_id, 16).unwrap(), 
                &self.sapling_spend, &self.sapling_output,
                false,
                addrs,
                fee,
                |txbytes| broadcast_raw_tx(&self.get_server_uri(), self.config.no_cert_verification, txbytes)
            )
        };
        
        result.map(|(txid, _)| txid)
    }
}

#[cfg(test)]
pub mod tests {
    use lazy_static::lazy_static;
    use tempdir::TempDir;
    use super::{LightClient, LightClientConfig};

    lazy_static!{
        static ref TEST_SEED: String = "youth strong sweet gorilla hammer unhappy congress stamp left stereo riot salute road tag clean toilet artefact fork certain leopard entire civil degree wonder".to_string();
    }

    #[test]
    pub fn test_encrypt_decrypt() {
        let lc = super::LightClient::unconnected(TEST_SEED.to_string(), None).unwrap();

        assert!(!lc.do_export(None).is_err());
        assert!(!lc.do_new_address("zs").is_err());
        assert!(!lc.do_new_address("R").is_err());
        assert_eq!(lc.do_seed_phrase().unwrap()["seed"], TEST_SEED.to_string());

        // Encrypt and Lock the wallet
        lc.wallet.write().unwrap().encrypt("password".to_string()).unwrap();
        assert!(lc.do_export(None).is_err());
        assert!(lc.do_seed_phrase().is_err());
        assert!(lc.do_new_address("R").is_err());
        assert!(lc.do_new_address("zs").is_err());
        assert!(lc.do_send(vec![("zs", 0, None)]).is_err());

        // Do a unlock, and make sure it all works now
        lc.wallet.write().unwrap().unlock("password".to_string()).unwrap();
        assert!(!lc.do_export(None).is_err());
        assert!(!lc.do_seed_phrase().is_err());

        // Can't add keys while unlocked but encrypted
        assert!(lc.do_new_address("R").is_err());
        assert!(lc.do_new_address("zs").is_err());

        // Remove encryption, which will allow adding 
        lc.wallet.write().unwrap().remove_encryption("password".to_string()).unwrap();
        
        assert!(lc.do_new_address("R").is_ok());
        assert!(lc.do_new_address("zs").is_ok());
    }

    #[test]
    pub fn test_addresses() {
        let lc = super::LightClient::unconnected(TEST_SEED.to_string(), None).unwrap();

        {
            let addresses = lc.do_address();
            // When restoring from seed, there should be 5+1 addresses
            assert_eq!(addresses["z_addresses"].len(), 51);
            assert_eq!(addresses["t_addresses"].len(), 6);
        }

        // Add new z and t addresses
            
        let taddr1 = lc.do_new_address("R").unwrap()[0].as_str().unwrap().to_string();
        let taddr2 = lc.do_new_address("R").unwrap()[0].as_str().unwrap().to_string();        
        let zaddr1 = lc.do_new_address("zs").unwrap()[0].as_str().unwrap().to_string();
        let zaddr2 = lc.do_new_address("zs").unwrap()[0].as_str().unwrap().to_string();
        
        let addresses = lc.do_address();
        
        assert_eq!(addresses["z_addresses"].len(), 52);
        assert_eq!(addresses["z_addresses"][49], zaddr1);
        assert_eq!(addresses["z_addresses"][50], zaddr2);

        assert_eq!(addresses["t_addresses"].len(), 8);
        assert_eq!(addresses["t_addresses"][6], taddr1);
        assert_eq!(addresses["t_addresses"][7], taddr2);

        use std::sync::{Arc, RwLock, Mutex};
        use crate::lightclient::{WalletStatus, LightWallet};

        // When creating a new wallet, there is only 1 address
        let config = LightClientConfig::create_unconnected("test".to_string(), None);
        let lc = LightClient {
            wallet          : Arc::new(RwLock::new(LightWallet::new(None, &config, 0).unwrap())),
            config          : config,
            sapling_output  : vec![], 
            sapling_spend   : vec![],
            sync_lock       : Mutex::new(()),
            sync_status     : Arc::new(RwLock::new(WalletStatus::new())),
            shutdown_flag   : Arc::new(AtomicBool::new(false)),
        };
        {
            let addresses = lc.do_address();
            // New wallets have only 1 address
            assert_eq!(addresses["z_addresses"].len(), 1);
            assert_eq!(addresses["t_addresses"].len(), 1);
    }
}

#[test]
pub fn test_bad_import() {
    let lc = super::LightClient::unconnected(TEST_SEED.to_string(), None).unwrap();

    assert!(lc.do_import_sk("bad_priv_key".to_string(), 0).is_err());
    assert!(lc.do_import_vk("bad_view_key".to_string(), 0).is_err());
}

    #[test]
    pub fn test_wallet_creation() {
        // Create a new tmp director
        {
            let tmp = TempDir::new("lctest").unwrap();
            let dir_name = tmp.path().to_str().map(|s| s.to_string());

            // A lightclient to a new, empty directory works.
            let config = LightClientConfig::create_unconnected("test".to_string(), dir_name);
            let lc = LightClient::new(&config, 0).unwrap();
            let seed = lc.do_seed_phrase().unwrap()["seed"].as_str().unwrap().to_string();
            lc.do_save().unwrap();

            // Doing another new will fail, because the wallet file now already exists
            assert!(LightClient::new(&config, 0).is_err());

            // new_from_phrase will not work either, again, because wallet file exists
            assert!(LightClient::new_from_phrase(TEST_SEED.to_string(), &config, 0, false).is_err());

            // Creating a lightclient to the same dir without a seed should re-read the same wallet
            // file and therefore the same seed phrase
            let lc2 = LightClient::read_from_disk(&config).unwrap();
            assert_eq!(seed, lc2.do_seed_phrase().unwrap()["seed"].as_str().unwrap().to_string());
        }

        // Now, get a new directory, and try to read from phrase
        {
            let tmp = TempDir::new("lctest").unwrap();
            let dir_name = tmp.path().to_str().map(|s| s.to_string());

            let config = LightClientConfig::create_unconnected("test".to_string(), dir_name);

            // read_from_disk will fail, because the dir doesn't exist
            assert!(LightClient::read_from_disk(&config).is_err());

            // New from phrase should work becase a file doesn't exist already
            let lc = LightClient::new_from_phrase(TEST_SEED.to_string(), &config, 0, false).unwrap();
            assert_eq!(TEST_SEED.to_string(), lc.do_seed_phrase().unwrap()["seed"].as_str().unwrap().to_string());
            lc.do_save().unwrap();

            // Now a new will fail because wallet exists
            assert!(LightClient::new(&config, 0).is_err());
        }
    }

    #[test]
    pub fn test_recover_seed() {
        // Create a new tmp director
        {
            let tmp = TempDir::new("lctest").unwrap();
            let dir_name = tmp.path().to_str().map(|s| s.to_string());

            // A lightclient to a new, empty directory works.
            let config = LightClientConfig::create_unconnected("test".to_string(), dir_name);
            let lc = LightClient::new(&config, 0).unwrap();
            let seed = lc.do_seed_phrase().unwrap()["seed"].as_str().unwrap().to_string();
            lc.do_save().unwrap();

            assert_eq!(seed, LightClient::attempt_recover_seed(&config, None).unwrap());

            // Now encrypt and save the file
            let pwd = "password".to_string();
            lc.wallet.write().unwrap().encrypt(pwd.clone()).unwrap();

            assert_eq!(seed, LightClient::attempt_recover_seed(&config, Some(pwd)).unwrap());
        }
    }

}
