use log::info;
use proxy_wasm::traits::*;
use proxy_wasm::types::*;
use serde::Deserialize;

proxy_wasm::main! {{
    proxy_wasm::set_log_level(LogLevel::Trace);
    proxy_wasm::set_root_context(|_| -> Box<dyn RootContext> {
        Box::new(CustomPolicyHeaderRoot {
            config: CustomPolicyConfig::default()
        })
    });
}}

// ---- CustomPolicyConfig ----

#[derive(Default, Clone, Deserialize)]
struct CustomPolicyConfig {
    #[serde(alias = "username")]
    username: String,
}


// ---- CustomPolicyHeaderRoot ----

struct CustomPolicyHeaderRoot {
    config: CustomPolicyConfig,
}

impl Context for CustomPolicyHeaderRoot {}

impl RootContext for CustomPolicyHeaderRoot {
    fn on_configure(&mut self, _: usize) -> bool {
        if let Some(config_bytes) = self.get_plugin_configuration() {
            self.config = serde_json::from_slice(config_bytes.as_slice()).unwrap()
        }
        true
    }

    fn create_http_context(&self, _: u32) -> Option<Box<dyn HttpContext>> {
        Some(Box::new(CustomPolicyHeader {
            config: self.config.clone()
        }))
    }

    fn get_type(&self) -> Option<ContextType> {
        Some(ContextType::HttpContext)
    }
}


// ---- CustomPolicyHeader ----

struct CustomPolicyHeader {
     config: CustomPolicyConfig,
}

impl Context for CustomPolicyHeader {}

impl HttpContext for CustomPolicyHeader {
    fn on_http_request_headers(&mut self, _num_headers: usize, _end_of_stream: bool) -> Action {
        match self.get_http_request_header("token") {
            Some(token) if is_prime(token.parse().unwrap()) => {
                info!("Executing the Prime Number Logic");
                self.resume_http_request();
                Action::Continue
            }
            _ => {
                self.send_http_response(
                    401,
                    vec![
                        ("Powered-By", "MuleSoft"),
                        ("Custom Policy", "Execution Failed"),
                    ],
                    Some(b"Access forbidden. Custom policy Logic Execution Failed. Try giving a Prime Number Greater than 1\n"),
                );
                Action::Pause
            }
        }
    }
}

pub fn is_prime(number: i32) -> bool {
    for i in 2..(number / 2 + 1) {
        if number % i == 0 {
            return false;
        }
    }
    return number > 1;
}