// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! This module implements record handling for a simple asynchronized mysql database

use std::convert::TryInto;

use akd::storage::types::{DbRecord, StorageType};
use akd::storage::Storable;
use akd::tree_node::{NodeKey, TreeNodeWithPreviousValue};
use akd::NodeLabel;
use mysql_async::prelude::*;
use mysql_async::*;

type MySqlError = mysql_async::Error;

pub(crate) const TABLE_AZKS: &str = "azks";
pub(crate) const TABLE_HISTORY_TREE_NODES: &str = "history";
pub(crate) const TABLE_USER: &str = "users";
pub(crate) const TEMP_IDS_TABLE: &str = "temp_ids_table";

const SELECT_AZKS_DATA: &str = "`epoch`, `num_nodes`";
const SELECT_HISTORY_TREE_NODE_DATA: &str =
    "`label_len`, `label_val`, `last_epoch`, `least_descendant_ep`, `parent_label_len`, `parent_label_val`, `node_type`, `left_child_len`, `left_child_label_val`, `right_child_len`, `right_child_label_val`, `hash`, `p_last_epoch`, `p_least_descendant_ep`, `p_parent_label_len`, `p_parent_label_val`, `p_node_type`, `p_left_child_len`, `p_left_child_label_val`, `p_right_child_len`, `p_right_child_label_val`, `p_hash`";
const SELECT_USER_DATA: &str =
    "`username`, `epoch`, `version`, `node_label_val`, `node_label_len`, `data`";

pub(crate) trait MySqlStorable {
    fn set_statement(&self) -> String;

    fn set_params(&self) -> Option<mysql_async::Params>;

    fn set_batch_statement<St: Storable>(items: usize) -> String;

    fn set_batch_params(items: &[DbRecord]) -> Result<mysql_async::Params>;

    fn get_statement<St: Storable>() -> String;

    fn get_batch_create_temp_table<St: Storable>() -> Option<String>;

    fn get_batch_fill_temp_table<St: Storable>(num_items: Option<usize>) -> String;

    fn get_batch_statement<St: Storable>() -> String;

    fn get_specific_statement<St: Storable>() -> String;

    fn get_specific_params<St: Storable>(key: &St::StorageKey) -> Option<mysql_async::Params>;

    fn get_multi_row_specific_params<St: Storable>(
        keys: &[St::StorageKey],
    ) -> Option<mysql_async::Params>;

    fn from_row<St: Storable>(row: &mut mysql_async::Row) -> core::result::Result<Self, MySqlError>
    where
        Self: std::marker::Sized;
}

