use crate::session::session::Session;
use crate::session::session::SessionDirection;
use crate::session::session::SessionState;
use crate::utils::utils::reverse_flow_id;
use chrono::Utc;
use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;

/******************************************************************************
 * Struct
 ******************************************************************************/

pub struct SessionManager {
    sessions: HashMap<String, Rc<RefCell<Session>>>,
}

/******************************************************************************
 * API
 ******************************************************************************/

impl SessionManager {
    pub fn new(capacity: usize) -> SessionManager {
        SessionManager {
            sessions: HashMap::with_capacity(capacity),
        }
    }

    pub fn update(&mut self, session_id: String) -> Rc<RefCell<Session>> {
        let result = self.get_session(&session_id);
        if result.is_none() {
            // New Session
            let mut new_session = Session::new(session_id.clone());
            new_session.update_session_current_dir(SessionDirection::C2S);

            let rc_new_session = Rc::new(RefCell::new(new_session));
            self.insert_session(session_id.clone(), rc_new_session.clone());
            rc_new_session
        } else {
            // Update Session
            let rc_old_session = result.unwrap();
            rc_old_session
                .borrow_mut()
                .set_session_state(SessionState::Active);
            rc_old_session.borrow_mut().update_session_last_ts();
            rc_old_session.borrow_mut().update_session_expire_ts();

            // Update Session Direction
            if rc_old_session.borrow().get_session_id() == session_id {
                rc_old_session
                    .borrow_mut()
                    .update_session_current_dir(SessionDirection::C2S);
            } else {
                rc_old_session
                    .borrow_mut()
                    .update_session_current_dir(SessionDirection::S2C);
            }

            rc_old_session
        }
    }

    pub fn get_session(&self, session_id: &String) -> Option<Rc<RefCell<Session>>> {
        let result = self.sessions.get(session_id).map(|session| session.clone());
        if result.is_some() {
            return result;
        }

        let reversed_id = reverse_flow_id(session_id);
        self.sessions
            .get(reversed_id.as_str())
            .map(|session| session.clone())
    }

    pub fn insert_session(&mut self, session_id: String, session: Rc<RefCell<Session>>) {
        self.sessions.insert(session_id, session);
    }

    pub fn remove_session(&mut self, session_id: &str) -> Option<Rc<RefCell<Session>>> {
        let result = self.sessions.remove(session_id);
        if result.is_some() {
            return result;
        }

        let reversed_id = reverse_flow_id(&session_id.to_string());
        self.sessions.remove(reversed_id.as_str())
    }

    pub fn expire_oldest_session(&mut self) -> Option<Rc<RefCell<Session>>> {
        let mut oldest_session: Option<Rc<RefCell<Session>>> = None;
        let mut oldest_session_expire_ts = Utc::now().timestamp();

        for (_, session) in &self.sessions {
            let session_expire_ts = session.borrow().get_session_expire_ts();
            if session_expire_ts < oldest_session_expire_ts {
                oldest_session = Some(session.clone());
                oldest_session_expire_ts = session_expire_ts;
            }
        }

        if oldest_session.is_some() {
            let oldest_session_id = oldest_session.clone().unwrap().borrow().get_session_id();
            oldest_session
                .clone()
                .unwrap()
                .borrow_mut()
                .set_session_state(SessionState::Expired);
            oldest_session
                .clone()
                .unwrap()
                .borrow_mut()
                .set_session_end_ts(oldest_session_expire_ts);
            self.remove_session(&oldest_session_id);
        }

        oldest_session
    }
}

/******************************************************************************
 * TEST
 ******************************************************************************/

#[cfg(test)]
mod tests {
    use super::SessionManager;
    use crate::session::session::Session;
    use std::cell::RefCell;
    use std::rc::Rc;

    #[test]
    fn test_session_manager() {
        let session_id = "192.168.0.1->192.168.0.2;UDP->UDP;2345->80;".to_string();
        let reversed_id = "192.168.0.2->192.168.0.1;UDP->UDP;80->2345;".to_string();
        let session = Rc::new(RefCell::new(Session::new(session_id.clone())));

        // Create Session Manager
        let mut session_mgr = SessionManager::new(1024);

        // Search Session
        assert_eq!(session_mgr.get_session(&session_id).is_none(), true);

        // Insert Session
        session_mgr.insert_session(session_id.clone(), session);

        // Update Session
        assert_eq!(session_mgr.get_session(&session_id).is_some(), true);
        assert_eq!(session_mgr.get_session(&reversed_id).is_some(), true);

        // Delete session
        assert_eq!(session_mgr.remove_session(&reversed_id).is_some(), true);

        // Research Session
        assert_eq!(session_mgr.get_session(&session_id).is_none(), true);
        assert_eq!(session_mgr.get_session(&reversed_id).is_none(), true);
    }
}