#![allow(clippy::needless_return)]

use std::{collections::HashMap, path::PathBuf, sync::Arc};

use persia_libs::{
    indexmap,
    once_cell::sync::OnceCell,
    serde::{self, Deserialize, Serialize},
    serde_yaml,
    thiserror::Error,
    tracing,
};

use persia_speedy::{Readable, Writable};

#[derive(Readable, Writable, Error, Debug, Clone)]
pub enum PersiaGlobalConfigError {
    #[error("global config not ready error")]
    NotReadyError,
    #[error("global config set error")]
    SetError,
    #[error("failed to open/read config file error")]
    ConfigFileError(String),
}

#[derive(Serialize, Deserialize, Readable, Writable, Debug, Default, Clone)]
#[serde(crate = "self::serde")]
pub struct BoundedUniformInitialization {
    pub lower: f32,
    pub upper: f32,
}

impl BoundedUniformInitialization {
    pub fn new(lower: f32, upper: f32) -> Self {
        BoundedUniformInitialization { lower, upper }
    }
}

#[derive(Serialize, Deserialize, Readable, Writable, Debug, Default, Clone)]
#[serde(crate = "self::serde")]
pub struct BoundedGammaInitialization {
    pub shape: f32,
    pub scale: f32,
}

impl BoundedGammaInitialization {
    pub fn new(shape: f32, scale: f32) -> Self {
        BoundedGammaInitialization { shape, scale }
    }
}

#[derive(Serialize, Deserialize, Readable, Writable, Debug, Default, Clone)]
#[serde(crate = "self::serde")]
pub struct BoundedPoissonInitialization {
    pub lambda: f32,
}

impl BoundedPoissonInitialization {
    pub fn new(lambda: f32) -> Self {
        BoundedPoissonInitialization { lambda }
    }
}

#[derive(Serialize, Deserialize, Readable, Writable, Debug, Default, Clone)]
#[serde(crate = "self::serde")]
pub struct BoundedNormalInitialization {
    pub mean: f32,
    pub standard_deviation: f32,
}

impl BoundedNormalInitialization {
    pub fn new(mean: f32, standard_deviation: f32) -> Self {
        BoundedNormalInitialization {
            mean,
            standard_deviation,
        }
    }
}

#[derive(Serialize, Deserialize, Readable, Writable, Debug, Clone)]
#[serde(crate = "self::serde")]
pub enum InitializationMethod {
    InverseEmbeddingSizeSqrt,
    BoundedUniform(BoundedUniformInitialization),
    BoundedGamma(BoundedGammaInitialization),
    BoundedPoisson(BoundedPoissonInitialization),
    BoundedNormal(BoundedNormalInitialization),
}

impl Default for InitializationMethod {
    fn default() -> Self {
        Self::BoundedUniform(BoundedUniformInitialization {
            lower: -0.01,
            upper: 0.01,
        })
    }
}

#[derive(Readable, Writable, Debug, Clone)]
pub struct PersiaEmbeddingModelHyperparameters {
    pub initialization_method: InitializationMethod,
    pub admit_probability: f32,
    pub weight_bound: f32,
    pub enable_weight_bound: bool,
}

static PERSIA_EMBEDDING_SEVER_CONFIG: OnceCell<Arc<EmbeddingParameterServerConfig>> =
    OnceCell::new();

static PERSIA_EMBEDDING_WORKER_CONFIG: OnceCell<Arc<EmbeddingWorkerConfig>> = OnceCell::new();

static PERSIA_COMMON_CONFIG: OnceCell<Arc<PersiaCommonConfig>> = OnceCell::new();

static PERSIA_REPLICA_INFO: OnceCell<Arc<PersiaReplicaInfo>> = OnceCell::new();

static PERSIA_INSTANCE_INFO: OnceCell<Arc<InstanceInfo>> = OnceCell::new();

static PERSIA_EMBEDDING_CONFIG: OnceCell<Arc<EmbeddingConfig>> = OnceCell::new();