impl MySqlStorable for DbRecord {
    fn set_statement(&self) -> String {
        match &self {
            DbRecord::Azks(_) => format!("INSERT INTO `{TABLE_AZKS}` (`key`, {SELECT_AZKS_DATA})
            VALUES (:key, :epoch, :num_nodes)
            ON DUPLICATE KEY UPDATE
                `epoch` = :epoch
                , `num_nodes` = :num_nodes"),
            DbRecord::TreeNode(_) => format!("INSERT INTO `{TABLE_HISTORY_TREE_NODES}` ({SELECT_HISTORY_TREE_NODE_DATA})
            VALUES (:label_len
                , :label_val
                , :last_epoch
                , :least_descendant_ep
                , :parent_label_len
                , :parent_label_val
                , :node_type
                , :left_child_len
                , :left_child_label_val
                , :right_child_len
                , :right_child_label_val
                , :hash
                , :p_last_epoch
                , :p_least_descendant_ep
                , :p_parent_label_len
                , :p_parent_label_val
                , :p_node_type
                , :p_left_child_len
                , :p_left_child_label_val
                , :p_right_child_len
                , :p_right_child_label_val
                , :p_hash)
            ON DUPLICATE KEY UPDATE
                `label_len` = :label_len
                , `label_val` = :label_val
                , `last_epoch` = :last_epoch
                , `least_descendant_ep` = :least_descendant_ep
                , `parent_label_len` = :parent_label_len
                , `parent_label_val` = :parent_label_val
                , `node_type` = :node_type
                , `left_child_len` = :left_child_len
                , `left_child_label_val` = :left_child_label_val
                , `right_child_len` = :right_child_len
                , `right_child_label_val` = :right_child_label_val
                , `hash` = :hash
                , `p_last_epoch` = :p_last_epoch
                , `p_least_descendant_ep` = :p_least_descendant_ep
                , `p_parent_label_len` = :p_parent_label_len
                , `p_parent_label_val` = :p_parent_label_val
                , `p_node_type` = :p_node_type
                , `p_left_child_len` = :p_left_child_len
                , `p_left_child_label_val` = :p_left_child_label_val
                , `p_right_child_len` = :p_right_child_len
                , `p_right_child_label_val` = :p_right_child_label_val
                , `p_hash` = :p_hash"),
            DbRecord::ValueState(_) => format!("INSERT INTO `{TABLE_USER}` ({SELECT_USER_DATA}) VALUES (:username, :epoch, :version, :node_label_val, :node_label_len, :data)"),
        }
    }

    fn set_params(&self) -> Option<mysql_async::Params> {
        match &self {
            DbRecord::Azks(azks) => Some(
                params! { "key" => 1u8, "epoch" => azks.latest_epoch, "num_nodes" => azks.num_nodes },
            ),
            DbRecord::TreeNode(node) => Some(params! {
                "label_len" => node.label.label_len,
                "label_val" => node.label.label_val,
                // "Latest" node values
                "last_epoch" => node.latest_node.last_epoch,
                "least_descendant_ep" => node.latest_node.min_descendant_epoch,
                "parent_label_len" => node.latest_node.parent.label_len,
                "parent_label_val" => node.latest_node.parent.label_val,
                "node_type" => node.latest_node.node_type as u8,
                "left_child_len" => node.latest_node.left_child.map(|lc| lc.label_len),
                "left_child_label_val" => node.latest_node.left_child.map(|lc| lc.label_val),
                "right_child_len" => node.latest_node.right_child.map(|rc| rc.label_len),
                "right_child_label_val" => node.latest_node.right_child.map(|rc| rc.label_val),
                "hash" => node.latest_node.hash.0,
                // "Previous" node values
                "p_last_epoch" => node.previous_node.clone().map(|a| a.last_epoch),
                "p_least_descendant_ep" => node.previous_node.clone().map(|a| a.min_descendant_epoch),
                "p_parent_label_len" => node.previous_node.clone().map(|a| a.parent.label_len),
                "p_parent_label_val" => node.previous_node.clone().map(|a| a.parent.label_val),
                "p_node_type" => node.previous_node.clone().map(|a| a.node_type as u8),
                "p_left_child_len" => node.previous_node.clone().and_then(|a| a.left_child.map(|lc| lc.label_len)),
                "p_left_child_label_val" => node.previous_node.clone().and_then(|a| a.left_child.map(|lc| lc.label_val)),
                "p_right_child_len" => node.previous_node.clone().and_then(|a| a.right_child.map(|rc| rc.label_len)),
                "p_right_child_label_val" => node.previous_node.clone().and_then(|a| a.right_child.map(|rc| rc.label_val)),
                "p_hash" => node.previous_node.clone().map(|a| a.hash.0),
            }),
            DbRecord::ValueState(state) => Some(
                params! { "username" => state.get_id().0, "epoch" => state.epoch, "version" => state.version, "node_label_len" => state.label.label_len, "node_label_val" => state.label.label_val, "data" => state.value.0.clone() },
            ),
        }
    }

    fn set_batch_statement<St: Storable>(items: usize) -> String {
        let mut parts = "".to_string();
        for i in 0..items {
            match St::data_type() {
                StorageType::TreeNode => {
                    parts = format!(
                        "{parts}(:label_len{i}
                            , :label_val{i}
                            , :last_epoch{i}
                            , :least_descendant_ep{i}
                            , :parent_label_len{i}
                            , :parent_label_val{i}
                            , :node_type{i}
                            , :left_child_len{i}
                            , :left_child_label_val{i}
                            , :right_child_len{i}
                            , :right_child_label_val{i}
                            , :hash{i}
                            , :p_last_epoch{i}
                            , :p_least_descendant_ep{i}
                            , :p_parent_label_len{i}
                            , :p_parent_label_val{i}
                            , :p_node_type{i}
                            , :p_left_child_len{i}
                            , :p_left_child_label_val{i}
                            , :p_right_child_len{i}
                            , :p_right_child_label_val{i}
                            , :p_hash{i})"
                    );
                }
                StorageType::ValueState => {
                    parts = format!(
                        "{parts}(:username{i}, :epoch{i}, :version{i}, :node_label_val{i}, :node_label_len{i}, :data{i})"
                    );
                }
                _ => {
                    // azks
                }
            }

            if i < items - 1 {
                parts += ", ";
            }
        }

        match St::data_type() {
            StorageType::Azks => format!(
                "INSERT INTO `{TABLE_AZKS}` (`key`, {SELECT_AZKS_DATA})
            VALUES (:key, :epoch, :num_nodes) as new
            ON DUPLICATE KEY UPDATE `epoch` = new.epoch, `num_nodes` = new.num_nodes"
            ),
            StorageType::TreeNode => format!(
                "INSERT INTO `{TABLE_HISTORY_TREE_NODES}` ({SELECT_HISTORY_TREE_NODE_DATA})
            VALUES {parts} as new
            ON DUPLICATE KEY UPDATE
                `label_len` = new.label_len
                , `label_val` = new.label_val
                , `least_descendant_ep` = new.least_descendant_ep
                , `last_epoch` = new.last_epoch
                , `parent_label_len` = new.parent_label_len
                , `parent_label_val` = new.parent_label_val
                , `node_type` = new.node_type
                , `left_child_len` = new.left_child_len
                , `left_child_label_val` = new.left_child_label_val
                , `right_child_len` = new.right_child_len
                , `right_child_label_val` = new.right_child_label_val
                , `hash` = new.hash
                , `p_last_epoch` = new.p_last_epoch
                , `p_least_descendant_ep` = new.p_least_descendant_ep
                , `p_parent_label_len` = new.p_parent_label_len
                , `p_parent_label_val` = new.p_parent_label_val
                , `p_node_type` = new.p_node_type
                , `p_left_child_len` = new.p_left_child_len
                , `p_left_child_label_val` = new.p_left_child_label_val
                , `p_right_child_len` = new.p_right_child_len
                , `p_right_child_label_val` = new.p_right_child_label_val
                , `p_hash` = new.p_hash"
            ),
            StorageType::ValueState => format!(
                "INSERT INTO `{TABLE_USER}` ({SELECT_USER_DATA})
            VALUES {parts} as new
            ON DUPLICATE KEY UPDATE
                `data` = new.data
                , `node_label_val` = new.node_label_val
                , `node_label_len` = new.node_label_len
                , `version` = new.version"
            ),
        }
    }

    fn set_batch_params(items: &[DbRecord]) -> Result<mysql_async::Params> {
        let param_batch = items
            .iter()
            .enumerate()
            .map(|(idx, item)| match &item {
                DbRecord::Azks(azks) => Ok(vec![
                    ("key".to_string(), Value::from(1u8)),
                    ("epoch".to_string(), Value::from(azks.latest_epoch)),
                    ("num_nodes".to_string(), Value::from(azks.num_nodes)),
                ]),
                DbRecord::TreeNode(node) => {
                    let pnode = &node.previous_node;
                    Ok(vec![
                        (format!("label_len{idx}"), Value::from(node.label.label_len)),
                        (format!("label_val{idx}"), Value::from(node.label.label_val)),
                        (
                            format!("last_epoch{idx}"),
                            Value::from(node.latest_node.last_epoch),
                        ),
                        (
                            format!("least_descendant_ep{idx}"),
                            Value::from(node.latest_node.min_descendant_epoch),
                        ),
                        (
                            format!("parent_label_len{idx}"),
                            Value::from(node.latest_node.parent.label_len),
                        ),
                        (
                            format!("parent_label_val{idx}"),
                            Value::from(node.latest_node.parent.label_val),
                        ),
                        (
                            format!("node_type{idx}"),
                            Value::from(node.latest_node.node_type as u8),
                        ),
                        (
                            format!("left_child_len{idx}"),
                            Value::from(node.latest_node.left_child.map(|lc| lc.label_len)),
                        ),
                        (
                            format!("left_child_label_val{idx}"),
                            Value::from(node.latest_node.left_child.map(|lc| lc.label_val)),
                        ),
                        (
                            format!("right_child_len{idx}"),
                            Value::from(node.latest_node.right_child.map(|rc| rc.label_len)),
                        ),
                        (
                            format!("right_child_label_val{idx}"),
                            Value::from(node.latest_node.right_child.map(|rc| rc.label_val)),
                        ),
                        (format!("hash{idx}"), Value::from(node.latest_node.hash.0)),
                        (
                            format!("p_last_epoch{idx}"),
                            Value::from(pnode.clone().map(|a| a.last_epoch)),
                        ),
                        (
                            format!("p_least_descendant_ep{idx}"),
                            Value::from(pnode.clone().map(|a| a.min_descendant_epoch)),
                        ),
                        (
                            format!("p_parent_label_len{idx}"),
                            Value::from(pnode.clone().map(|a| a.parent.label_len)),
                        ),
                        (
                            format!("p_parent_label_val{idx}"),
                            Value::from(pnode.clone().map(|a| a.parent.label_val)),
                        ),
                        (
                            format!("p_node_type{idx}"),
                            Value::from(pnode.clone().map(|a| a.node_type as u8)),
                        ),
                        (
                            format!("p_left_child_len{idx}"),
                            Value::from(
                                pnode
                                    .clone()
                                    .and_then(|a| a.left_child.map(|lc| lc.label_len)),
                            ),
                        ),
                        (
                            format!("p_left_child_label_val{idx}"),
                            Value::from(
                                pnode
                                    .clone()
                                    .and_then(|a| a.left_child.map(|lc| lc.label_val)),
                            ),
                        ),
                        (
                            format!("p_right_child_len{idx}"),
                            Value::from(
                                pnode
                                    .clone()
                                    .and_then(|a| a.right_child.map(|rc| rc.label_len)),
                            ),
                        ),
                        (
                            format!("p_right_child_label_val{idx}"),
                            Value::from(
                                pnode
                                    .clone()
                                    .and_then(|a| a.right_child.map(|rc| rc.label_val)),
                            ),
                        ),
                        (
                            format!("p_hash{idx}"),
                            Value::from(pnode.clone().map(|a| a.hash.0)),
                        ),
                    ])
                }
                DbRecord::ValueState(state) => Ok(vec![
                    (format!("username{idx}"), Value::from(state.get_id().0)),
                    (format!("epoch{idx}"), Value::from(state.epoch)),
                    (format!("version{idx}"), Value::from(state.version)),
                    (
                        format!("node_label_len{idx}"),
                        Value::from(state.label.label_len),
                    ),
                    (
                        format!("node_label_val{idx}"),
                        Value::from(state.label.label_val),
                    ),
                    (format!("data{idx}"), Value::from(state.value.0.clone())),
                ]),
            })
            .collect::<Result<Vec<_>>>()?
            .into_iter()
            .flatten()
            .collect::<Vec<_>>();
        Ok(mysql_async::Params::from(param_batch))
    }

    fn get_statement<St: Storable>() -> String {
        match St::data_type() {
            StorageType::Azks => format!("SELECT {SELECT_AZKS_DATA} FROM `{TABLE_AZKS}`"),
            StorageType::TreeNode => {
                format!("SELECT {SELECT_HISTORY_TREE_NODE_DATA} FROM `{TABLE_HISTORY_TREE_NODES}`")
            }
            StorageType::ValueState => format!("SELECT {SELECT_USER_DATA} FROM `{TABLE_USER}`"),
        }
    }

    fn get_batch_create_temp_table<St: Storable>() -> Option<String> {
        match St::data_type() {
            StorageType::Azks => None,
            StorageType::TreeNode => {
                Some(
                    format!(
                        "CREATE TEMPORARY TABLE `{TEMP_IDS_TABLE}`(`label_len` INT UNSIGNED NOT NULL, `label_val` VARBINARY(32) NOT NULL, PRIMARY KEY(`label_len`, `label_val`))"
                    )
                )
            },
            StorageType::ValueState => {
                Some(
                    format!(
                        "CREATE TEMPORARY TABLE `{TEMP_IDS_TABLE}`(`username` VARCHAR(256) NOT NULL, `epoch` BIGINT UNSIGNED NOT NULL, PRIMARY KEY(`username`, `epoch`))"
                    )
                )
            },
        }
    }

    fn get_batch_fill_temp_table<St: Storable>(num_items: Option<usize>) -> String {
        let mut statement = match St::data_type() {
            StorageType::Azks => "".to_string(),
            StorageType::TreeNode => {
                format!("INSERT INTO `{TEMP_IDS_TABLE}` (`label_len`, `label_val`) VALUES ")
            }
            StorageType::ValueState => {
                format!("INSERT INTO `{TEMP_IDS_TABLE}` (`username`, `epoch`) VALUES ")
            }
        };
        if let Some(item_count) = num_items {
            for i in 0..item_count {
                let append = match St::data_type() {
                    StorageType::Azks => String::from(""),
                    StorageType::TreeNode => {
                        format!("(:label_len{i}, :label_val{i})")
                    }
                    StorageType::ValueState => {
                        format!("(:username{i}, :epoch{i})")
                    }
                };
                statement = format!("{statement}{append}");

                if i < item_count - 1 {
                    // inner-item, append a comma
                    statement += ", ";
                }
            }
        } else {
            statement += match St::data_type() {
                StorageType::Azks => "",
                StorageType::TreeNode => "(:label_len, :label_val)",
                StorageType::ValueState => "(:username, :epoch)",
            };
        }
        statement
    }

    fn get_batch_statement<St: Storable>() -> String {
        match St::data_type() {
            StorageType::Azks => {
                format!("SELECT {SELECT_AZKS_DATA} FROM `{TABLE_AZKS}` LIMIT 1")
            }
            StorageType::TreeNode => {
                format!(
                    "SELECT
                        a.`label_len`
                        , a.`label_val`
                        , a.`last_epoch`
                        , a.`least_descendant_ep`
                        , a.`parent_label_len`
                        , a.`parent_label_val`
                        , a.`node_type`
                        , a.`left_child_len`
                        , a.`left_child_label_val`
                        , a.`right_child_len`
                        , a.`right_child_label_val`
                        , a.`hash`, a.`p_last_epoch`
                        , a.`p_least_descendant_ep`
                        , a.`p_parent_label_len`
                        , a.`p_parent_label_val`
                        , a.`p_node_type`
                        , a.`p_left_child_len`
                        , a.`p_left_child_label_val`
                        , a.`p_right_child_len`
                        , a.`p_right_child_label_val`
                        , a.`p_hash`
                    FROM `{TABLE_HISTORY_TREE_NODES}` a
                    INNER JOIN {TEMP_IDS_TABLE} ids
                        ON ids.`label_len` = a.`label_len`
                        AND ids.`label_val` = a.`label_val`"
                )
            }
            StorageType::ValueState => {
                format!(
                    "SELECT
                        a.`username`
                        , a.`epoch`
                        , a.`version`
                        , a.`node_label_val`
                        , a.`node_label_len`
                        , a.`data`
                    FROM `{TABLE_USER}` a
                    INNER JOIN {TEMP_IDS_TABLE} ids
                        ON ids.`username` = a.`username`
                        AND ids.`epoch` = a.`epoch`"
                )
            }
        }
    }

    fn get_specific_statement<St: Storable>() -> String {
        match St::data_type() {
            StorageType::Azks => {
                format!("SELECT {SELECT_AZKS_DATA} FROM `{TABLE_AZKS}` LIMIT 1")
            }
            StorageType::TreeNode => format!(
                "SELECT {SELECT_HISTORY_TREE_NODE_DATA} FROM `{TABLE_HISTORY_TREE_NODES}` WHERE `label_len` = :label_len AND `label_val` = :label_val"
            ),
            StorageType::ValueState => format!(
                "SELECT {SELECT_USER_DATA} FROM `{TABLE_USER}` WHERE `username` = :username AND `epoch` = :epoch"
            ),
        }
    }

    fn get_specific_params<St: Storable>(key: &St::StorageKey) -> Option<mysql_async::Params> {
        match St::data_type() {
            StorageType::Azks => None,
            StorageType::TreeNode => {
                let bin = St::get_full_binary_key_id(key);
                if let Ok(back) = TreeNodeWithPreviousValue::key_from_full_binary(&bin) {
                    Some(params! {
                        "label_len" => back.0.label_len,
                        "label_val" => back.0.label_val,
                    })
                } else {
                    None
                }
            }
            StorageType::ValueState => {
                let bin = St::get_full_binary_key_id(key);
                if let Ok(back) = akd::storage::types::ValueState::key_from_full_binary(&bin) {
                    Some(params! {
                        "username" => back.0,
                        "epoch" => back.1
                    })
                } else {
                    None
                }
            }
        }
    }

    fn get_multi_row_specific_params<St: Storable>(
        keys: &[St::StorageKey],
    ) -> Option<mysql_async::Params> {
        match St::data_type() {
            StorageType::Azks => None,
            StorageType::TreeNode => {
                let pvec = keys
                    .iter()
                    .enumerate()
                    .flat_map(|(idx, key)| {
                        let bin = St::get_full_binary_key_id(key);
                        // Since these are constructed from a safe key, they should never fail
                        // so we'll leave the unwrap to simplify
                        let back: NodeKey =
                            TreeNodeWithPreviousValue::key_from_full_binary(&bin).unwrap();
                        vec![
                            (format!("label_len{idx}"), Value::from(back.0.label_len)),
                            (format!("label_val{idx}"), Value::from(back.0.label_val)),
                        ]
                    })
                    .collect::<Vec<_>>();
                Some(mysql_async::Params::from(pvec))
            }
            StorageType::ValueState => {
                let pvec = keys
                    .iter()
                    .enumerate()
                    .flat_map(|(idx, key)| {
                        let bin = St::get_full_binary_key_id(key);
                        // Since these are constructed from a safe key, they should never fail
                        // so we'll leave the unwrap to simplify
                        let back: akd::storage::types::ValueStateKey =
                            akd::storage::types::ValueState::key_from_full_binary(&bin).unwrap();
                        vec![
                            (format!("username{idx}"), Value::from(back.0.clone())),
                            (format!("epoch{idx}"), Value::from(back.1)),
                        ]
                    })
                    .collect::<Vec<_>>();
                Some(mysql_async::Params::from(pvec))
            }
        }
    }

    fn from_row<St: Storable>(row: &mut mysql_async::Row) -> core::result::Result<Self, MySqlError>
    where
        Self: std::marker::Sized,
    {
        fn cast_err() -> MySqlError {
            MySqlError::from(mysql_async::ServerError {
                state: "".to_string(),
                code: 0,
                message: "Failed to cast label:val into [u8; 32]".to_string(),
            })
        }

        fn optional_child_label(
            child_val: Option<Value>,
            child_len: Option<Value>,
        ) -> core::result::Result<Option<NodeLabel>, MySqlError> {
            match (child_val, child_len) {
                (Some(Value::NULL), _) => Ok(None),
                (_, Some(Value::NULL)) => Ok(None),
                (None, _) => Ok(None),
                (_, None) => Ok(None),
                (Some(possible_bytes), Some(possible_u32)) => {
                    match (
                        from_value_opt::<u32>(possible_u32),
                        from_value_opt::<Vec<u8>>(possible_bytes),
                    ) {
                        (Ok(len), Ok(val)) => Ok(Some(NodeLabel::new(
                            val.try_into().map_err(|_| cast_err())?,
                            len,
                        ))),
                        (Err(_len_err), _) => Err(Error::from(mysql_async::ServerError {
                            state: "".to_string(),
                            code: 0,
                            message: "Error decoding length".to_string(),
                        })),
                        (_, Err(_val_err)) => Err(Error::from(mysql_async::ServerError {
                            state: "".to_string(),
                            code: 0,
                            message: "Error decoding value".to_string(),
                        })),
                    }
                }
            }
        }

        match St::data_type() {
            StorageType::Azks => {
                // epoch, num_nodes
                if let (Some(Ok(epoch)), Some(Ok(num_nodes))) = (row.take_opt(0), row.take_opt(1)) {
                    let azks = DbRecord::build_azks(epoch, num_nodes);
                    return Ok(DbRecord::Azks(azks));
                }
            }
            StorageType::TreeNode => {
                // A NOTE ABOUT THE SYNTAX HERE: The outer-most Some(**) of the row.take(..) indicates the COLUMN EXISTS, not that it actually has a value.
                // Therefore the signature of these fields you want is Some(Option<value>) for a nullable column. The inner option type is not necessary
                // for types which have NOT NULL in the SQL definition
                if let (
                    Some(Ok(label_len)),
                    Some(Ok(label_val)),
                    Some(Ok(last_epoch)),
                    Some(Ok(least_descendant_ep)),
                    Some(Ok(parent_label_len)),
                    Some(Ok(parent_label_val)),
                    Some(Ok(node_type)),
                    left_child_len_res,
                    left_child_val_res,
                    right_child_len_res,
                    right_child_val_res,
                    Some(Ok(hash)),
                    Some(p_last_epoch),
                    Some(p_least_descendant_ep),
                    Some(p_parent_label_len),
                    Some(p_parent_label_val),
                    Some(p_node_type),
                    p_left_child_len,
                    p_left_child_label_val,
                    p_right_child_len,
                    p_right_child_label_val,
                    Some(p_hash),
                ) = (
                    row.take_opt(0),
                    row.take_opt(1),
                    row.take_opt(2),
                    row.take_opt(3),
                    row.take_opt(4),
                    row.take_opt(5),
                    row.take_opt(6),
                    row.take(7),
                    row.take(8),
                    row.take(9),
                    row.take(10),
                    row.take_opt(11),
                    row.take(12),
                    row.take(13),
                    row.take(14),
                    row.take(15),
                    row.take(16),
                    row.take(17),
                    row.take(18),
                    row.take(19),
                    row.take(20),
                    row.take(21),
                ) {
                    let left_child = optional_child_label(left_child_val_res, left_child_len_res)?;
                    let right_child =
                        optional_child_label(right_child_val_res, right_child_len_res)?;
                    let p_left_child =
                        optional_child_label(p_left_child_label_val, p_left_child_len)?;
                    let p_right_child =
                        optional_child_label(p_right_child_label_val, p_right_child_len)?;

                    let label_val_vec: Vec<u8> = label_val;

                    let parent_label_val_vec: Vec<u8> = parent_label_val;
                    let prev_parent_label_val_vec: Option<Vec<u8>> = p_parent_label_val;
                    let hash_vec: Vec<u8> = hash;
                    let prev_hash_vec: Option<Vec<u8>> = p_hash;

                    let massaged_prev_parent_label_val: Option<[u8; 32]> =
                        match prev_parent_label_val_vec {
                            Some(v) => Some(v.try_into().map_err(|_| cast_err())?),
                            None => None,
                        };
                    let massaged_hash_vec: akd::Digest =
                        akd::hash::try_parse_digest(&hash_vec).map_err(|_| cast_err())?;
                    let massaged_prev_hash_vec: Option<akd::Digest> = match prev_hash_vec {
                        Some(v) => match akd::hash::try_parse_digest(&v).map_err(|_| cast_err()) {
                            Ok(r) => Some(r),
                            Err(err) => return Err(err),
                        },
                        None => None,
                    };

                    let node = DbRecord::build_tree_node_with_previous_value(
                        label_val_vec.try_into().map_err(|_| cast_err())?,
                        label_len,
                        last_epoch,
                        least_descendant_ep,
                        parent_label_val_vec.try_into().map_err(|_| cast_err())?,
                        parent_label_len,
                        node_type,
                        left_child,
                        right_child,
                        massaged_hash_vec,
                        p_last_epoch,
                        p_least_descendant_ep,
                        massaged_prev_parent_label_val,
                        p_parent_label_len,
                        p_node_type,
                        p_left_child,
                        p_right_child,
                        massaged_prev_hash_vec,
                    );
                    return Ok(DbRecord::TreeNode(node));
                }
            }
            StorageType::ValueState => {
                // `username`, `epoch`, `version`, `node_label_val`, `node_label_len`, `data`
                if let (
                    Some(Ok(username)),
                    Some(Ok(epoch)),
                    Some(Ok(version)),
                    Some(Ok(node_label_val)),
                    Some(Ok(node_label_len)),
                    Some(Ok(data)),
                ) = (
                    row.take_opt(0),
                    row.take_opt(1),
                    row.take_opt(2),
                    row.take_opt(3),
                    row.take_opt(4),
                    row.take_opt(5),
                ) {
                    let node_label_val_vec: Vec<u8> = node_label_val;
                    let state = DbRecord::build_user_state(
                        username,
                        data,
                        version,
                        node_label_len,
                        node_label_val_vec.try_into().map_err(|_| cast_err())?,
                        epoch,
                    );
                    return Ok(DbRecord::ValueState(state));
                }
            }
        }
        // fallback
        let err = MySqlError::Driver(mysql_async::DriverError::FromRow { row: row.clone() });
        Err(err)
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

pub mod fixture_generator;

pub mod test_suites;


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

use akd_test_tools::fixture_generator;

#[tokio::main]
async fn main() {
    fixture_generator::run().await;
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

extern crate thread_id;

use akd::ecvrf::VRFKeyStorage;
use akd::storage::Database;
use akd::Directory;
use akd::HistoryParams;
use akd::{AkdLabel, AkdValue};
use rand::distributions::Alphanumeric;
use rand::seq::IteratorRandom;
use rand::{thread_rng, Rng};

/// The suite of tests to run against a fully-instantated and storage-backed directory.
/// This will publish 3 epochs of ```num_users``` records and
/// perform 10 random lookup proofs + 2 random history proofs + and audit proof from epochs 1u64 -> 2u64
pub async fn directory_test_suite<S: Database + 'static, V: VRFKeyStorage>(
    mysql_db: &akd::storage::StorageManager<S>,
    num_users: usize,
    vrf: &V,
) {
    // generate the test data
    let mut rng = thread_rng();

    let mut users: Vec<String> = vec![];
    for _ in 0..num_users {
        users.push(
            thread_rng()
                .sample_iter(&Alphanumeric)
                .take(30)
                .map(char::from)
                .collect(),
        );
    }
    let mut root_hashes = vec![];
    // create & test the directory
    let maybe_dir = Directory::<_, _>::new(mysql_db.clone(), vrf.clone()).await;
    match maybe_dir {
        Err(akd_error) => panic!("Error initializing directory: {:?}", akd_error),
        Ok(dir) => {
            // Publish 3 epochs of user material
            for i in 1..=3 {
                let mut data = Vec::new();
                for value in users.iter() {
                    data.push((
                        AkdLabel::from(value),
                        AkdValue(format!("{i}").as_bytes().to_vec()),
                    ));
                }

                if let Err(error) = dir.publish(data).await {
                    panic!("Error publishing batch {:?}", error);
                }
                let root_hash = dir.get_epoch_hash().await.unwrap().1;
                root_hashes.push(root_hash);
            }

            // Perform 10 random lookup proofs on the published users
            for user in users.iter().choose_multiple(&mut rng, 10) {
                let key = AkdLabel::from(user);
                match dir.lookup(key.clone()).await {
                    Err(error) => panic!("Error looking up user information {:?}", error),
                    Ok((proof, root_hash)) => {
                        let vrf_pk = dir.get_public_key().await.unwrap();
                        if let Err(error) = akd::client::lookup_verify(
                            vrf_pk.as_bytes(),
                            root_hash.hash(),
                            key,
                            proof,
                        ) {
                            panic!("Lookup proof failed to verify {:?}", error);
                        }
                    }
                }
            }

            // Perform 2 random history proofs on the published material
            for user in users.iter().choose_multiple(&mut rng, 2) {
                let key = AkdLabel::from(user);
                match dir.key_history(&key, HistoryParams::default()).await {
                    Err(error) => panic!("Error performing key history retrieval {:?}", error),
                    Ok((proof, root_hash)) => {
                        let vrf_pk = dir.get_public_key().await.unwrap();
                        if let Err(error) = akd::client::key_history_verify(
                            vrf_pk.as_bytes(),
                            root_hash.hash(),
                            root_hash.epoch(),
                            key,
                            proof,
                            akd::HistoryVerificationParams::default(),
                        ) {
                            panic!("History proof failed to verify {:?}", error);
                        }
                    }
                }
            }

            // Perform an audit proof from 1u64 -> 2u64

            mysql_db.log_metrics(log::Level::Info).await;
            log::warn!("Beginning audit proof generation");
            mysql_db.flush_cache().await;
            match dir.audit(1u64, 2u64).await {
                Err(error) => panic!("Error perform audit proof retrieval {:?}", error),
                Ok(proof) => {
                    mysql_db.log_metrics(log::Level::Info).await;
                    log::warn!("Done with audit proof generation");
                    let start_root_hash = root_hashes[0];
                    let end_root_hash = root_hashes[1];
                    akd::auditor::audit_verify(vec![start_root_hash, end_root_hash], proof)
                        .await
                        .unwrap();
                }
            }
        }
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! This module contains the struct definitions of the tool output and main
//! fixture generation logic.

use std::collections::HashMap;
use std::env;
use std::fs::File;
use std::io::Write;

use akd::directory::Directory;
use akd::storage::types::DbRecord;
use akd::storage::{StorageManager, StorageUtil};
use akd::{AkdLabel, AkdValue};
use clap::Parser;
use rand::rngs::StdRng;
use rand::Rng;
use rand::SeedableRng;
use serde::{Deserialize, Serialize};

use crate::fixture_generator::parser::Args;
use crate::fixture_generator::writer::yaml::YamlWriter;
use crate::fixture_generator::writer::Writer;

/// Directory state comprises all database records at a particular epoch.
#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct State {
    pub epoch: u32,
    pub records: Vec<DbRecord>,
}

/// Delta comprises all key updates published to the directory to advance to an
/// epoch.
#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Delta {
    pub epoch: u32,
    pub updates: Vec<(AkdLabel, AkdValue)>,
}

/// Metadata about the output, including arguments passed to this tool and
/// the tool version.
#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Metadata {
    pub args: Args,
    pub version: String,
}

// "@" has to be separated from "generated" or linters might ignore this file
const HEADER_COMMENT: &str = concat!(
    "@",
    "generated This file was automatically generated by \n\
    the fixture generator tool with the following command:\n\n\
    cargo run -- \\"
);
const METADATA_COMMENT: &str = "Metadata";
const STATE_COMMENT: &str = "State - Epoch";
const DELTA_COMMENT: &str = "Delta - Epoch";

pub async fn run() {
    let args = Args::parse();
    generate(args).await;
}

pub(crate) async fn generate(args: Args) {
    let mut rng = StdRng::seed_from_u64(42);

    // args assertions
    assert!(args.max_updates >= args.min_updates);
    assert!(args
        .capture_states
        .as_ref()
        .map_or(true, |states| states.iter().max().unwrap() <= &args.epochs));
    assert!(args
        .capture_deltas
        .as_ref()
        .map_or(true, |deltas| deltas.iter().max().unwrap() <= &args.epochs));

    // process users
    let mut user_map = HashMap::new();
    for user in &args.users {
        let mut events_map = HashMap::new();
        for event in &user.events {
            events_map.insert(event.epoch, event.value.clone());
        }
        user_map.insert(user.label.clone(), events_map);
    }

    // initialize writer
    let buffer: Box<dyn Write> = if let Some(ref file_name) = args.out {
        Box::new(File::create(file_name).unwrap())
    } else {
        Box::new(std::io::stdout())
    };
    let mut writer = YamlWriter::new(buffer);

    // write raw args as comment
    let raw_args = format!(
        " {}",
        env::args().skip(1).collect::<Vec<String>>().join(" ")
    );
    writer.write_comment(HEADER_COMMENT);
    raw_args
        .split(" -")
        .skip(1)
        .map(|arg| format!("  -{arg} \\"))
        .for_each(|comment| writer.write_comment(&comment));

    // write fixture metadata
    let comment = METADATA_COMMENT.to_string();
    let metadata = Metadata {
        args: args.clone(),
        version: env!("CARGO_PKG_VERSION").to_string(),
    };
    writer.write_line();
    writer.write_comment(&comment);
    writer.write_object(metadata);

    // initialize directory
    let db = akd::storage::memory::AsyncInMemoryDatabase::new();
    let vrf = akd::ecvrf::HardCodedAkdVRF {};
    let storage_manager = StorageManager::new_no_cache(db);
    let akd = Directory::<_, _>::new(storage_manager.clone(), vrf)
        .await
        .unwrap();

    for epoch in 1..=args.epochs {
        // gather specified key updates
        let mut updates = vec![];
        for (label, events) in user_map.iter() {
            if let Some(maybe_value) = events.get(&epoch) {
                let value = maybe_value
                    .clone()
                    .unwrap_or_else(|| AkdValue::random(&mut rng));
                updates.push((label.clone(), value))
            }
        }

        // generate random key updates if allowed
        if !args.no_generated_updates {
            let num_updates = rng.gen_range(args.min_updates, args.max_updates);
            for _ in updates.len()..num_updates as usize {
                updates.push((AkdLabel::random(&mut rng), AkdValue::random(&mut rng)));
            }
        }

        // write delta if required
        if let Some(ref deltas) = args.capture_deltas {
            if deltas.contains(&epoch) {
                let comment = format!("{DELTA_COMMENT} {epoch}");
                let delta = Delta {
                    epoch,
                    updates: updates.clone(),
                };
                writer.write_line();
                writer.write_comment(&comment);
                writer.write_object(delta);
            }
        }

        // perform publish
        akd.publish(updates.clone()).await.unwrap();

        // write state if required
        if let Some(ref states) = args.capture_states {
            if states.contains(&epoch) {
                let comment = format!("{STATE_COMMENT} {epoch}");

                // Sort the records by label to make the output deterministic.
                let mut records = storage_manager
                    .get_db()
                    .batch_get_all_direct()
                    .await
                    .unwrap();
                records.sort();

                let state = State { epoch, records };

                writer.write_line();
                writer.write_comment(&comment);
                writer.write_object(state);
            }
        }
    }

    // flush writer and exit
    writer.flush();
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! A CLI tool for generating directory fixtures for debug and testing purposes.
//! Run cargo run -- --help for options. Example command:
//!
//!   cargo run -- \
//!     --user "User1: 1, (9, 'abc'), (10, 'def')" \
//!     --user "User2: 1, 2, 3, 4, 5, 6, 7, 8, 9, 10" \
//!     --epochs 10 \
//!     --max_updates 5 \
//!     --capture_states 9 10 \
//!     --capture_deltas 10
//!

mod examples;
mod generator;
mod parser;
mod writer;

pub mod reader;

/// Re-export generator run function.
pub use generator::run;


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! This module contains the CLI argument definitions and parser.

use akd::{AkdLabel, AkdValue};
use clap::{AppSettings, Parser};
use regex::Regex;
use serde::{Deserialize, Serialize};

/// Any alphanumeric string - spaces are allowed e.g. "User123" or "User 123"
const USER_PATTERN: &str = r"[\w\s]+";

/// A solo string of digits e.g. "10" or a tuple of digits and a string
/// e.g."(10, 'abc')"
const EVENT_PATTERN: &str = r"\d+|(\(\s*(\d+)\s*,\s*'(\w*)'\s*\))";

/// A key update the tool should include in the tree at the given epoch.
/// If "value" is None, the tool will randomly generate a value for the epoch.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct UserEvent {
    pub epoch: u32,
    pub value: Option<AkdValue>,
}

/// A user whose key update events should be included in the tree.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct User {
    pub label: AkdLabel,
    pub events: Vec<UserEvent>,
}

/// This tool allows a directory to be created with specified and random
/// contents, capturing the directory state and epoch-to-epoch delta in
/// an output file for use in debugging and as test fixtures.
#[derive(Parser, Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
#[clap(setting = AppSettings::DeriveDisplayOrder)]
pub struct Args {
    /// Users and their associated key update events.
    /// A username is expected, followed by a colon and a list of epochs OR
    /// (epoch, value). Usernames are expected to be utf-8 strings, which will
    /// be internally interpreted as bytes.
    /// The following are valid examples of user arguments:
    ///   --user "username: 1, 3, (5, 'xyz')"
    ///   --user="username: [(1,'abc'), 2]"
    ///   -u "some username: 1"
    #[clap(
        long = "user",
        short = 'u',
        multiple_occurrences = true,
        parse(try_from_str = parse_user_events),
    )]
    pub users: Vec<User>,

    /// Number of epochs to advance the tree by
    /// e.g. a value of 3 will perform 3 publishes on an empty directory.
    #[clap(long = "epochs", short = 'e')]
    pub epochs: u32,

    /// Maximum number of key updates **per epoch** the tool should perform.
    /// Note that all user events explicitly passed for an epoch will be
    /// included even if the number exceeds this value.
    #[clap(long = "max_updates", default_value = "10")]
    pub max_updates: u32,

    /// Minimum number of key updates **per epoch** the tool should perform.
    /// The tool will generate random labels and values to include in an epoch
    /// if the user events explicitly passed for an epoch are not sufficients.
    #[clap(long = "min_updates", default_value = "0")]
    pub min_updates: u32,

    /// Epochs where the state of the directory should be captured in the output
    /// e.g. the value 3 will output all db records after epoch 3 is performed.
    /// Multiple values are accepted e.g. --capture_states 9 10
    #[clap(long = "capture_states", short = 's', multiple_values = true)]
    pub capture_states: Option<Vec<u32>>,

    /// Epochs where the key updates required to bring the directory to the
    /// epoch should be captured in the output.
    /// e.g. the value 3 will output all key updates that were performed to
    /// advance the directory from epoch 2 to 3.
    /// Multiple values are accepted e.g. --capture_deltas 9 10
    #[clap(long = "capture_deltas", short = 'd', multiple_values = true)]
    pub capture_deltas: Option<Vec<u32>>,

    /// Name of output file.
    /// If omitted, output will be printed to stdout.
    #[clap(long = "out", short = 'o')]
    pub out: Option<String>,

    /// Stops tool from generating random key updates in publishes.
    /// Use this if you want the tool to only use explicitly passed key updates.
    /// Explicilty passed key updates without values would still use randomly
    /// generated values.
    #[clap(long = "no_generated_updates", short = 'n')]
    pub no_generated_updates: bool,
}

