use wasmer_deploy_schema::schema::{DeploymentV1, WorkloadRunnerV1};

/// Errors that can occur during module validation.
#[derive(Clone, Debug)]
pub enum ModuleValidationError {
    WebProxyMissingSyscalls { syscalls: Vec<String> },
    WcgiMissingSyscalls { syscalls: Vec<String> },
}

impl ModuleValidationError {
    pub fn explain(&self) -> String {
        match self {
            Self::WebProxyMissingSyscalls { syscalls } => {
                format!(
                    "web_proxy modules must read and write to sockets, \
                    but the module does not import the WASI functions required \
                    to do so. \\n
                    Missing functions: {}",
                    syscalls.join(", ")
                )
            }
            Self::WcgiMissingSyscalls { syscalls } => {
                format!(
                    "wcgi modules must read from stdin and write to stdout, \
                    but the module does not import the WASI functions required \
                    to do so. \n\
                    Missing functions: {}",
                    syscalls.join(", ")
                )
            }
        }
    }
}

impl std::fmt::Display for ModuleValidationError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::WebProxyMissingSyscalls { syscalls } => {
                write!(
                    f,
                    "Module does not import required syscalls: {}",
                    syscalls.join(", ")
                )
            }
            Self::WcgiMissingSyscalls { syscalls } => {
                write!(
                    f,
                    "Module does not import required syscalls: {}",
                    syscalls.join(", ")
                )
            }
        }
    }
}

impl std::error::Error for ModuleValidationError {}

/// Validate a WASM module for usage as a webproxy.
///
/// This only checks for the presence of some mandatory syscalls.
pub fn validate_parse_module_webproxy(bytes: &[u8]) -> Result<(), ModuleValidationError> {
    let mut sock_listen = false;
    let mut sock_accept = false;
    let mut sock_recv = false;
    let mut sock_send = false;

    for payload in wasmparser::Parser::new(0).parse_all(bytes) {
        // NOTE: validation deliberately ignores parse errors to be more compatible.
        if let Ok(wasmparser::Payload::ImportSection(imports)) = payload {
            for import in imports {
                if let Ok(import) = import {
                    if !matches!(import.ty, wasmparser::TypeRef::Func(_)) {
                        continue;
                    }

                    // Checking is deliberately permissive, just check that the
                    // absolutely necessary syscalls are used.
                    if import.name.contains("sock_listen") {
                        sock_listen = true;
                    } else if import.name.contains("sock_accept") {
                        sock_accept = true;
                    } else if import.name.contains("sock_recv") {
                        sock_recv = true;
                    } else if import.name.contains("sock_send") {
                        sock_send = true;
                    }
                }
            }
        }
    }

    let mut missing = Vec::new();
    if !sock_listen {
        // NOTE: not checking this, because WASI (non-wasix) would also work
        // without listen, if there was a pre-defined socket.
        // missing.push("sock_listen".to_string());
    }
    if !sock_accept {
        missing.push("sock_accept".to_string());
    }
    if !sock_recv {
        missing.push("sock_recv".to_string());
    }
    if !sock_send {
        missing.push("sock_send".to_string());
    }

    if missing.is_empty() {
        Ok(())
    } else {
        Err(ModuleValidationError::WebProxyMissingSyscalls { syscalls: missing })
    }
}

/// Validate a WASM module for usage with wcgi .
///
/// This only checks for the presence of some mandatory syscalls.
pub fn validate_parse_module_wcgi(bytes: &[u8]) -> Result<(), ModuleValidationError> {
    let mut fd_read = false;
    let mut fd_write = false;
    let mut fd_pipe = false;

    for payload in wasmparser::Parser::new(0).parse_all(bytes) {
        // NOTE: validation deliberately ignores parse errors to be more compatible.
        if let Ok(wasmparser::Payload::ImportSection(imports)) = payload {
            for import in imports {
                if let Ok(import) = import {
                    if !matches!(import.ty, wasmparser::TypeRef::Func(_)) {
                        continue;
                    }

                    // Checking is deliberately permissive, just check that the
                    // absolutely necessary syscalls are used.
                    if import.name.contains("fd_read") {
                        fd_read = true;
                    } else if import.name.contains("fd_write") {
                        fd_write = true;
                    } else if import.name.contains("fd_pipe") {
                        fd_pipe = true;
                    }
                }
            }
        }
    }

    let mut missing = Vec::new();
    if !fd_read {
        missing.push("fd_read".to_string());
    }
    if !(fd_write || fd_pipe) {
        missing.push("fd_write|fd_pipe".to_string());
    }

    if missing.is_empty() {
        Ok(())
    } else {
        Err(ModuleValidationError::WcgiMissingSyscalls { syscalls: missing })
    }
}

#[derive(Clone, Debug)]
pub enum DeploymentValidationError {
    Module(ModuleValidationError),
    UnsupportedRunner(String),
}

impl std::fmt::Display for DeploymentValidationError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Module(e) => e.fmt(f),
            Self::UnsupportedRunner(e) => write!(f, "Unsupported runner: {}", e),
        }
    }
}