/// get local ip address by specific socket name or
/// request the 8.8.8.8:80 to retrieve the ip address from response.
fn get_local_ip() -> String {
    let persia_socket_ifname =
        std::env::var("PERSIA_SOCKET_IFNAME").unwrap_or(String::from("eth0"));

    let local_ip_addr = match get_if_addrs::get_if_addrs() {
        Ok(addr_interfaces) => {
            if let Some(ip_addr) = addr_interfaces
                .iter()
                .find(|x| x.name == persia_socket_ifname.as_str())
            {
                let ip_addr = ip_addr.ip().to_string();
                Some(ip_addr)
            } else {
                None
            }
        }
        Err(_) => None,
    };

    match local_ip_addr {
        Some(ip_addr) => return ip_addr,
        None => {
            match local_ipaddress::get() {
                Some(ip_addr) => return ip_addr,
                None => panic!(
                    "failed to get local ip address by PERSIA_SOCKET_IFNAME {}, please specify a correct PERSIA_SOCKET_IFNAME",
                    persia_socket_ifname
                )
            }
        }
    }
}

#[derive(Deserialize, Serialize, Readable, Writable, Debug, Clone)]
#[serde(crate = "self::serde")]
pub struct InferConfig {
    pub servers: Vec<String>,
    pub embedding_checkpoint: String,
}

impl Default for InferConfig {
    fn default() -> Self {
        Self {
            servers: vec![],
            embedding_checkpoint: String::from(""),
        }
    }
}

#[derive(Deserialize, Serialize, Readable, Writable, Debug, Clone)]
#[serde(crate = "self::serde")]
pub enum PerisaJobType {
    Train,
    Eval,
    Infer,
}

#[derive(Deserialize, Serialize, Readable, Writable, Debug, Clone)]
#[serde(crate = "self::serde")]
pub struct PersiaReplicaInfo {
    pub replica_size: usize,
    pub replica_index: usize,
}

impl PersiaReplicaInfo {
    pub fn get() -> Result<Arc<Self>, PersiaGlobalConfigError> {
        let singleton = PERSIA_REPLICA_INFO.get();
        match singleton {
            Some(s) => Ok(s.clone()),
            None => Err(PersiaGlobalConfigError::NotReadyError),
        }
    }

    pub fn set(replica_size: usize, replica_index: usize) -> Result<(), PersiaGlobalConfigError> {
        let instance = Self {
            replica_size,
            replica_index,
        };
        PERSIA_REPLICA_INFO
            .set(Arc::new(instance))
            .map_err(|_| PersiaGlobalConfigError::SetError)
    }

    pub fn is_master(&self) -> bool {
        self.replica_index == 0
    }
}

#[derive(Deserialize, Serialize, Readable, Writable, Debug, Clone)]
#[serde(crate = "self::serde")]
pub struct InstanceInfo {
    pub ip_address: String,
    pub port: u16,
}

impl InstanceInfo {
    pub fn new(port: u16) -> Self {
        Self {
            ip_address: get_local_ip(),
            port,
        }
    }

    pub fn get() -> Result<Arc<Self>, PersiaGlobalConfigError> {
        let singleton = PERSIA_INSTANCE_INFO.get();
        match singleton {
            Some(s) => Ok(s.clone()),
            None => Err(PersiaGlobalConfigError::NotReadyError),
        }
    }
}

fn get_true() -> bool {
    true
}

fn get_false() -> bool {
    false
}

fn get_four() -> usize {
    4
}

fn get_eight() -> usize {
    8
}

fn get_zero() -> u64 {
    0
}

fn get_ten() -> usize {
    10
}

fn get_hundred() -> usize {
    100
}

fn get_thousand() -> usize {
    1000
}

fn get_million() -> usize {
    1_000_000
}

fn get_billion() -> usize {
    1_000_000_000
}

fn get_default_incremental_dir() -> String {
    String::from("/workspace/incremental_dir/")
}

fn get_default_job_name() -> String {
    String::from("persia_default_jobname")
}

fn get_default_common_config() -> PersiaCommonConfig {
    PersiaCommonConfig::default()
}

fn get_default_embedding_worker_config() -> EmbeddingWorkerConfig {
    EmbeddingWorkerConfig::default()
}

fn get_default_embedding_parameter_server_config() -> EmbeddingParameterServerConfig {
    EmbeddingParameterServerConfig::default()
}