fn parse_user_events(s: &str) -> Result<User, String> {
    let mut split = s.split(':');
    let username_text = split.next().unwrap();
    let maybe_events_text = split.next();

    let username = Regex::new(USER_PATTERN)
        .unwrap()
        .captures(username_text)
        .unwrap()
        .get(0)
        .unwrap()
        .as_str();

    let events = if let Some(events_text) = maybe_events_text {
        Regex::new(EVENT_PATTERN)
            .unwrap()
            .captures_iter(events_text)
            .map(|event| {
                let epoch: u32;
                let value: Option<AkdValue>;
                if event.get(1).is_some() {
                    epoch = event.get(2).unwrap().as_str().parse().unwrap();
                    value = Some(AkdValue::from(event.get(3).unwrap().as_str()));
                } else {
                    epoch = event.get(0).unwrap().as_str().parse().unwrap();
                    value = None;
                }
                UserEvent { epoch, value }
            })
            .collect::<Vec<_>>()
    } else {
        vec![]
    };

    Ok(User {
        label: AkdLabel::from(username),
        events,
    })
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! Example test utilizing a fixture file.

use std::fs::File;

use akd::{
    directory::Directory,
    ecvrf::HardCodedAkdVRF,
    storage::{memory::AsyncInMemoryDatabase, Database, StorageManager, StorageUtil},
};

use crate::fixture_generator::reader::yaml::YamlFileReader;
use crate::fixture_generator::reader::Reader;

// Contains two consecutive states and the delta between them
const TEST_FILE: &str = "src/fixture_generator/examples/test.yaml";

#[tokio::test]
async fn test_use_fixture() {
    // load fixture
    let mut reader = YamlFileReader::new(File::open(TEST_FILE).unwrap()).unwrap();
    let metadata = reader.read_metadata().unwrap();
    let epochs = metadata.args.capture_states.unwrap();

    // prepare directory with initial state
    let initial_state = reader.read_state(epochs[0]).unwrap();
    let db = AsyncInMemoryDatabase::new();
    db.batch_set(initial_state.records, akd::storage::DbSetState::General)
        .await
        .unwrap();
    let vrf = HardCodedAkdVRF {};
    let storage_manager = StorageManager::new_no_cache(db);
    let akd = Directory::<_, _>::new(storage_manager.clone(), vrf)
        .await
        .unwrap();

    // publish delta updates
    let delta = reader.read_delta(epochs[1]).unwrap();
    akd.publish(delta.updates).await.unwrap();

    // assert final directory state
    let final_state = reader.read_state(epochs[1]).unwrap();
    let records = storage_manager
        .get_db()
        .batch_get_all_direct()
        .await
        .unwrap();
    assert_eq!(final_state.records.len(), records.len());
    assert!(records.iter().all(|r| final_state.records.contains(r)));
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! This module contains examples of how fixtures can be read and used in tests.

#[cfg(test)]
mod example_tests;


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! This module contains the Reader trait to deserialize the tool's serde-compatible
//! objects from a formatted file, as well as implementations of the trait.

use std::result::Result;

use crate::fixture_generator::generator::{Delta, Metadata, State};

/// Interface for reading output generated by the tool.
pub trait Reader {
    /// Reads a metadata object.
    fn read_metadata(&mut self) -> Result<Metadata, ReaderError>;

    /// Reads a state object for a given epoch.
    fn read_state(&mut self, epoch: u32) -> Result<State, ReaderError>;

    /// Reads a delta object for a given epoch.
    fn read_delta(&mut self, epoch: u32) -> Result<Delta, ReaderError>;
}

#[derive(Debug, PartialEq, Eq)]
pub enum ReaderError {
    NotFound,
    Format(String),
    Input(String),
}

impl std::error::Error for ReaderError {}

impl std::fmt::Display for ReaderError {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            ReaderError::NotFound => write!(f, "Object not found"),
            ReaderError::Format(message) => write!(f, "Unexpected format: {message}"),
            ReaderError::Input(message) => write!(f, "Input stream error: {message}"),
        }
    }
}

