use alloc::collections::{BTreeMap, BTreeSet};
use alloc::string::{String, ToString};
use alloc::vec::Vec;

use miden_client::Word;
use miden_client::block::BlockHeader;
use miden_client::crypto::{Forest, InOrderIndex, MmrPeaks};
use miden_client::note::BlockNumber;
use miden_client::store::{BlockRelevance, PartialBlockchainFilter, StoreError};
use miden_client::utils::Deserializable;

use super::IdxdbStore;
use crate::promise::{await_js, await_js_value, await_ok};

mod js_bindings;
use js_bindings::{
    idxdb_get_block_headers,
    idxdb_get_current_blockchain_peaks,
    idxdb_get_partial_blockchain_nodes,
    idxdb_get_partial_blockchain_nodes_all,
    idxdb_get_partial_blockchain_nodes_up_to_inorder_index,
    idxdb_get_tracked_block_header_numbers,
    idxdb_get_tracked_block_headers,
    idxdb_insert_block_header,
    idxdb_insert_partial_blockchain_nodes,
    idxdb_prune_irrelevant_blocks,
};

mod models;
use models::{
    BlockHeaderIdxdbObject,
    PartialBlockchainNodeIdxdbObject,
    PartialBlockchainPeaksIdxdbObject,
};

pub mod utils;
use utils::{
    SerializedBlockHeaderData,
    SerializedPartialBlockchainNodeData,
    process_partial_blockchain_nodes_from_js_value,
    serialize_block_header,
    serialize_partial_blockchain_node,
};

impl IdxdbStore {
    pub(crate) async fn insert_block_header(
        &self,
        block_header: &BlockHeader,
        has_client_notes: bool,
    ) -> Result<(), StoreError> {
        let SerializedBlockHeaderData { block_num, header, has_client_notes } =
            serialize_block_header(block_header, has_client_notes);

        let promise = idxdb_insert_block_header(self.db_id(), block_num, header, has_client_notes);
        await_ok(promise, "failed to insert block header").await?;

        Ok(())
    }

    pub(crate) async fn get_block_headers(
        &self,
        block_numbers: &BTreeSet<BlockNumber>,
    ) -> Result<Vec<(BlockHeader, BlockRelevance)>, StoreError> {
        let formatted_block_numbers_list: Vec<u32> =
            block_numbers.iter().map(BlockNumber::as_u32).collect();

        let promise = idxdb_get_block_headers(self.db_id(), formatted_block_numbers_list);
        let block_headers_idxdb: Vec<Option<BlockHeaderIdxdbObject>> =
            await_js(promise, "failed to get block headers").await?;

        // Transform the list of Option<BlockHeaderIdxdbObject> to a list of results
        let results: Result<Vec<(BlockHeader, BlockRelevance)>, StoreError> = block_headers_idxdb
            .into_iter()
            .filter_map(|record_option| record_option.map(Ok))
            .map(|record_result: Result<BlockHeaderIdxdbObject, StoreError>| {
                let record = record_result?;
                let block_header = BlockHeader::read_from_bytes(&record.header)?;
                let has_client_notes = record.has_client_notes.into();

                Ok((block_header, has_client_notes))
            })
            .collect(); // Collects into Result<Vec<(BlockHeader, bool)>, StoreError>

        results
    }

    pub(crate) async fn get_tracked_block_headers(&self) -> Result<Vec<BlockHeader>, StoreError> {
        let promise = idxdb_get_tracked_block_headers(self.db_id());
        let block_headers_idxdb: Vec<BlockHeaderIdxdbObject> =
            await_js(promise, "failed to get tracked block headers").await?;

        let results: Result<Vec<BlockHeader>, StoreError> = block_headers_idxdb
            .into_iter()
            .map(|record| {
                let block_header = BlockHeader::read_from_bytes(&record.header)?;

                Ok(block_header)
            })
            .collect();

        results
    }

    pub(crate) async fn get_tracked_block_header_numbers(
        &self,
    ) -> Result<BTreeSet<usize>, StoreError> {
        let promise = idxdb_get_tracked_block_header_numbers(self.db_id());
        let block_nums: Vec<u32> =
            await_js(promise, "failed to get tracked block header numbers").await?;

        Ok(block_nums.into_iter().map(|n| n as usize).collect())
    }

