dce-router 1.8.0

use std::{any::type_name, collections::{HashMap, HashSet}, sync::RwLock};

use dce_util::{arena_tree::{ArenaTree, KeyFactory}, result::{DceError, DceResult}};
use log::{debug, warn};

use crate::{api::{Api, Handler, Hook, Suffix, MARK_PATH_PART_SEPARATOR, MARK_SUFFIX_BOUNDARY, MARK_VARIABLE_CLOSING, MARK_VARIABLE_OPENER, MARK_VAR_TYPE_EMPTABLE_VECTOR, MARK_VAR_TYPE_OPTIONAL, MARK_VAR_TYPE_VECTOR}, context::Param, protocol::RoutableProtocol};

const CODE_NOT_FOUND: isize = 404;
const HOOK_PATH_SUFFIXES: [&str; 2] = ["+", "*"];

pub struct Router<Rp: RoutableProtocol + 'static> {
    path_separator: &'static str,
    suffix_boundary: &'static str,
    api_buffer: RwLock<Vec<&'static Api<Rp>>>,
    raw_omitted_paths: Vec<&'static str>,
    id_api_mapping: HashMap<&'static str, &'static Api<Rp>>,
    apis_mapping: HashMap<&'static str, Vec<&'static Api<Rp>>>,
    apis_tree: ArenaTree<ApiBranch<Rp>, &'static str>,
    before_mapping: HashMap<&'static str, Hook<Rp>>,
    after_mapping: HashMap<&'static str, Hook<Rp>>,
    path_before_mapping: HashMap<&'static str, &'static str>,
    path_after_mapping: HashMap<&'static str, &'static str>,
}

impl <Rp: RoutableProtocol> Router<Rp> {
    pub fn new() -> Self {
        Router {
            path_separator: MARK_PATH_PART_SEPARATOR,
            suffix_boundary: MARK_SUFFIX_BOUNDARY,
            api_buffer: Default::default(),
            raw_omitted_paths: Default::default(),
            id_api_mapping: Default::default(),
            apis_mapping: Default::default(),
            apis_tree: ArenaTree::new(ApiBranch::new("")),
            before_mapping: Default::default(),
            after_mapping: Default::default(),
            path_before_mapping: Default::default(),
            path_after_mapping: Default::default(),
        }
    }

    pub fn bind(&mut self, path: &'static str, handler: Handler<Rp>) -> &mut Self {
        let api = Api::new(path);
        let api = api.bind_handler(handler);
        self.add(api)
    }

    pub fn bind_api(&mut self, api: Api<Rp>, handler: Handler<Rp>) -> &mut Self {
        let api = api.bind_handler(handler);
        self.add(api)
    }

    pub fn register(&mut self, api_supplier: fn() -> Api<Rp>) -> &mut Self {
        let api = api_supplier();
        self.add(api)
    }

    pub fn add(&mut self, api: Api<Rp>) -> &mut Self {
        let api = Box::leak(Box::new(api.upgrade()));
        let mut api_buffer = self.api_buffer.write().unwrap();
        api_buffer.push(api);
        if api.omission {
            self.raw_omitted_paths.push(api.path);
        }
        if let Some(id) = &api.id {
            self.id_api_mapping.insert(id, api);
        }
        drop(api_buffer);
        self
    }

    pub fn ready(&mut self) {
        self.build_tree();
        let mut api_buffer = self.api_buffer.write().unwrap();
        while api_buffer.len() > 0 {
            let api = api_buffer.remove(0);
            let path = self.omitted_path(api.path);
            let mut apis = vec![api];
            let mut suffixes = api.suffixes.iter().map(|s| s).collect::<HashSet<_>>();
            let mut i = 0;
            while i < api_buffer.len() {
			    // collect the omitted same path into an array
                if path == self.omitted_path(api_buffer.get(i).unwrap().path) {
                    let omitted_api = api_buffer.remove(i);
                    apis.push(omitted_api);
                    suffixes.extend(omitted_api.suffixes.iter().collect::<Vec<_>>());                    
                } else {
                    i += 1;
                }
            }
            // append suffix to path as api mapping key to grouping the apis
            for suffix in suffixes {
                // insert suffix matched apis into the mapping
                self.apis_mapping.insert(Box::leak(Self::suffixed_path(path.as_str(), Some(suffix)).into_boxed_str()),
                    apis.iter().filter(|a| a.suffixes.contains(suffix)).map(|a| *a).collect::<Vec<_>>());
            }
        }
        // self.api_buffer.read().unwrap().iter().for_each(|a| println!("{}", a.path()));
        self.path_before_mapping = self.map_middleware(&self.before_mapping, true);
        self.path_after_mapping = self.map_middleware(&self.after_mapping, false);
    }