/// YAML implementor of Reader trait.
pub mod yaml;

#[cfg(test)]
mod tests;


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! Tests basic reader behavior.

use std::env;
use std::fs::File;

use assert_fs::fixture::{FileWriteStr, NamedTempFile, TempDir};
use clap::Parser;

use crate::fixture_generator::generator;
use crate::fixture_generator::parser::Args;
use crate::fixture_generator::reader::yaml::YamlFileReader;
use crate::fixture_generator::reader::{Reader, ReaderError};

#[tokio::test]
async fn test_read() {
    // generate a temp fixture file
    let file = TempDir::new().unwrap().with_file_name("test.yaml");
    let args = Args::parse_from(vec![
        env!("CARGO_CRATE_NAME"),
        "--epochs",
        "10",
        "--capture_deltas",
        "10",
        "--capture_states",
        "9",
        "10",
        "--out",
        &format!("{}", file.display()),
    ]);
    generator::generate(args).await;

    // initialize reader
    let mut reader = YamlFileReader::new(File::open(file).unwrap()).unwrap();

    // objects can be read in any order
    assert!(reader.read_state(10).is_ok());
    assert!(reader.read_delta(10).is_ok());
    assert!(reader.read_state(9).is_ok());
    assert!(reader.read_metadata().is_ok());

    // reading a non-existent object will return a NotFound error
    assert_eq!(Err(ReaderError::NotFound), reader.read_delta(9));
    assert_eq!(Err(ReaderError::NotFound), reader.read_state(11));

    // reading an already read object is OK
    assert!(reader.read_metadata().is_ok());
}