fn get_default_metrics_config() -> PersiaMetricsConfig {
    PersiaMetricsConfig::default()
}

fn get_default_job_type() -> PerisaJobType {
    PerisaJobType::Train
}

fn get_default_infer_config() -> InferConfig {
    InferConfig::default()
}

fn get_default_hashstack_config() -> HashStackConfig {
    HashStackConfig {
        hash_stack_rounds: 0,
        embedding_size: 0,
    }
}

fn get_default_feature_groups() -> indexmap::IndexMap<String, Vec<String>> {
    indexmap::IndexMap::new()
}

fn get_default_checkpointing_config() -> CheckpointingConfig {
    CheckpointingConfig::default()
}

#[derive(Deserialize, Serialize, Readable, Writable, Debug, Clone)]
#[serde(crate = "self::serde")]
pub struct PersiaMetricsConfig {
    #[serde(default = "get_true")]
    pub enable_metrics: bool,
    #[serde(default = "get_default_job_name")]
    pub job_name: String,
    #[serde(default = "get_ten")]
    pub push_interval_seconds: usize,
}

impl Default for PersiaMetricsConfig {
    fn default() -> Self {
        Self {
            enable_metrics: true,
            job_name: get_default_job_name(),
            push_interval_seconds: 10,
        }
    }
}

#[derive(Deserialize, Serialize, Readable, Writable, Debug, Clone)]
#[serde(crate = "self::serde")]
pub struct CheckpointingConfig {
    #[serde(default = "get_four")]
    pub num_workers: usize,
}

impl Default for CheckpointingConfig {
    fn default() -> Self {
        Self { num_workers: 4 }
    }
}

#[derive(Deserialize, Serialize, Readable, Writable, Debug, Clone)]
#[serde(crate = "self::serde")]
pub struct PersiaCommonConfig {
    #[serde(default = "get_default_metrics_config")]
    pub metrics_config: PersiaMetricsConfig,
    #[serde(default = "get_default_job_type")]
    pub job_type: PerisaJobType,
    #[serde(default = "get_default_infer_config")]
    pub infer_config: InferConfig,
    #[serde(default = "get_default_checkpointing_config")]
    pub checkpointing_config: CheckpointingConfig,
}

impl Default for PersiaCommonConfig {
    fn default() -> Self {
        Self {
            metrics_config: PersiaMetricsConfig::default(),
            job_type: PerisaJobType::Train,
            infer_config: InferConfig::default(),
            checkpointing_config: CheckpointingConfig::default(),
        }
    }
}

impl PersiaCommonConfig {
    pub fn get() -> Result<Arc<Self>, PersiaGlobalConfigError> {
        let singleton = PERSIA_COMMON_CONFIG.get();
        match singleton {
            Some(s) => Ok(s.clone()),
            None => Err(PersiaGlobalConfigError::NotReadyError),
        }
    }
}

#[derive(Deserialize, Serialize, Readable, Writable, Debug, Clone)]
#[serde(crate = "self::serde")]
pub struct EmbeddingWorkerConfig {
    #[serde(default = "get_thousand")]
    pub forward_buffer_size: usize,
    #[serde(default = "get_thousand")]
    pub buffered_data_expired_sec: usize,
}

impl Default for EmbeddingWorkerConfig {
    fn default() -> Self {
        Self {
            forward_buffer_size: 1000,
            buffered_data_expired_sec: 1000,
        }
    }
}

impl EmbeddingWorkerConfig {
    pub fn get() -> Result<Arc<Self>, PersiaGlobalConfigError> {
        let singleton = PERSIA_EMBEDDING_WORKER_CONFIG.get();
        match singleton {
            Some(s) => Ok(s.clone()),
            None => Err(PersiaGlobalConfigError::NotReadyError),
        }
    }
}

#[derive(Deserialize, Serialize, Readable, Writable, Debug, Clone)]
#[serde(crate = "self::serde")]
pub struct EmbeddingParameterServerConfig {
    // Eviction map config
    #[serde(default = "get_billion")]
    pub capacity: usize,
    #[serde(default = "get_hundred")]
    pub num_hashmap_internal_shards: usize,
    // incremental dump config
    #[serde(default = "get_false")]
    pub enable_incremental_update: bool,
    #[serde(default = "get_million")]
    pub incremental_buffer_size: usize,
    #[serde(default = "get_default_incremental_dir")]
    pub incremental_dir: String,
    #[serde(default = "get_thousand")]
    pub incremental_channel_capacity: usize,
}