/// Validate a module against the given deployment config.
///
/// Will check that the module is compatible with the runner.
pub fn validate_deployment_module(
    depl: &DeploymentV1,
    module_wasm: &[u8],
) -> Result<(), DeploymentValidationError> {
    match &depl.workload.runner {
        WorkloadRunnerV1::Wasm(_) => Err(DeploymentValidationError::UnsupportedRunner(
            "wasm".to_string(),
        )),
        WorkloadRunnerV1::WebcCommand(_) => Err(DeploymentValidationError::UnsupportedRunner(
            "webc_command".to_string(),
        )),
        WorkloadRunnerV1::TcpProxy(_) => Err(DeploymentValidationError::UnsupportedRunner(
            "tcp_proxy".to_string(),
        )),
        WorkloadRunnerV1::WCgi(_) => {
            validate_parse_module_wcgi(module_wasm).map_err(DeploymentValidationError::Module)
        }
        WorkloadRunnerV1::WebProxy(_) => {
            validate_parse_module_webproxy(module_wasm).map_err(DeploymentValidationError::Module)
        }
    }
}

#[cfg(test)]
mod tests {
    use std::path::PathBuf;

    use wasmer_deploy_schema::schema::{
        RunnerWCgiV1, WorkloadRunnerWasmSourceLocalPathV1, WorkloadRunnerWasmSourceV1,
    };

    use super::*;

    fn root_path() -> PathBuf {
        std::env::var("CARGO_MANIFEST_DIR")
            .map(PathBuf::from)
            .unwrap()
            .parent()
            .unwrap()
            .parent()
            .unwrap()
            .to_owned()
    }

    fn tests_path() -> PathBuf {
        root_path().join("wasm-tests").join("compiled")
    }

    #[test]
    fn test_validate_wcgi_valid() {
        let path = tests_path().join("local").join("wcgi-hello.wasm");
        let contents = std::fs::read(&path).unwrap();
        validate_parse_module_wcgi(&contents).unwrap();
    }

    #[test]
    fn test_validate_wcgi_invalid() {
        let path = tests_path().join("vendor").join("empty.wasm");
        let contents = std::fs::read(&path).unwrap();

        let res = validate_parse_module_wcgi(&contents);
        assert!(matches!(
            res,
            Err(ModuleValidationError::WcgiMissingSyscalls { .. })
        ));
    }

    #[test]
    fn test_validate_webproxy_valid() {
        let path = tests_path().join("vendor").join("static-web-server.wasm");
        let contents = std::fs::read(&path).unwrap();

        validate_parse_module_webproxy(&contents).unwrap();
    }

    #[test]
    fn test_validate_webproxy_invalid() {
        let path = tests_path().join("local").join("wcgi-hello.wasm");
        let contents = std::fs::read(&path).unwrap();

        let res = validate_parse_module_webproxy(&contents);
        assert!(matches!(
            res,
            Err(ModuleValidationError::WebProxyMissingSyscalls { .. })
        ));
    }

    #[test]
    fn test_validate_deployment_wcgi() {
        // Valid

        let path = tests_path().join("local").join("wcgi-hello.wasm");
        let contents = std::fs::read(&path).unwrap();

        let depl = DeploymentV1 {
            name: "name".to_string(),
            workload: wasmer_deploy_schema::schema::WorkloadV1 {
                name: None,
                capabilities: Default::default(),
                runner: WorkloadRunnerV1::WCgi(RunnerWCgiV1 {
                    source: WorkloadRunnerWasmSourceV1::LocalPath(
                        WorkloadRunnerWasmSourceLocalPathV1 {
                            path: ".".to_string(),
                        },
                    ),
                    dialect: None,
                }),
            },
        };

        validate_deployment_module(&depl, &contents).unwrap();

        // Invalid.

        let path = tests_path().join("vendor").join("empty.wasm");
        let contents = std::fs::read(&path).unwrap();

        assert!(matches!(
            validate_deployment_module(&depl, &contents),
            Err(DeploymentValidationError::Module(
                ModuleValidationError::WcgiMissingSyscalls { .. }
            ))
        ));
    }

    #[test]
    fn test_validate_deployment_web_proxy() {
        // Valid

        let path = tests_path().join("vendor").join("static-web-server.wasm");
        let contents = std::fs::read(&path).unwrap();

        let depl = DeploymentV1 {
            name: "name".to_string(),
            workload: wasmer_deploy_schema::schema::WorkloadV1 {
                name: None,
                capabilities: Default::default(),
                runner: WorkloadRunnerV1::WebProxy(
                    wasmer_deploy_schema::schema::RunnerWebProxyV1 {
                        source: WorkloadRunnerWasmSourceV1::LocalPath(
                            WorkloadRunnerWasmSourceLocalPathV1 {
                                path: ".".to_string(),
                            },
                        ),
                    },
                ),
            },
        };

        validate_deployment_module(&depl, &contents).unwrap();

        // Invalid.

        let path = tests_path().join("local").join("wcgi-hello.wasm");
        let contents = std::fs::read(&path).unwrap();

        assert!(matches!(
            validate_deployment_module(&depl, &contents),
            Err(DeploymentValidationError::Module(
                ModuleValidationError::WebProxyMissingSyscalls { .. }
            ))
        ));
    }
}