#[tokio::test]
async fn test_read_invalid_format() {
    // create an invalid file with no YAML separators
    let file = NamedTempFile::new("invalid.yaml").unwrap();
    file.write_str("a\nb\nc\n").unwrap();

    // initialize reader
    let mut reader = YamlFileReader::new(File::open(file).unwrap()).unwrap();

    // reading any object will return a Format error
    assert!(matches!(
        reader.read_metadata(),
        Err(ReaderError::Format(_))
    ));
    assert!(matches!(reader.read_state(0), Err(ReaderError::Format(_))));
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! This module contains an implementor of the Reader trait for the YAML format.

use std::fmt::Write as _;
use std::fs::File;
use std::io::{BufRead, BufReader, Lines, Seek};
use std::iter::Peekable;
use std::result::Result; // import without risk of name clashing

use serde::de::DeserializeOwned;

use crate::fixture_generator::generator::{Delta, Metadata, State};
use crate::fixture_generator::reader::{Reader, ReaderError};

const YAML_SEPARATOR: &str = "---";

impl From<std::io::Error> for ReaderError {
    fn from(error: std::io::Error) -> Self {
        ReaderError::Input(error.to_string())
    }
}

/// YAML format file reader.
pub struct YamlFileReader {
    file: File,
    index: u32,
    buffer: Peekable<Lines<BufReader<File>>>,
}

impl YamlFileReader {
    pub fn new(file: File) -> Result<Self, ReaderError> {
        let index = 0;
        let buffer = Self::buffer(&file)?;

        Ok(Self {
            file,
            index,
            buffer,
        })
    }

    // Instantiates a new buffer for a given file.
    fn buffer(file: &File) -> Result<Peekable<Lines<BufReader<File>>>, ReaderError> {
        let mut file_ref_copy = file.try_clone()?;
        file_ref_copy.rewind()?;

        Ok(BufReader::new(file_ref_copy).lines().peekable())
    }

    // Returns the next YAML "doc" in the file, looping back to the start of the
    // file if EOF is encountered.
    fn next_doc(&mut self) -> Result<String, ReaderError> {
        // find start of doc
        loop {
            match self.buffer.peek() {
                Some(Ok(sep)) if sep.trim_end() == YAML_SEPARATOR => {
                    self.buffer.next();
                    break;
                }
                Some(Ok(_)) => {
                    self.buffer.next();
                }
                None => {
                    return Err(ReaderError::Format(
                        "EOF encountered while looking for start of YAML doc".to_string(),
                    ))
                }
                Some(Err(err)) => return Err(ReaderError::Input(err.to_string())),
            }
        }

        // collect lines until end of doc
        let mut doc = String::new();
        loop {
            match self.buffer.peek() {
                Some(Ok(sep)) if sep.trim_end() == YAML_SEPARATOR => {
                    self.index += 1;
                    return Ok(doc);
                }
                Some(Ok(line)) => {
                    // avoid the extra allocation call with a format!
                    let _ = writeln!(doc, "{line}");
                    self.buffer.next();
                }
                None => {
                    // EOF encountered, reset buffer before returning
                    self.index = 0;
                    self.buffer = Self::buffer(&self.file)?;
                    return Ok(doc);
                }
                Some(Err(err)) => return Err(ReaderError::Input(err.to_string())),
            }
        }
    }

    // Reads an object from the YAML file, utilizing validate_fun to validate
    // the object before returning it.
    fn read_impl<T: DeserializeOwned, F: Fn(&T) -> bool>(
        &mut self,
        validate_fun: F,
    ) -> Result<T, ReaderError> {
        let start = self.index;
        loop {
            if let Ok(object) = serde_yaml::from_str::<T>(&self.next_doc()?) {
                if validate_fun(&object) {
                    return Ok(object);
                }
            }
            // exit if all docs have been checked
            if self.index == start {
                return Err(ReaderError::NotFound);
            }
        }
    }
}

impl Reader for YamlFileReader {
    fn read_metadata(&mut self) -> Result<Metadata, ReaderError> {
        self.read_impl(|_: &Metadata| true)
    }

    fn read_state(&mut self, epoch: u32) -> Result<State, ReaderError> {
        self.read_impl(|state: &State| state.epoch == epoch)
    }

    fn read_delta(&mut self, epoch: u32) -> Result<Delta, ReaderError> {
        self.read_impl(|delta: &Delta| delta.epoch == epoch)
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! This module contains the Writer trait to serialize the tool's serde-compatible
//! objects to a format, as well as implementations of the trait.

use serde::Serialize;

/// Interface for writing output generated by the tool.
pub(crate) trait Writer {
    /// Writes a serde serializable object.
    fn write_object(&mut self, object: impl Serialize);

    /// Writes a comment that should be ignored by parsers.
    fn write_comment(&mut self, comment: &str);

    /// Writes a newline.
    fn write_line(&mut self);

    /// Flushes the internal buffer.
    fn flush(&mut self);
}

/// YAML implementor of Writer trait.
pub(crate) mod yaml;


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! This module contains an implementor of the Writer trait for the YAML format.

use std::io::Write;

use serde::Serialize;

use crate::fixture_generator::writer::Writer;

/// YAML format writer.
pub(crate) struct YamlWriter<T: Write> {
    out: T,
}

impl<T: Write> YamlWriter<T> {
    pub fn new(out: T) -> Self {
        Self { out }
    }
}

impl<T: Write> Writer for YamlWriter<T> {
    fn write_object(&mut self, object: impl Serialize) {
        serde_yaml::to_writer(&mut self.out, &object).unwrap();
    }

    fn write_comment(&mut self, comment: &str) {
        let lines = comment.split('\n');
        lines.for_each(|line| writeln!(self.out, "# {line}").unwrap());
    }

    fn write_line(&mut self) {
        writeln!(self.out).unwrap()
    }

    fn flush(&mut self) {
        self.out.flush().unwrap();
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

#[cfg(test)]
mod mysql_tests;
#[cfg(test)]
mod test_util;

#[cfg(test)]
mod memory_tests;


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

use akd::{ecvrf::HardCodedAkdVRF, storage::StorageManager};
use log::info;
type InMemoryDb = akd::storage::memory::AsyncInMemoryDatabase;

#[tokio::test]
async fn test_directory_operations() {
    crate::test_util::log_init(log::Level::Info);

    info!("\n\n******** Starting In-Memory Directory Operations Integration Test ********\n\n");

    let db = InMemoryDb::new();

    let vrf = HardCodedAkdVRF {};
    let storage_manager = StorageManager::new_no_cache(db);
    akd_test_tools::test_suites::directory_test_suite::<_, HardCodedAkdVRF>(
        &storage_manager,
        500,
        &vrf,
    )
    .await;

    info!("\n\n******** Finished In-Memory Directory Operations Integration Test ********\n\n");
}

#[tokio::test]
async fn test_directory_operations_with_caching() {
    crate::test_util::log_init(log::Level::Info);

    info!("\n\n******** Starting In-Memory Directory Operations (w/caching) Integration Test ********\n\n");

    let db = InMemoryDb::new();

    let vrf = HardCodedAkdVRF {};
    let storage_manager = StorageManager::new(db, None, None, None);
    akd_test_tools::test_suites::directory_test_suite::<_, HardCodedAkdVRF>(
        &storage_manager,
        500,
        &vrf,
    )
    .await;

    info!("\n\n******** Finished In-Memory Directory Operations (w/caching) Integration Test ********\n\n");
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

use akd::ecvrf::HardCodedAkdVRF;
use akd::storage::StorageManager;
use akd_mysql::mysql::*;
use log::{error, info, warn};

#[tokio::test]
#[serial_test::serial]
async fn test_directory_operations() {
    crate::test_util::log_init(log::Level::Info);

    info!("\n\n******** Starting MySQL Directory Operations Integration Test ********\n\n");

    if AsyncMySqlDatabase::test_guard() {
        // create the "test" database
        if let Err(error) = AsyncMySqlDatabase::create_test_db(
            "localhost",
            Option::from("root"),
            Option::from("example"),
            Option::from(8001),
        )
        .await
        {
            panic!("Error creating test database: {}", error);
        }

        // connect to the newly created test db
        let mysql_db = AsyncMySqlDatabase::new(
            "localhost",
            "test_db",
            Option::from("root"),
            Option::from("example"),
            Option::from(8001),
            200,
        )
        .await
        .expect("Failed to create async mysql db");

        // delete all data from the db
        if let Err(error) = mysql_db.delete_data().await {
            error!("Error cleaning mysql prior to test suite: {}", error);
        }

        let vrf = HardCodedAkdVRF {};
        let storage_manager = StorageManager::new_no_cache(mysql_db.clone());
        akd_test_tools::test_suites::directory_test_suite::<_, HardCodedAkdVRF>(
            &storage_manager,
            50,
            &vrf,
        )
        .await;

        storage_manager.log_metrics(log::Level::Trace).await;

        // clean the test infra
        if let Err(mysql_async::Error::Server(error)) = storage_manager.get_db().drop_tables().await
        {
            error!(
                "ERROR: Failed to clean MySQL test database with error {}",
                error
            );
        }
    } else {
        warn!("WARN: Skipping MySQL test due to test guard noting that the docker container appears to not be running.");
    }

    info!("\n\n******** Completed MySQL Directory Operations Integration Test ********\n\n");
}

#[tokio::test]
#[serial_test::serial]
async fn test_directory_operations_with_caching() {
    crate::test_util::log_init(log::Level::Info);

    info!("\n\n******** Starting MySQL Directory Operations (w/caching) Integration Test ********\n\n");

    if AsyncMySqlDatabase::test_guard() {
        // create the "test" database
        if let Err(error) = AsyncMySqlDatabase::create_test_db(
            "localhost",
            Option::from("root"),
            Option::from("example"),
            Option::from(8001),
        )
        .await
        {
            panic!("Error creating test database: {}", error);
        }

        // connect to the newly created test db
        let mysql_db = AsyncMySqlDatabase::new(
            "localhost",
            "test_db",
            Option::from("root"),
            Option::from("example"),
            Option::from(8001),
            200,
        )
        .await
        .expect("Failed to create async mysql db");

        // delete all data from the db
        if let Err(error) = mysql_db.delete_data().await {
            error!("Error cleaning mysql prior to test suite: {}", error);
        }

        let vrf = HardCodedAkdVRF {};
        let storage_manager = StorageManager::new(mysql_db.clone(), None, None, None);
        akd_test_tools::test_suites::directory_test_suite::<_, HardCodedAkdVRF>(
            &storage_manager,
            50,
            &vrf,
        )
        .await;

        storage_manager.log_metrics(log::Level::Trace).await;

        // clean the test infra
        if let Err(mysql_async::Error::Server(error)) = storage_manager.get_db().drop_tables().await
        {
            error!(
                "ERROR: Failed to clean MySQL test database with error {}",
                error
            );
        }
    } else {
        warn!("WARN: Skipping MySQL test due to test guard noting that the docker container appears to not be running.");
    }

    info!("\n\n******** Completed MySQL Directory Operations (w/caching) Integration Test ********\n\n");
}

#[tokio::test]
#[serial_test::serial]
async fn test_lookups() {
    crate::test_util::log_init(log::Level::Info);

    info!("\n\n******** Starting MySQL Lookup Tests ********\n\n");

    if AsyncMySqlDatabase::test_guard() {
        // create the "test" database
        if let Err(error) = AsyncMySqlDatabase::create_test_db(
            "localhost",
            Option::from("root"),
            Option::from("example"),
            Option::from(8001),
        )
        .await
        {
            panic!("Error creating test database: {}", error);
        }

        // connect to the newly created test db
        let mysql_db = AsyncMySqlDatabase::new(
            "localhost",
            "test_db",
            Option::from("root"),
            Option::from("example"),
            Option::from(8001),
            200,
        )
        .await
        .expect("Failed to create async mysql db");

        // delete all data from the db
        if let Err(error) = mysql_db.delete_data().await {
            error!("Error cleaning mysql prior to test suite: {}", error);
        }

        let vrf = HardCodedAkdVRF {};
        let storage_manager = StorageManager::new(mysql_db, None, None, None);

        crate::test_util::test_lookups::<_, HardCodedAkdVRF>(&storage_manager, &vrf, 50, 5, 100)
            .await;

        // clean the test infra
        if let Err(mysql_async::Error::Server(error)) = storage_manager.get_db().drop_tables().await
        {
            error!(
                "ERROR: Failed to clean MySQL test database with error {}",
                error
            );
        }
    } else {
        warn!("WARN: Skipping MySQL test due to test guard noting that the docker container appears to not be running.");
    }

    info!("\n\n******** Completed MySQL Lookup Tests ********\n\n");
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

extern crate thread_id;

use akd::ecvrf::VRFKeyStorage;
use akd::storage::{Database, StorageManager};
use akd::Directory;
use akd::{AkdLabel, AkdValue};
use log::{info, Level, Metadata, Record};
use once_cell::sync::OnceCell;
use rand::distributions::Alphanumeric;
use rand::seq::IteratorRandom;
use rand::{thread_rng, Rng};
use std::fs::File;
use std::io;
use std::io::Write;
use std::path::Path;
use std::sync::Mutex;
use tokio::time::{Duration, Instant};

static EPOCH: OnceCell<Instant> = OnceCell::new();

static LOG: OnceCell<u64> = OnceCell::new();

// ================== Logging ================== //

pub(crate) fn log_init(level: Level) {
    EPOCH.get_or_init(Instant::now);
    LOG.get_or_init(|| {
        if let Ok(logger) = FileLogger::new(String::from("integration_test.log")) {
            let loggers: Vec<Box<dyn log::Log>> = vec![Box::new(logger)];
            let mlogger = multi_log::MultiLogger::new(loggers);

            log::set_max_level(level.to_level_filter());
            if let Err(error) = log::set_boxed_logger(Box::new(mlogger)) {
                panic!("Error initializing multi-logger {}", error);
            }
        } else {
            panic!("Error creating file logger!");
        }
        0
    });
}

pub(crate) fn format_log_record(io: &mut (dyn Write + Send), record: &Record) {
    let target = {
        if let Some(target_str) = record.target().split(':').last() {
            if let Some(line) = record.line() {
                format!(" ({target_str}:{line})")
            } else {
                format!(" ({target_str})")
            }
        } else {
            "".to_string()
        }
    };

    let toc = if let Some(epoch) = EPOCH.get() {
        Instant::now() - *epoch
    } else {
        Duration::from_millis(0)
    };

    let seconds = toc.as_secs();
    let hours = seconds / 3600;
    let minutes = (seconds / 60) % 60;
    let seconds = seconds % 60;
    let miliseconds = toc.subsec_millis();

    let msg = format!(
        "[{:02}:{:02}:{:02}.{:03}] ({:x}) {:6} {}{}",
        hours,
        minutes,
        seconds,
        miliseconds,
        thread_id::get(),
        record.level(),
        record.args(),
        target
    );
    let _ = writeln!(io, "{msg}");
}

pub(crate) struct FileLogger {
    sink: Mutex<File>,
}

impl FileLogger {
    pub(crate) fn new<T: AsRef<Path>>(path: T) -> io::Result<Self> {
        let file = File::create(path)?;
        Ok(Self {
            sink: Mutex::new(file),
        })
    }
}

impl log::Log for FileLogger {
    fn enabled(&self, _metadata: &Metadata) -> bool {
        // use the global log-level
        true
    }

    fn log(&self, record: &Record) {
        if !self.enabled(record.metadata()) {
            return;
        }
        let mut sink = &*self.sink.lock().unwrap();
        format_log_record(&mut sink, record);
    }

    fn flush(&self) {
        let _ = std::io::stdout().flush();
    }
}

// ================== Test Helpers ================== //

pub(crate) async fn test_lookups<S: Database + 'static, V: VRFKeyStorage>(
    mysql_db: &StorageManager<S>,
    vrf: &V,
    num_users: u64,
    num_epochs: u64,
    num_lookups: usize,
) {
    // generate the test data
    let mut rng = thread_rng();

    let mut users: Vec<String> = vec![];
    for _ in 0..num_users {
        users.push(
            thread_rng()
                .sample_iter(&Alphanumeric)
                .take(30)
                .map(char::from)
                .collect(),
        );
    }

    // create & test the directory
    let maybe_dir = Directory::<_, _>::new(mysql_db.clone(), vrf.clone()).await;
    match maybe_dir {
        Err(akd_error) => panic!("Error initializing directory: {:?}", akd_error),
        Ok(dir) => {
            info!("AKD Directory started. Beginning tests");

            // Publish `num_epochs` epochs of user material
            for i in 1..=num_epochs {
                let mut data = Vec::new();
                for value in users.iter() {
                    data.push((
                        AkdLabel::from(value),
                        AkdValue(format!("{i}").as_bytes().to_vec()),
                    ));
                }

                if let Err(error) = dir.publish(data).await {
                    panic!("Error publishing batch {:?}", error);
                } else {
                    info!("Published epoch {}", i);
                }
            }

            // Perform `num_lookup` random lookup proofs on the published users

            // Pick a set of users to lookup
            let mut labels = Vec::new();
            for user in users.iter().choose_multiple(&mut rng, num_lookups) {
                let label = AkdLabel::from(user);
                labels.push(label);
            }

            log::warn!("Metrics after publish(es).");
            reset_mysql_db::<S>(mysql_db).await;

            let start = Instant::now();
            // Lookup selected users one by one
            for label in labels.clone() {
                match dir.lookup(label.clone()).await {
                    Err(error) => panic!("Error looking up user information {:?}", error),
                    Ok((proof, root_hash)) => {
                        let vrf_pk = dir.get_public_key().await.unwrap();
                        if let Err(error) = akd::client::lookup_verify(
                            vrf_pk.as_bytes(),
                            root_hash.hash(),
                            label,
                            proof,
                        ) {
                            panic!("Lookup proof failed to verify {:?}", error);
                        }
                    }
                }
            }
            log::warn!(
                "Individual {} lookups took {}ms.",
                num_lookups,
                start.elapsed().as_millis()
            );

            log::warn!("Metrics after individual lookups:");
            reset_mysql_db::<S>(mysql_db).await;

            let start = Instant::now();
            // Bulk lookup selected users
            match dir.batch_lookup(&labels).await {
                Err(error) => panic!("Error batch looking up user information {:?}", error),
                Ok((proofs, root_hash)) => {
                    assert_eq!(labels.len(), proofs.len());

                    let vrf_pk = dir.get_public_key().await.unwrap();
                    for i in 0..proofs.len() {
                        let label = labels[i].clone();
                        let proof = proofs[i].clone();
                        if let Err(error) = akd::client::lookup_verify(
                            vrf_pk.as_bytes(),
                            root_hash.hash(),
                            label,
                            proof,
                        ) {
                            panic!("Batch lookup failed to verify for index {} {:?}", i, error);
                        }
                    }
                }
            }
            log::warn!(
                "Bulk {} lookups took {}ms.",
                num_lookups,
                start.elapsed().as_millis()
            );

            log::warn!("Metrics after lookup proofs: ");
            reset_mysql_db::<S>(mysql_db).await;
        }
    }
}

// Reset MySQL database by logging metrics which resets the metrics, and flushing cache.
// These allow us to accurately assess the additional efficiency of
// bulk lookup proofs.
async fn reset_mysql_db<S: Database>(mysql_db: &StorageManager<S>) {
    mysql_db.log_metrics(Level::Warn).await;
    mysql_db.flush_cache().await;
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

use crate::directory_host::DirectoryCommand;
use colored::*;

pub(crate) enum Command {
    Help,
    Exit,
    Flush,
    Info,
    Directory(DirectoryCommand),
    InvalidArgs(String),
    Unknown(String),
}

impl Command {
    pub(crate) fn parse(text: &mut String) -> Command {
        trim_newline(text);
        let parts: Vec<&str> = text.split(' ').collect();

        let mut command = String::new();
        if let Some(head) = parts.first() {
            command = String::from(*head);
        }

        match command.to_lowercase().as_ref() {
            "exit" | "x" => Command::Exit,
            "help" | "?" => Command::Help,
            "flush" => Command::Flush,
            "info" => Command::Info,
            cmd => Command::handle_dir_cmd(cmd, parts, text),
        }
    }

    pub(crate) fn print_help_menu() {
        println!(
            "{}",
            "*************************** Help menu ***************************".red()
        );
        println!(
            "{} are commands, {} are mandatory args, {} are optional args",
            "green".green(),
            "blue".blue(),
            "magenta".magenta()
        );
        println!("=============================================================");
        println!("  {}|{}:\t\t\tprint this menu", "help".green(), "?".green());
        println!(
            "  {}|{}:\t\t\texit the application",
            "exit".green(),
            "x".green()
        );
        println!("  {}\t\t\t\tflush the database entries", "flush".green());
        println!(
            "  {}\t\t\t\tprints information about the running instance",
            "info".green()
        );
        println!(
            "  {} {} {}:\t\tpublish key material (value) for user",
            "publish".green(),
            "user".blue(),
            "value".blue()
        );
        println!(
            "  {} {}:\t\t\tlookup a proof for user",
            "lookup".green(),
            "user".blue()
        );
        println!(
            "  {} {}:\t\t\tlookup key history for user",
            "history".green(),
            "user".blue()
        );
        println!(
            "  {} {} {}:\t\tretrieve audit proof between start and end epochs",
            "audit".green(),
            "start".blue(),
            "end".blue()
        );
        println!(
            "  {}|{} {}:\t\tretrieve the root hash at given epoch (default = latest epoch)",
            "root".green(),
            "root_hash".green(),
            "epoch".magenta()
        );
    }

    // ==== Helpers for managing directory commands ==== //
    fn handle_dir_cmd(command: &str, parts: Vec<&str>, full_text: &str) -> Command {
        let dir_cmd: Option<Option<DirectoryCommand>> = match command {
            "publish" => Some(Command::publish(parts)),
            "lookup" => Some(Command::lookup(parts)),
            "history" => Some(Command::history(parts)),
            "audit" => Some(Command::audit(parts)),
            "root" | "root_hash" => Some(Command::root_hash(parts)),
            _ => None,
        };
        match dir_cmd {
            Some(Some(cmd)) => Command::Directory(cmd),
            Some(None) => {
                let msg = format!(
                    "Command {} received invalid argments. Check {} for syntax",
                    command,
                    "help".green()
                );
                Command::InvalidArgs(msg)
            }
            _ => Command::Unknown(String::from(full_text)),
        }
    }

    fn publish(parts: Vec<&str>) -> Option<DirectoryCommand> {
        if parts.len() < 3 {
            return None;
        }
        let (a, b) = (parts[1], parts[2]);
        let cmd = DirectoryCommand::Publish(String::from(a), String::from(b));
        Some(cmd)
    }

    fn lookup(parts: Vec<&str>) -> Option<DirectoryCommand> {
        if parts.len() < 2 {
            return None;
        }
        let a = parts[1];
        let cmd = DirectoryCommand::Lookup(String::from(a));
        Some(cmd)
    }

    fn history(parts: Vec<&str>) -> Option<DirectoryCommand> {
        if parts.len() < 2 {
            return None;
        }
        let a = parts[1];
        let cmd = DirectoryCommand::KeyHistory(String::from(a));
        Some(cmd)
    }

    fn audit(parts: Vec<&str>) -> Option<DirectoryCommand> {
        if parts.len() < 3 {
            return None;
        }
        let (a, b) = (parts[1], parts[2]);
        match (a.parse::<u64>(), b.parse::<u64>()) {
            (Ok(u_a), Ok(u_b)) => {
                let cmd = DirectoryCommand::Audit(u_a, u_b);
                Some(cmd)
            }
            _ => None,
        }
    }

    fn root_hash(_parts: Vec<&str>) -> Option<DirectoryCommand> {
        let cmd = DirectoryCommand::RootHash;
        Some(cmd)
    }
}

fn trim_newline(s: &mut String) {
    if s.ends_with('\n') {
        s.pop();
        if s.ends_with('\r') {
            s.pop();
        }
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

use akd::ecvrf::VRFKeyStorage;
use akd::errors::AkdError;
use akd::storage::types::*;
use akd::storage::{Database, StorageManager};
use akd::HistoryParams;
use akd::{AkdLabel, AkdValue, Digest};
use akd::{Directory, EpochHash};
use log::{debug, error, info};
use std::marker::{Send, Sync};
use tokio::sync::mpsc::*;
use tokio::time::Instant;

pub(crate) struct Rpc(
    pub(crate) DirectoryCommand,
    pub(crate) Option<tokio::sync::oneshot::Sender<Result<String, String>>>,
);

#[derive(Debug)]
pub enum DirectoryCommand {
    Publish(String, String),
    PublishBatch(Vec<(String, String)>),
    Lookup(String),
    KeyHistory(String),
    Audit(u64, u64),
    RootHash,
    Terminate,
}

pub(crate) async fn init_host<S, V>(rx: &mut Receiver<Rpc>, directory: &mut Directory<S, V>)
where
    S: Database + 'static,
    V: VRFKeyStorage,
{
    info!("Starting the verifiable directory host");

    while let Some(Rpc(message, channel)) = rx.recv().await {
        match (message, channel) {
            (DirectoryCommand::Terminate, _) => {
                break;
            }
            (DirectoryCommand::Publish(a, b), Some(response)) => {
                let tic = Instant::now();
                match directory
                    .publish(vec![(AkdLabel::from(&a), AkdValue::from(&b))])
                    .await
                {
                    Ok(EpochHash(epoch, hash)) => {
                        let toc = Instant::now() - tic;
                        let msg = format!(
                            "PUBLISHED '{}' = '{}' in {} s (epoch: {}, root hash: {})",
                            a,
                            b,
                            toc.as_secs_f64(),
                            epoch,
                            hex::encode(hash)
                        );
                        response.send(Ok(msg)).unwrap()
                    }
                    Err(error) => {
                        let msg = format!("Failed to publish with error: {error:?}");
                        response.send(Err(msg)).unwrap();
                    }
                }
            }
            (DirectoryCommand::PublishBatch(batches), Some(response)) => {
                let tic = Instant::now();
                let len = batches.len();
                match directory
                    .publish(
                        batches
                            .into_iter()
                            .map(|(key, value)| (AkdLabel::from(&key), AkdValue::from(&value)))
                            .collect(),
                    )
                    .await
                {
                    Ok(_) => {
                        let toc = Instant::now() - tic;
                        let msg = format!("PUBLISHED {} records in {} s", len, toc.as_secs_f64());
                        response.send(Ok(msg)).unwrap()
                    }
                    Err(error) => {
                        let msg = format!("Failed to publish with error: {error:?}");
                        response.send(Err(msg)).unwrap();
                    }
                }
            }
            (DirectoryCommand::Lookup(a), Some(response)) => {
                match directory.lookup(AkdLabel::from(&a)).await {
                    Ok((proof, root_hash)) => {
                        let vrf_pk = directory.get_public_key().await.unwrap();
                        let verification = akd::client::lookup_verify(
                            vrf_pk.as_bytes(),
                            root_hash.hash(),
                            AkdLabel::from(&a),
                            proof,
                        );
                        if verification.is_err() {
                            let msg = format!("WARN: Lookup proof failed verification for '{a}'");
                            response.send(Err(msg)).unwrap();
                        } else {
                            let msg = format!("Lookup proof verified for user '{a}'");
                            response.send(Ok(msg)).unwrap();
                        }
                    }
                    Err(error) => {
                        let msg = format!("Failed to lookup with error {error:?}");
                        response.send(Err(msg)).unwrap();
                    }
                }
            }
            (DirectoryCommand::KeyHistory(a), Some(response)) => {
                match directory
                    .key_history(&AkdLabel::from(&a), HistoryParams::default())
                    .await
                {
                    Ok(_proof) => {
                        let msg = format!("GOT KEY HISTORY FOR '{a}'");
                        response.send(Ok(msg)).unwrap();
                    }
                    Err(error) => {
                        let msg = format!("Failed to lookup with error {error:?}");
                        response.send(Err(msg)).unwrap();
                    }
                }
            }
            (DirectoryCommand::Audit(start, end), Some(response)) => {
                match directory.audit(start, end).await {
                    Ok(_proof) => {
                        let msg = format!("GOT AUDIT PROOF BETWEEN ({start}, {end})");
                        response.send(Ok(msg)).unwrap();
                    }
                    Err(error) => {
                        let msg = format!("Failed to get audit proof with error {error:?}");
                        response.send(Err(msg)).unwrap();
                    }
                }
            }
            (DirectoryCommand::RootHash, Some(response)) => {
                let hash = directory.get_epoch_hash().await;
                match hash {
                    Ok(EpochHash(_, hash)) => {
                        let msg = format!("Retrieved root hash {}", hex::encode(hash));
                        response.send(Ok(msg)).unwrap();
                    }
                    Err(error) => {
                        let msg = format!("Failed to retrieve root hash with error {error:?}");
                        response.send(Err(msg)).unwrap();
                    }
                }
            }
            (_, None) => {
                error!("A channel was not provided to the directory server to process a command!");
            }
        }
    }

    info!("AKD host shutting down");
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

extern crate thread_id;

use colored::*;
use log::{Level, LevelFilter, Metadata, Record};
use once_cell::sync::Lazy;
use once_cell::sync::OnceCell;
use tokio::time::{Duration, Instant};

use std::fs::File;
use std::io;
use std::io::Write;
use std::path::Path;
use std::sync::Mutex;

static EPOCH: OnceCell<Instant> = OnceCell::new();

pub(crate) struct ConsoleLogger {
    pub(crate) level: Level,
}

impl ConsoleLogger {
    pub(crate) fn touch() {
        EPOCH.get_or_init(Instant::now);
    }

    pub(crate) fn format_log_record(io: &mut (dyn Write + Send), record: &Record, no_color: bool) {
        let target = {
            if let Some(target_str) = record.target().split(':').last() {
                if let Some(line) = record.line() {
                    format!(" ({target_str}:{line})")
                } else {
                    format!(" ({target_str})")
                }
            } else {
                "".to_string()
            }
        };

        let toc = if let Some(epoch) = EPOCH.get() {
            Instant::now() - *epoch
        } else {
            Duration::from_millis(0)
        };

        let seconds = toc.as_secs();
        let hours = seconds / 3600;
        let minutes = (seconds / 60) % 60;
        let seconds = seconds % 60;
        let miliseconds = toc.subsec_millis();

        let msg = format!(
            "[{:02}:{:02}:{:02}.{:03}] ({:x}) {:6} {}{}",
            hours,
            minutes,
            seconds,
            miliseconds,
            thread_id::get(),
            record.level(),
            record.args(),
            target
        );
        if no_color {
            let _ = writeln!(io, "{msg}");
        } else {
            let msg = match record.level() {
                Level::Trace | Level::Debug => msg.white(),
                Level::Info => msg.blue(),
                Level::Warn => msg.yellow(),
                Level::Error => msg.red(),
            };
            let _ = writeln!(io, "{msg}");
        }
    }
}

impl log::Log for ConsoleLogger {
    fn enabled(&self, metadata: &Metadata) -> bool {
        metadata.level() <= self.level
    }

    fn log(&self, record: &Record) {
        if !self.enabled(record.metadata()) {
            return;
        }
        let mut io = std::io::stdout();
        ConsoleLogger::format_log_record(&mut io, record, false);
    }

    fn flush(&self) {
        let _ = std::io::stdout().flush();
    }
}

pub(crate) struct FileLogger {
    sink: Mutex<File>,
}

impl FileLogger {
    pub(crate) fn new<T: AsRef<Path>>(path: T) -> io::Result<Self> {
        let file = File::create(path)?;
        Ok(Self {
            sink: Mutex::new(file),
        })
    }
}

impl log::Log for FileLogger {
    fn enabled(&self, metadata: &Metadata) -> bool {
        // use the global log-level
        true
    }

    fn log(&self, record: &Record) {
        if !self.enabled(record.metadata()) {
            return;
        }
        let mut sink = &*self.sink.lock().unwrap();
        ConsoleLogger::format_log_record(&mut sink, record, true);
    }

    fn flush(&self) {
        let _ = std::io::stdout().flush();
    }
}