impl Default for EmbeddingParameterServerConfig {
    fn default() -> Self {
        Self {
            capacity: 1_000_000_000,
            num_hashmap_internal_shards: 1000,
            enable_incremental_update: false,
            incremental_buffer_size: 1_000_000,
            incremental_dir: get_default_incremental_dir(),
            incremental_channel_capacity: 1000,
        }
    }
}

impl EmbeddingParameterServerConfig {
    pub fn get() -> Result<Arc<Self>, PersiaGlobalConfigError> {
        let singleton = PERSIA_EMBEDDING_SEVER_CONFIG.get();
        match singleton {
            Some(s) => Ok(s.clone()),
            None => Err(PersiaGlobalConfigError::NotReadyError),
        }
    }
}

#[derive(Deserialize, Serialize, Readable, Writable, Debug, Clone)]
#[serde(crate = "self::serde")]
pub struct PersiaGlobalConfig {
    #[serde(default = "get_default_common_config")]
    pub common_config: PersiaCommonConfig,
    #[serde(default = "get_default_embedding_worker_config")]
    pub embedding_worker_config: EmbeddingWorkerConfig,
    #[serde(default = "get_default_embedding_parameter_server_config")]
    pub embedding_parameter_server_config: EmbeddingParameterServerConfig,
}

impl PersiaGlobalConfig {
    pub fn set_configures(
        file_path: &PathBuf,
        port: u16,
        replica_index: usize,
        replica_size: usize,
    ) -> Result<(), PersiaGlobalConfigError> {
        if !file_path.is_file() {
            tracing::error!("global config yaml file NOT found");
            std::thread::sleep(std::time::Duration::from_secs(120));
            panic!("global config yaml file NOT found")
        }

        let global_config: PersiaGlobalConfig = serde_yaml::from_reader(
            std::fs::File::open(file_path).expect("cannot read config file"),
        )
        .expect("cannot parse config file");

        tracing::info!(
            "setting embedding_parameter_server_config {:?}",
            global_config.embedding_parameter_server_config
        );
        PERSIA_EMBEDDING_SEVER_CONFIG
            .set(Arc::new(global_config.embedding_parameter_server_config))
            .map_err(|_| PersiaGlobalConfigError::SetError)?;

        tracing::info!(
            "setting embedding_worker_config {:?}",
            global_config.embedding_worker_config
        );
        PERSIA_EMBEDDING_WORKER_CONFIG
            .set(Arc::new(global_config.embedding_worker_config))
            .map_err(|_| PersiaGlobalConfigError::SetError)?;

        tracing::info!("setting common_config {:?}", global_config.common_config);
        PERSIA_COMMON_CONFIG
            .set(Arc::new(global_config.common_config))
            .map_err(|_| PersiaGlobalConfigError::SetError)?;

        let instance_info = InstanceInfo::new(port);
        tracing::info!("setting instance_info {:?}", instance_info);
        PERSIA_INSTANCE_INFO
            .set(Arc::new(instance_info))
            .map_err(|_| PersiaGlobalConfigError::SetError)?;

        let replica_info = PersiaReplicaInfo {
            replica_index,
            replica_size,
        };
        tracing::info!("setting replica_info {:?}", replica_info);
        PERSIA_REPLICA_INFO
            .set(Arc::new(replica_info))
            .map_err(|_| PersiaGlobalConfigError::SetError)?;

        Ok(())
    }
}

#[derive(Deserialize, Serialize, Readable, Writable, Debug, Clone)]
#[serde(crate = "self::serde")]
pub struct HashStackConfig {
    pub hash_stack_rounds: usize,
    pub embedding_size: usize,
}