    fn build_tree(&mut self) {
    	// 1. make apis to ApiBranches
        let api_buffer = self.api_buffer.write().unwrap();
        let mut path_apis = api_buffer.iter().map(|a| a.path).collect::<HashSet<_>>()
            .into_iter().map(|p| (p, ApiBranch::new(p)))
            .collect::<HashMap<_, _>>();
        for api in api_buffer.iter() {
            path_apis.get_mut(api.path).map(|v| v.push(*api));
        }
	    // 2. init the apis_tree
        self.apis_tree.fill(path_apis.into_values().collect::<Vec<_>>());
	    // 3. fill the ApiBranche properties
        for i in 1..self.apis_tree.len() {
            let mut is_omitted_child = false;
            let node = (*self.apis_tree).get(i).unwrap();
            let node_var_type = node.element().var_type.clone();
            let mut parent_index = *node.parent_index();
            while let Some(Some(parent)) = parent_index.map(|pi| (*self.apis_tree).get_mut(pi)) {
                parent_index = *parent.parent_index();
                if parent.element().is_omission {
                    is_omitted_child = true;
                    continue;
                }
                match parent.element().var_type {
                    VarType::Required(_) => parent.element_mut().is_mid_var = true,
                    VarType::NotVar => break,
                    _ => panic!(r#"Ambiguous type var "{}" cannot in middle."#, parent.element().id()),
                };
                if !matches!(node_var_type, VarType::NotVar) {
                    parent.element_mut().var_children.push(i);
                } else if is_omitted_child {
                    parent.element_mut().omitted_children.push(i);
                }
            }
        }
    }

    fn omitted_path(&self, path: &str) -> String {
        // Path in api field should always be `MARK_PATH_PART_SEPARATOR`
        let parts = path.split(MARK_PATH_PART_SEPARATOR).collect::<Vec<_>>();
        parts.iter().enumerate()
            .filter(|&(i, _)| !self.raw_omitted_paths.contains(&parts[..=i].join(MARK_PATH_PART_SEPARATOR).as_str()))
            .map(|(_, p)| *p).collect::<Vec<_>>()
            .join(MARK_PATH_PART_SEPARATOR)
    }

    fn suffixed_path(path: &str, suffix: Option<&Suffix>) -> String {
        suffix.filter(|s| !s.as_ref().is_empty())
            .map_or_else(|| path.to_owned(), |s| format!("{}{}{}", path, MARK_SUFFIX_BOUNDARY, s.as_ref()))
    }

    fn map_middleware(&self, handler_mapping: &HashMap<&'static str, Hook<Rp>>, pre: bool) -> HashMap<&'static str, &'static str> {
        let mut apis_paths = self.apis_mapping.keys().collect::<Vec<_>>();
        let mut path_mapping: HashMap<&'static str, &'static str> = Default::default();

        for (key, _) in handler_mapping {
            let mut path = key.to_string();
            let mut suffix = String::new();
            let wildcard = HOOK_PATH_SUFFIXES.iter().find(|s| key.ends_with(*s));

            if let Some(w) = wildcard {
                suffix = key[key.len() - w.len()..].to_string();
                path = key[..key.len() - w.len()].to_string();
            }
            apis_paths = apis_paths
                .into_iter()
                .filter(|api_path| {
                    if {
                        if suffix != HOOK_PATH_SUFFIXES[0] && path.eq(*api_path) {
                            true
                        } else if matches!(wildcard, Some(_)) {
                            path.is_empty() || api_path.starts_with(&(path.clone() + MARK_PATH_PART_SEPARATOR))
                        } else {
                            false
                        }
                    } {
                        Self::hook_override_warn(api_path, &path_mapping, pre);
                        path_mapping.insert(api_path, key);
                        false
                    } else {
                        true
                    }
                })
                .collect();
        }
        path_mapping
    }

    fn hook_override_warn(path: &str, mapping: &HashMap<&'static str, &'static str>, pre: bool) {
        if mapping.contains_key(path) {            
            let hook = if pre { "pre-handler" } else { "post-handler" };
            warn!(r#"Path "{}" already has a {}; reassigning it will overwrite the current one."#, path, hook);
        }
    }

    pub fn lookup(&self, rp: &Rp) -> DceResult<RouteMatch<'_, Rp>> {
        self.lookup_api(rp)
            .map(|(api, params, suffix)| {
                let pre_hook = self.path_before_mapping.get(api.path)
                    .map(|p| self.before_mapping.get(p)).flatten();
                let post_hook = self.path_after_mapping.get(api.path)
                    .map(|p| self.after_mapping.get(p)).flatten();
                RouteMatch{api, params, suffix, pre_hook, post_hook}
            })
    }

    fn lookup_api(&self, rp: &Rp) -> DceResult<(&'static Api<Rp>, HashMap<String, Param>, Option<Suffix>)> {
        let req_path = rp.path();
        let mut api = None;
        let mut suffix = None;
        let mut params = HashMap::new();
        let mut path = req_path;
        while api.is_none() {
            let mut apis = self.apis_mapping.get(path);
            if apis.is_none() {
                if let Some((tmp_path, tmp_params, tmp_suffix)) = self.lookup_var(path) {
                    apis = self.apis_mapping.get(Self::suffixed_path(tmp_path, tmp_suffix.as_ref()).as_str());
                    params = tmp_params;
                    suffix = tmp_suffix;
                }
            }
            if let Some(apis) = apis {
                if let Some(matched) = rp.match_api(apis) {
                    if let Some(redirect) = matched.redirect {
                        path = redirect;
                        continue;
                    }
                    api = Some(matched);
                    break;
                }
            }
            if self.apis_mapping.is_empty() {
                if self.api_buffer.read().unwrap().is_empty() {
				    panic!(r#"Lookup failed, "Router.apiBuffer" is empty, you may need to call the "Router.Push()" to bind apis"#);
                } else {
				    panic!(r#"Router is not ready, please call "Router.ready()" first before lookup"#);
                }
            } else {
                break;
            }
        }
        api.map(|a| {
            debug!(r#"{}: path "{}" matched api "{}""#, type_name::<Rp>(), req_path, a.path);
            (a, params, suffix)
        }).ok_or_else(|| DceError::pub_msg(CODE_NOT_FOUND, format!(r#"path "{}" route failed, could not matched by Router"#, path)))
    }

    fn lookup_var(&self, path: &str) -> Option<(&'static str, HashMap<String, Param>, Option<Suffix>)> {
        let mut path_parts = path.split(self.path_separator).collect::<Vec<_>>();
        let mut branch_and_part_indexes = vec![(0, 0)];
        let mut params = HashMap::new();
        let mut target_api_branch = None;
        let mut suffix = None;
        'outer: while let Some((branch_index, part_index)) = branch_and_part_indexes.pop() {
            let api_branch = (*self.apis_tree).get(branch_index).unwrap();
            let is_last_part = part_index == path_parts.len() - 1;
            let is_overflowed = part_index >= path_parts.len();
            if is_overflowed && ! api_branch.element().apis.is_empty() {
                // should be finished at last request path part if not a bare tree
                target_api_branch = Some(api_branch);
			    break;
            }
            if ! is_overflowed {
                // if not overflow and request path matched, then it must be a normal path
                if let Some((sub_api_index, matched_suffix)) = self.find_consider_suffix(path_parts[part_index], is_last_part, api_branch.child_indexes(), &api_branch.element().omitted_children) {
                    branch_and_part_indexes.push((sub_api_index, part_index + 1));
                    suffix = matched_suffix;
                    continue;
                }
            }

            let insert_pos = branch_and_part_indexes.len();
            for (var_branch_index, var_api_branch) in api_branch.element().var_children.iter().filter_map(
                |i| self.apis_tree.by_index(*i).map(|n| (i, n))).collect::<Vec<_>>() {
                if ! var_api_branch.element().is_mid_var {
                    // if not a middle var, then should finish var path match and collect vars and end the outer loop
                    match &var_api_branch.element().var_type {
                        VarType::Optional(_) if is_overflowed => {},
                        VarType::Optional(var_name) | VarType::Required(var_name) if is_last_part =>
                            suffix = self.suffix_trimmer(&mut path_parts, var_api_branch.element(),
                                &mut |ps|{ params.insert(var_name.clone(), Param::Scalar(ps.get(0).map(|p| p.to_string()).unwrap())); }),
                        VarType::EmptableVector(_) if is_overflowed => {},
                        VarType::EmptableVector(var_name) | VarType::Vector(var_name) if ! is_overflowed =>
                            suffix = self.suffix_trimmer(&mut path_parts, var_api_branch.element(),
                                &mut |ps|{ params.insert(var_name.clone(), Param::Vector(ps.iter().map(|p| p.to_string()).collect::<Vec<_>>())); }),
                        _ => continue,
                    };
                    target_api_branch = Some(var_api_branch);
                    break 'outer
                } else if let VarType::Required(var_name) = &var_api_branch.element().var_type {
				    // if it's middle var then insert to loop queue to handle it next cycle
                    params.insert(var_name.clone(), Param::Scalar(path_parts[part_index].to_owned()));
                    branch_and_part_indexes.insert(insert_pos, (*var_branch_index, part_index + 1));
                }
            }
        }
        target_api_branch.map(|b| (b.element().path, params, suffix))
    }

    fn suffix_trimmer(&self, parts: &mut Vec<&str>, branch: &ApiBranch<Rp>, consumer: &mut dyn FnMut(&Vec<&str>)) -> Option<Suffix> {
        // just need to check is_last_part because should already handle suffix if overflowed
        // pop out the last part to clean (cut off the suffix)
        let mut suffix = None;
        if let Some(mut last_part) = parts.pop() {
            if let Some(ts) = branch.apis.iter().flat_map(|a| &a.suffixes).find(|s| last_part.ends_with(format!("{}{}", self.suffix_boundary, s.as_ref()).as_str())) {
                last_part = &last_part[..last_part.len() - self.suffix_boundary.len() - ts.as_ref().len()];
                suffix = Some(ts.clone());
            }
            // Push last part (basename without suffix) back into path parts
            parts.push(last_part);
        }
        consumer(parts);
        suffix
    }

    fn find_consider_suffix(&self, path_part: &str, is_last_part: bool, child_indexes: &Vec<usize>, omitted_indexes: &Vec<usize>) -> Option<(usize, Option<Suffix>)> {
        let mut matches = self.find_from_indexes_by_part(child_indexes, path_part);
        if matches.is_none() {
            matches = self.find_from_indexes_by_part(omitted_indexes, path_part);
        }
        let mut suffix = None;
        if matches.is_none() && is_last_part {
            let mut boundary = Some(path_part.len());
            loop {
                if let Some((base_part, index)) = boundary.map(|i| (&path_part[..i], i)) {
                    matches = self.find_from_indexes_by_part(child_indexes, base_part);
                    if matches.is_none() {
                        matches = self.find_from_indexes_by_part(omitted_indexes, base_part);
                    }
                    // 如果从剩余的路径找到了后缀边界符，且为找到匹配的API分支，则尝试以切割剩余的尾部与 Apis 里的 suffixes 匹配，两者都匹配才是真正找到
                    suffix = matches.map(|(_, ab)| ab.apis.iter().flat_map(|a| &a.suffixes)

                        .find(|s| path_part[index + 1..].eq(s.as_ref())))
                        .flatten().map(|s| s.clone());
                    if suffix.is_none() {
                        // 如果未找到API分支或后缀不匹配，则跳过此轮循环以尝试前移边界匹配API分支
                        boundary = path_part[..index].rfind(self.suffix_boundary);
                        continue;
                    }
                }
                break;
            }
        }
        matches.map(|(i, _)| (i, suffix))
    }

    fn find_from_indexes_by_part(&self, indexes: &Vec<usize>, part: &str) -> Option<(usize, &ApiBranch<Rp>)> {
        indexes.iter()
            .filter_map(|i| self.apis_tree.by_index(*i).map(|n| (*i, n.element())))
            .find(|(_, ab)| ab.path_part().eq(part))
    }
}

pub struct RouteMatch<'a, Rp: RoutableProtocol + 'static> {
    pub api: &'static Api<Rp>,
    pub params: HashMap<String, Param>,
    pub suffix: Option<Suffix>,
    pub pre_hook: Option<&'a Hook<Rp>>,
    pub post_hook: Option<&'a Hook<Rp>>,
}

#[derive(Clone)]
pub enum VarType {
    NotVar,
    Required(String),
    Optional(String),
    Vector(String),
    EmptableVector(String),
}

struct ApiBranch<Rp: RoutableProtocol + 'static> {
    path: &'static str,
    var_type: VarType,
    is_mid_var: bool,
    is_omission: bool,
    apis: Vec<&'static Api<Rp>>,
    var_children: Vec<usize>,
    omitted_children: Vec<usize>,
}

impl <Rp: RoutableProtocol + 'static> ApiBranch<Rp> {
    fn push(&mut self, api: &'static Api<Rp>) {
        self.apis.push(api);
        if api.omission {
            self.is_omission = true;
        }
    }

    fn path_part(&self) -> &str {
        self.path.rfind(MARK_PATH_PART_SEPARATOR).map_or(self.path, |i| &self.path[i+1..])
    }

    fn new(path: &'static str) -> Self {
        let mut var_type = VarType::NotVar;
        if path.starts_with(MARK_VARIABLE_OPENER) && path.ends_with(MARK_VARIABLE_CLOSING) {
            let mut var_name = path[MARK_VARIABLE_OPENER.len() .. path.len() - MARK_VARIABLE_CLOSING.len()].to_string();
            if var_name.ends_with(MARK_VAR_TYPE_OPTIONAL) {
                var_name = var_name[..var_name.len() - MARK_VAR_TYPE_OPTIONAL.len()].to_string();
                var_type = VarType::Optional(var_name);
            } else if var_name.ends_with(MARK_VAR_TYPE_EMPTABLE_VECTOR) {
                var_name = var_name[..var_name.len() - MARK_VAR_TYPE_EMPTABLE_VECTOR.len()].to_string();
                var_type = VarType::EmptableVector(var_name);
            } else if var_name.ends_with(MARK_VAR_TYPE_VECTOR) {
                var_name = var_name[..var_name.len() - MARK_VAR_TYPE_VECTOR.len()].to_string();
                var_type = VarType::Vector(var_name);
            } else {
                var_type = VarType::Required(var_name);
            }
        }
        Self {
            path,
            var_type,
            is_mid_var: Default::default(),
            is_omission: Default::default(),
            apis: Default::default(),
            var_children: Default::default(),
            omitted_children: Default::default(),
        }
    }
}

impl <Rp: RoutableProtocol + 'static> KeyFactory<&'static str> for ApiBranch<Rp> {
    fn id(&self) -> &'static str {
        self.path
    }

    fn child_of(&self, parent: &Self) -> bool {
        self.path.rfind(MARK_PATH_PART_SEPARATOR).map_or_else(|| parent.path.is_empty(), |i| self.path[..i].eq(parent.path))
    }
    
    fn new_parent(&self) -> Self {
        Self::new(self.path.rfind(MARK_PATH_PART_SEPARATOR).map_or("", |i| &self.path[..i]))
    }
}