    pub(crate) async fn get_partial_blockchain_nodes(
        &self,
        filter: PartialBlockchainFilter,
    ) -> Result<BTreeMap<InOrderIndex, Word>, StoreError> {
        match filter {
            PartialBlockchainFilter::All => {
                let promise = idxdb_get_partial_blockchain_nodes_all(self.db_id());
                let js_value =
                    await_js_value(promise, "failed to get all partial blockchain nodes").await?;
                process_partial_blockchain_nodes_from_js_value(js_value)
            },
            PartialBlockchainFilter::List(ids) => {
                let formatted_list: Vec<String> =
                    ids.iter().map(|id| (Into::<usize>::into(*id)).to_string()).collect();

                let promise = idxdb_get_partial_blockchain_nodes(self.db_id(), formatted_list);
                let js_value =
                    await_js_value(promise, "failed to get partial blockchain nodes").await?;
                let nodes = process_partial_blockchain_nodes_from_js_value(js_value)?;

                // Verify that all requested nodes were found. Missing nodes indicate
                // that MMR authentication nodes were not persisted during a previous
                // sync (e.g. the browser extension was closed mid-sync).
                for id in &ids {
                    if !nodes.contains_key(id) {
                        return Err(StoreError::PartialBlockchainNodeNotFound(id.inner() as u64));
                    }
                }

                Ok(nodes)
            },
            PartialBlockchainFilter::Forest(forest) => {
                if forest.is_empty() {
                    return Ok(BTreeMap::new());
                }

                let max_in_order_index = forest.rightmost_in_order_index().inner().to_string();
                let promise = idxdb_get_partial_blockchain_nodes_up_to_inorder_index(
                    self.db_id(),
                    max_in_order_index,
                );
                let js_value =
                    await_js_value(promise, "failed to get partial blockchain nodes up to index")
                        .await?;
                process_partial_blockchain_nodes_from_js_value(js_value)
            },
        }
    }

    pub(crate) async fn get_current_blockchain_peaks(&self) -> Result<MmrPeaks, StoreError> {
        let promise = idxdb_get_current_blockchain_peaks(self.db_id());
        let peaks_idxdb: PartialBlockchainPeaksIdxdbObject =
            await_js(promise, "failed to get current blockchain peaks").await?;

        let PartialBlockchainPeaksIdxdbObject { block_num, peaks } = peaks_idxdb;

        if let Some(peaks) = peaks {
            let mmr_peaks_nodes: Vec<Word> = Vec::<Word>::read_from_bytes(&peaks)?;

            return MmrPeaks::new(Forest::new(block_num as usize)?, mmr_peaks_nodes)
                .map_err(StoreError::MmrError);
        }

        Ok(MmrPeaks::new(Forest::empty(), vec![])?)
    }

    pub(crate) async fn insert_partial_blockchain_nodes(
        &self,
        nodes: &[(InOrderIndex, Word)],
    ) -> Result<(), StoreError> {
        let mut serialized_node_ids = Vec::new();
        let mut serialized_nodes = Vec::new();
        for (id, node) in nodes {
            let SerializedPartialBlockchainNodeData { id, node } =
                serialize_partial_blockchain_node(*id, *node)?;
            serialized_node_ids.push(id);
            serialized_nodes.push(node);
        }

        let promise = idxdb_insert_partial_blockchain_nodes(
            self.db_id(),
            serialized_node_ids,
            serialized_nodes,
        );
        await_ok(promise, "failed to insert partial blockchain nodes").await?;

        Ok(())
    }

    pub(crate) async fn prune_irrelevant_blocks(
        &self,
        blocks_to_untrack: &[BlockNumber],
        node_indices_to_remove: &[InOrderIndex],
    ) -> Result<(), StoreError> {
        let block_nums_vec: Vec<u32> = blocks_to_untrack.iter().map(BlockNumber::as_u32).collect();
        let node_ids_vec: Vec<String> = node_indices_to_remove
            .iter()
            .map(|id| (Into::<usize>::into(*id)).to_string())
            .collect();

        let promise = idxdb_prune_irrelevant_blocks(self.db_id(), block_nums_vec, node_ids_vec);
        await_ok(promise, "failed to prune irrelevant blocks").await?;

        Ok(())
    }
}