#[derive(Deserialize, Serialize, Readable, Writable, Debug, Clone)]
#[serde(crate = "self::serde")]
pub struct SlotConfig {
    pub dim: usize,
    #[serde(default = "get_ten")]
    pub sample_fixed_size: usize, // raw embedding placeholder size to fill 3d tensor -> (bs, sample_fix_sized, dim)
    #[serde(default = "get_true")]
    pub embedding_summation: bool,
    #[serde(default = "get_false")]
    pub sqrt_scaling: bool,
    #[serde(default = "get_default_hashstack_config")]
    pub hash_stack_config: HashStackConfig,
    // index_prefix: different prefix add to index of different features, to prevent bucket conflict for each feature embedding.
    #[serde(default = "get_zero")]
    pub index_prefix: u64,
}

#[derive(Debug, Serialize, Deserialize, Readable, Writable, Clone)]
#[serde(crate = "self::serde")]
pub struct EmbeddingConfig {
    #[serde(default = "get_eight")]
    pub feature_index_prefix_bit: usize,
    pub slots_config: indexmap::IndexMap<String, SlotConfig>,
    #[serde(default = "get_default_feature_groups")]
    pub feature_groups: indexmap::IndexMap<String, Vec<String>>,
}

impl EmbeddingConfig {
    pub fn set(file_path: &PathBuf) -> Result<(), PersiaGlobalConfigError> {
        if !file_path.is_file() {
            tracing::error!("embedding config yaml file NOT found");
            std::thread::sleep(std::time::Duration::from_secs(120));
            panic!("embedding config yaml file NOT found")
        }

        let embedding_config: EmbeddingConfig = serde_yaml::from_reader(
            std::fs::File::open(file_path).expect("cannot read config file"),
        )
        .expect("cannot parse config file");

        let embedding_config = parse_embedding_config(embedding_config);

        tracing::info!("setting embedding_config {:?}", embedding_config,);
        PERSIA_EMBEDDING_CONFIG
            .set(Arc::new(embedding_config))
            .map_err(|_| PersiaGlobalConfigError::SetError)?;

        Ok(())
    }

    pub fn get() -> Result<Arc<Self>, PersiaGlobalConfigError> {
        let singleton = PERSIA_EMBEDDING_CONFIG.get();
        match singleton {
            Some(s) => Ok(s.clone()),
            None => Err(PersiaGlobalConfigError::NotReadyError),
        }
    }

    pub fn get_slot_by_feature_name(&self, feature_name: &str) -> &SlotConfig {
        self.slots_config
            .get(feature_name)
            .expect(format!("slot: {} not found", feature_name).as_str())
    }
}

pub fn parse_embedding_config(config: EmbeddingConfig) -> EmbeddingConfig {
    let mut config = config;
    let slots_config = &mut config.slots_config;
    let feature_groups = &mut config.feature_groups;
    let mut slot_name_to_feature_group = HashMap::new();

    feature_groups
        .iter()
        .for_each(|(feature_group_name, slot_names)| {
            slot_names.iter().for_each(|slot_name| {
                slot_name_to_feature_group.insert(slot_name.clone(), feature_group_name.clone());
            })
        });

    slots_config.iter().for_each(|(slot_name, _)| {
        if !slot_name_to_feature_group.contains_key(slot_name) {
            let res = feature_groups.insert(slot_name.clone(), vec![slot_name.clone()]);
            if res.is_some() {
                panic!("a slot name can not same with feature group name");
            }
            slot_name_to_feature_group.insert(slot_name.clone(), slot_name.clone());
        }
    });

    assert_ne!(
        config.feature_index_prefix_bit, 0,
        "feature_index_prefix_bit must > 0"
    );
    let feature_prefix_bias = u64::BITS - config.feature_index_prefix_bit as u32;
    slots_config
        .iter_mut()
        .for_each(|(slot_name, slot_config)| {
            assert_eq!(
                slot_config.index_prefix, 0,
                "please do not set index_prefix manually"
            );
            let feature_group_name = slot_name_to_feature_group
                .get(slot_name)
                .expect("feature group not found");
            let feature_group_index = feature_groups
                .get_index_of(feature_group_name)
                .expect("feature group not found");
            slot_config.index_prefix = num_traits::CheckedShl::checked_shl(
                &(feature_group_index as u64 + 1),
                feature_prefix_bias,
            )
            .expect("slot index_prefix overflow, please try a bigger feature_index_prefix_bit");
        });

    config
}
