/* mCaptcha - A proof of work based DoS protection system
 * Copyright © 2021 Aravinth Manivannan <realravinth@batsense.net>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
use std::collections::HashMap;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use std::time::Duration;

use dashmap::DashMap;
use serde::{Deserialize, Serialize};

use super::defense::Defense;
use libmcaptcha::errors::*;
use libmcaptcha::master::messages as ManagerMessages;

/// Builder for [MCaptcha]
#[derive(Clone, Serialize, Deserialize, Debug)]
pub struct MCaptchaBuilder {
    visitor_threshold: u32,
    defense: Option<Defense>,
    duration: Option<u64>,
}

impl Default for MCaptchaBuilder {
    fn default() -> Self {
        MCaptchaBuilder {
            visitor_threshold: 0,
            defense: None,
            duration: None,
        }
    }
}

impl MCaptchaBuilder {
    /// set defense
    pub fn defense(&mut self, d: Defense) -> &mut Self {
        self.defense = Some(d);
        self
    }

    /// set duration
    pub fn duration(&mut self, d: u64) -> &mut Self {
        self.duration = Some(d);
        self
    }

    /// Builds new [MCaptcha]
    pub fn build(self: &mut MCaptchaBuilder) -> CaptchaResult<MCaptcha> {
        if self.duration.is_none() {
            Err(CaptchaError::PleaseSetValue("duration".into()))
        } else if self.defense.is_none() {
            Err(CaptchaError::PleaseSetValue("defense".into()))
        } else if self.duration <= Some(0) {
            Err(CaptchaError::CaptchaDurationZero)
        } else {
            let m = MCaptcha {
                duration: self.duration.unwrap(),
                defense: self.defense.clone().unwrap(),
                visitor_threshold: Arc::new(AtomicU32::new(self.visitor_threshold)),
            };
            Ok(m)
        }
    }
}

#[derive(Clone, Serialize, Deserialize, Debug)]
pub struct MCaptcha {
    visitor_threshold: Arc<AtomicU32>,
    defense: Defense,
    duration: u64,
}

impl MCaptcha {
    /// increments the visitor count by one
    #[inline]
    pub fn add_visitor(&self) -> u32 {
        //        self.visitor_threshold += 1;
        let current_visitor_level = self.visitor_threshold.fetch_add(1, Ordering::SeqCst) + 1;
        let current_level = self.defense.current_level(current_visitor_level);
        current_level.difficulty_factor
    }

    /// decrements the visitor count by specified count
    #[inline]
    pub fn set_visitor_count(&self, new_current: u32) {
        self.visitor_threshold
            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |mut current| {
                if current != new_current {
                    Some(new_current)
                } else {
                    None
                }
            });
    }

    /// decrements the visitor count by specified count
    #[inline]
    pub fn decrement_visitor_by(&self, count: u32) {
        self.visitor_threshold
            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |mut current| {
                if current > 0 {
                    if current >= count {
                        current -= count;
                    } else {
                        current = 0;
                    }
                    Some(current)
                } else {
                    None
                }
            });
    }

    /// get [Counter]'s current visitor_threshold
    pub fn get_visitors(&self) -> u32 {
        self.visitor_threshold.load(Ordering::SeqCst)
    }
}

#[derive(Clone, Serialize, Deserialize)]
pub struct Manager {
    pub captchas: Arc<DashMap<String, Arc<MCaptcha>>>,
    pub gc: u64,
}

impl Manager {
    /// add [Counter] actor to [Manager]
    pub fn add_captcha(&self, m: Arc<MCaptcha>, id: String) {
        self.captchas.insert(id, m);
    }

    /// create new master
    /// accepts a `u64` to configure garbage collection period
    pub fn new(gc: u64) -> Self {
        Manager {
            captchas: Arc::new(DashMap::new()),
            gc,
        }
    }

    fn gc(captchas: Arc<DashMap<String, Arc<MCaptcha>>>) {
        for captcha in captchas.iter() {
            let visitor = { captcha.value().get_visitors() };
            if visitor == 0 {
                captchas.remove(captcha.key());
            }
        }
    }

    /// get [Counter] actor from [Manager]
    pub fn get_captcha(&self, id: &str) -> Option<Arc<MCaptcha>> {
        if let Some(captcha) = self.captchas.get(id) {
            Some(captcha.clone())
        } else {
            None
        }
    }

    /// removes [Counter] actor from [Manager]
    pub fn rm_captcha(&self, id: &str) -> Option<(String, Arc<MCaptcha>)> {
        self.captchas.remove(id)
    }

    /// renames [Counter] actor
    pub fn rename(&self, current_id: &str, new_id: String) {
        // If actor isn't present, it's okay to not throw an error
        // since actors are lazyily initialized and are cleaned up when inactive
        if let Some((_, captcha)) = self.captchas.remove(current_id) {
            self.add_captcha(captcha, new_id);
        }
    }

    pub async fn clean_all_after_cold_start(&self, updated: Manager) {
        updated.captchas.iter().for_each(|x| {
            self.captchas
                .insert(x.key().to_owned(), x.value().to_owned());
        });
        let captchas = self.clone();
        let keys: Vec<String> = captchas
            .captchas
            .clone()
            .iter()
            .map(|x| x.key().to_owned())
            .collect();
        let fut = async move {
            tokio::time::sleep(Duration::new(captchas.gc, 0)).await;
            for key in keys.iter() {
                captchas.rm_captcha(key);
            }
        };

        tokio::spawn(fut);
    }

    pub fn add_visitor(
        &self,
        msg: &ManagerMessages::AddVisitor,
    ) -> Option<libmcaptcha::master::AddVisitorResult> {
        if let Some(captcha) = self.captchas.get(&msg.0) {
            let difficulty_factor = captcha.add_visitor();
            //           let id = msg.0.clone();

            let c = captcha.clone();
            let captchas = self.captchas.clone();
            let fut = async move {
                tokio::time::sleep(Duration::new(c.duration, 0)).await;
                c.decrement_visitor_by(1);
                //                Self::gc(captchas);
                //                if c.get_visitors() == 0 {
                //                    println!("Removing captcha addvivi");
                //                    captchas.remove(&id);
                //                }
            };

            tokio::spawn(fut);

            Some(libmcaptcha::master::AddVisitorResult {
                duration: captcha.duration,
                difficulty_factor,
            })
        } else {
            None
        }
    }

    pub fn get_internal_data(&self) -> HashMap<String, libmcaptcha::mcaptcha::MCaptcha> {
        let mut res = HashMap::with_capacity(self.captchas.len());
        for value in self.captchas.iter() {
            res.insert(value.key().to_owned(), value.value().as_ref().into());
        }
        res
    }

    pub fn set_internal_data(&self, mut map: HashMap<String, libmcaptcha::mcaptcha::MCaptcha>) {
        for (id, captcha) in map.drain() {
            let visitors = captcha.get_visitors();
            let new_captcha: MCaptcha = (&captcha).into();
            let new_captcha = Arc::new(new_captcha);
            self.captchas.insert(id.clone(), new_captcha.clone());
            let msg = ManagerMessages::AddVisitor(id);
            for _ in 0..visitors {
                self.add_visitor(&msg);
            }
        }
    }
}

impl From<&libmcaptcha::mcaptcha::MCaptcha> for MCaptcha {
    fn from(value: &libmcaptcha::mcaptcha::MCaptcha) -> Self {
        let mut defense = super::defense::DefenseBuilder::default();
        for level in value.get_defense().get_levels() {
            let _ = defense.add_level(level);
        }
        let defense = defense.build().unwrap();
        let new_captcha = MCaptchaBuilder::default()
            .defense(defense)
            .duration(value.get_duration())
            .build()
            .unwrap();
        //        for _ in 0..value.get_visitors() {
        //            new_captcha.add_visitor();
        //        }

        new_captcha
    }
}

impl From<&MCaptcha> for libmcaptcha::mcaptcha::MCaptcha {
    fn from(value: &MCaptcha) -> Self {
        let mut defense = libmcaptcha::defense::DefenseBuilder::default();
        for level in value.defense.get_levels().drain(0..) {
            let _ = defense.add_level(level);
        }
        let defense = defense.build().unwrap();
        let mut new_captcha = libmcaptcha::mcaptcha::MCaptchaBuilder::default()
            .defense(defense)
            .duration(value.duration)
            .build()
            .unwrap();
        for _ in 0..value.get_visitors() {
            new_captcha.add_visitor();
        }
        new_captcha
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use libmcaptcha::defense::LevelBuilder;
    use libmcaptcha::master::messages::*;

    pub const LEVEL_1: (u32, u32) = (50, 50);
    pub const LEVEL_2: (u32, u32) = (500, 500);
    pub const DURATION: u64 = 5;

    use crate::mcaptcha::defense::*;

    pub fn get_defense() -> Defense {
        DefenseBuilder::default()
            .add_level(
                LevelBuilder::default()
                    .visitor_threshold(LEVEL_1.0)
                    .difficulty_factor(LEVEL_1.1)
                    .unwrap()
                    .build()
                    .unwrap(),
            )
            .unwrap()
            .add_level(
                LevelBuilder::default()
                    .visitor_threshold(LEVEL_2.0)
                    .difficulty_factor(LEVEL_2.1)
                    .unwrap()
                    .build()
                    .unwrap(),
            )
            .unwrap()
            .build()
            .unwrap()
    }

    async fn race(manager: &Manager, id: String, count: (u32, u32)) {
        let msg = ManagerMessages::AddVisitor(id);
        for _ in 0..count.0 as usize - 1 {
            manager.add_visitor(&msg);
        }
    }

    //    pub fn get_counter() -> Counter {
    //        get_mcaptcha().into()
    //    }

    pub fn get_mcaptcha() -> MCaptcha {
        MCaptchaBuilder::default()
            .defense(get_defense())
            .duration(DURATION)
            .build()
            .unwrap()
    }

    #[actix_rt::test]
    async fn manager_works() {
        let manager = Manager::new(1);

        //        let get_add_site_msg = |id: String, mcaptcha: MCaptcha| {
        //            AddSiteBuilder::default()
        //                .id(id)
        //                .mcaptcha(mcaptcha)
        //                .build()
        //                .unwrap()
        //        };

        let id = "yo";
        manager.add_captcha(Arc::new(get_mcaptcha()), id.into());

        let mcaptcha_addr = manager.get_captcha(id);
        assert!(mcaptcha_addr.is_some());

        let mut mcaptcha_data = manager.get_internal_data();
        mcaptcha_data.get_mut(id).unwrap().add_visitor();
        mcaptcha_data.get_mut(id).unwrap().add_visitor();
        mcaptcha_data.get_mut(id).unwrap().add_visitor();
        //        let mcaptcha_data: HashMap<String, libmcaptcha::mcaptcha::MCaptcha> = {
        //        let serialized  = serde_json::to_string(&mcaptcha_data).unwrap();
        //        serde_json::from_str(&serialized).unwrap()
        //        };
        //        println!("{:?}", mcaptcha_data);
        manager.set_internal_data(mcaptcha_data);

        let mcaptcha_data = manager.get_internal_data();
        assert_eq!(
            manager.get_captcha(id).unwrap().get_visitors(),
            mcaptcha_data.get(id).unwrap().get_visitors()
        );

        let new_id = "yoyo";
        manager.rename(id, new_id.into());

        {
            let mcaptcha_addr = manager.get_captcha(new_id);
            assert!(mcaptcha_addr.is_some());

            let addr_doesnt_exist = manager.get_captcha(id);
            assert!(addr_doesnt_exist.is_none());

            let timer_expire = Duration::new(DURATION, 0);
            tokio::time::sleep(timer_expire).await;
            tokio::time::sleep(timer_expire).await;
        }

        //        Manager::gc(manager.captchas.clone());
        //        let mcaptcha_addr = manager.get_captcha(new_id);
        //        assert_eq!(mcaptcha_addr.as_ref().unwrap().get_visitors(), 0);
        //        assert!(mcaptcha_addr.is_none());
        //
        //        assert!(
        //            manager.rm_captcha(new_id.into()).is_some());
    }

    #[actix_rt::test]
    async fn counter_defense_works() {
        let manager = Manager::new(1);
        let id = "yo";
        manager.add_captcha(Arc::new(get_mcaptcha()), id.into());

        let mut mcaptcha = manager
            .add_visitor(&ManagerMessages::AddVisitor(id.to_string()))
            .unwrap();
        assert_eq!(mcaptcha.difficulty_factor, LEVEL_1.0);

        race(&manager, id.to_string(), LEVEL_2).await;
        mcaptcha = manager
            .add_visitor(&ManagerMessages::AddVisitor(id.to_string()))
            .unwrap();
        assert_eq!(mcaptcha.difficulty_factor, LEVEL_2.1);
        tokio::time::sleep(Duration::new(DURATION * 2, 0)).await;
        assert_eq!(manager.get_captcha(id).unwrap().get_visitors(), 0);
    }
}
//
//#[cfg(test)]
//pub mod tests {
//    use super::*;
//    use crate::defense::*;
//    use crate::errors::*;
//    use crate::mcaptcha;
//    use crate::mcaptcha::MCaptchaBuilder;
//
//    // constants for testing
//    // (visitor count, level)
//    pub const LEVEL_1: (u32, u32) = (50, 50);
//    pub const LEVEL_2: (u32, u32) = (500, 500);
//    pub const DURATION: u64 = 5;
//
//    type MyActor = Addr<Counter>;
//
//    pub fn get_defense() -> Defense {
//        DefenseBuilder::default()
//            .add_level(
//                LevelBuilder::default()
//                    .visitor_threshold(LEVEL_1.0)
//                    .difficulty_factor(LEVEL_1.1)
//                    .unwrap()
//                    .build()
//                    .unwrap(),
//            )
//            .unwrap()
//            .add_level(
//                LevelBuilder::default()
//                    .visitor_threshold(LEVEL_2.0)
//                    .difficulty_factor(LEVEL_2.1)
//                    .unwrap()
//                    .build()
//                    .unwrap(),
//            )
//            .unwrap()
//            .build()
//            .unwrap()
//    }
//
//    async fn race(addr: Addr<Counter>, count: (u32, u32)) {
//        for _ in 0..count.0 as usize - 1 {
//            let _ = addr.send(AddVisitor).await.unwrap();
//        }
//    }
//
//    pub fn get_counter() -> Counter {
//        get_mcaptcha().into()
//    }
//
//    pub fn get_mcaptcha() -> MCaptcha {
//        MCaptchaBuilder::default()
//            .defense(get_defense())
//            .duration(DURATION)
//            .build()
//            .unwrap()
//    }
//
//    #[test]
//    fn mcaptcha_decrement_by_works() {
//        let mut m = get_mcaptcha();
//        for _ in 0..100 {
//            m.add_visitor();
//        }
//        m.decrement_visitor_by(50);
//        assert_eq!(m.get_visitors(), 50);
//        m.decrement_visitor_by(500);
//        assert_eq!(m.get_visitors(), 0);
//    }
//

//
//    #[actix_rt::test]
//    async fn counter_defense_loosenup_works() {
//        //use actix::clock::sleep;
//        //use actix::clock::delay_for;
//        let addr: MyActor = get_counter().start();
//
//        race(addr.clone(), LEVEL_2).await;
//        race(addr.clone(), LEVEL_2).await;
//        let mut mcaptcha = addr.send(AddVisitor).await.unwrap();
//        assert_eq!(mcaptcha.difficulty_factor, LEVEL_2.1);
//
//        let duration = Duration::new(DURATION, 0);
//        sleep(duration).await;
//        //delay_for(duration).await;
//
//        mcaptcha = addr.send(AddVisitor).await.unwrap();
//        assert_eq!(mcaptcha.difficulty_factor, LEVEL_1.1);
//    }
//
//    #[test]
//    fn test_mcatcptha_builder() {
//        let defense = get_defense();
//        let m = MCaptchaBuilder::default()
//            .duration(0)
//            .defense(defense.clone())
//            .build();
//
//        assert_eq!(m.err(), Some(CaptchaError::CaptchaDurationZero));
//
//        let m = MCaptchaBuilder::default().duration(30).build();
//        assert_eq!(
//            m.err(),
//            Some(CaptchaError::PleaseSetValue("defense".into()))
//        );
//
//        let m = MCaptchaBuilder::default().defense(defense).build();
//        assert_eq!(
//            m.err(),
//            Some(CaptchaError::PleaseSetValue("duration".into()))
//        );
//    }
//
//    #[actix_rt::test]
//    async fn get_current_visitor_count_works() {
//        let addr: MyActor = get_counter().start();
//
//        addr.send(AddVisitor).await.unwrap();
//        addr.send(AddVisitor).await.unwrap();
//        addr.send(AddVisitor).await.unwrap();
//        addr.send(AddVisitor).await.unwrap();
//        let count = addr.send(GetCurrentVisitorCount).await.unwrap();
//
//        assert_eq!(count, 4);
//    }
//
//    #[actix_rt::test]
//    #[should_panic]
//    async fn stop_works() {
//        let addr: MyActor = get_counter().start();
//        addr.send(Stop).await.unwrap();
//        addr.send(AddVisitor).await.unwrap();
//    }
//
//    #[actix_rt::test]
//    async fn get_set_internal_data_works() {
//        let addr: MyActor = get_counter().start();
//        let mut mcaptcha = addr.send(GetInternalData).await.unwrap();
//        mcaptcha.add_visitor();
//        addr.send(SetInternalData(mcaptcha.clone())).await.unwrap();
//        assert_eq!(
//            addr.send(GetInternalData).await.unwrap().get_visitors(),
//            mcaptcha.get_visitors()
//        );
//
//        let duration = Duration::new(mcaptcha.get_duration() + 3, 0);
//        sleep(duration).await;
//        assert_eq!(addr.send(GetCurrentVisitorCount).await.unwrap(), 0);
//    }
//
//    #[actix_rt::test]
//    async fn bulk_delete_works() {
//        let addr: MyActor = get_counter().start();
//        addr.send(AddVisitor).await.unwrap();
//        addr.send(AddVisitor).await.unwrap();
//        assert_eq!(addr.send(GetCurrentVisitorCount).await.unwrap(), 2);
//        addr.send(BulkDecrement(3)).await.unwrap();
//        assert_eq!(addr.send(GetCurrentVisitorCount).await.unwrap(), 0);
//    }
//}