// Copyright 2021 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

pub mod registration;

use std::{
    borrow::Cow,
    collections::HashMap,
    convert::TryFrom,
    ffi::{OsStr, OsString},
    io::{LineWriter, Write},
    ops::Not,
    path::{Path, PathBuf},
    str::FromStr,
};

use anyhow::{anyhow, Context, Error};
use itertools::Itertools;
use lazy_static::lazy_static;
use log::{debug, warn};
use serde::{Deserialize, Serialize};

use registration::{ensure_user_exists, User};

lazy_static! {
    /// Environment variable: the directory where a given module should be copied.
    ///
    /// Passed to `build` scripts.
    static ref MX_TEST_MODULE_DIR: OsString = OsString::from_str("MX_TEST_MODULE_DIR").unwrap();

    /// Environment variable: the directory where the synapse modules are placed in sub directories.
    ///
    /// Passed to `build` scripts.
    static ref MX_TEST_SYNAPSE_DIR: OsString = OsString::from_str("MX_TEST_SYNAPSE_DIR").unwrap();

    /// Environment variable: a temporary directory where scripts can store data.
    ///
    /// Passed to `build`, `up`, `run`, `down` scripts.
    static ref MX_TEST_SCRIPT_TMPDIR: OsString = OsString::from_str("MX_TEST_SCRIPT_TMPDIR").unwrap();

    /// Environment variable: the directory where we launched the test.
    ///
    /// Passed to `build`, `up`, `run`, `down` scripts.
    static ref MX_TEST_CWD: OsString = OsString::from_str("MX_TEST_CWD").unwrap();

    /// The docker tag used for the Synapse image we produce.
    static ref PATCHED_IMAGE_DOCKER_TAG: OsString = OsString::from_str("mx-tester/synapse").unwrap();
}

#[derive(Debug, Deserialize)]
pub struct DockerConfig {
    #[serde(default)]
    /// A docker network to run the synape containers on.
    /// If the network does not exist it will be created.
    /// If a network is provided, the synapse container will be given the hostname `synapse`.
    /// Defaults to None.
    pub docker_network: Option<String>,

    /// The hostname to give the synapse container on the docker network, if the docker network has been provided.
    /// Defaults to `synapse` but will not be used unless a network is provided in docker_network.
    pub hostname: String,

    /// The docker port mapping configuration to use for the synapse container e.g. `9999:9999`.
    /// Defaults to 9999:9999
    pub port_mapping: String,
}

impl Default for DockerConfig {
    fn default() -> DockerConfig {
        DockerConfig {
            docker_network: None,
            hostname: "synapse".to_string(),
            port_mapping: "9999:9999".to_string(),
        }
    }
}
#[derive(Debug, Deserialize, Serialize)]
pub struct HomeserverConfig {
    /// The name of the homeserver.
    server_name: String,

    /// The URL to communicate to the server with.
    public_baseurl: String,

    #[serde(default = "HomeserverConfig::registration_shared_secret_default")]
    /// The registration shared secret, if provided.
    registration_shared_secret: String,

    #[serde(flatten)]
    /// Any extra fields in the homeserver config
    extra_fields: HashMap<String, serde_yaml::Value>,
}

impl Default for HomeserverConfig {
    fn default() -> HomeserverConfig {
        HomeserverConfig {
            server_name: "localhost:9999".to_string(),
            public_baseurl: "http://localhost:9999".to_string(),
            registration_shared_secret: HomeserverConfig::registration_shared_secret_default(),
            extra_fields: HashMap::new(),
        }
    }
}

impl HomeserverConfig {
    /// Update the homeserver.yaml at the given path (usually one that has been generated by synapse)
    /// with the properties in this struct (which will usually have been provided from mx-tester.yaml)
    pub fn update_homeserver_config(&self, target_homeserver_config: &Path) -> Result<(), Error> {
        let config_file = std::fs::File::open(target_homeserver_config)
            .context("Could not open the homeserver.yaml generated by synapse")?;
        let mut combined_config: serde_yaml::Mapping = serde_yaml::from_reader(config_file)
            .context("The homeserver.yaml generated by synapse is invalid")?;
        let mut insert_value = |key: &str, value: &str| {
            combined_config.insert(serde_yaml::Value::from(key), serde_yaml::Value::from(value));
        };
        insert_value("public_baseurl", &self.public_baseurl);
        insert_value("server_name", &self.server_name);
        insert_value(
            "registration_shared_secret",
            &self.registration_shared_secret,
        );
        for (key, value) in &self.extra_fields {
            combined_config.insert(serde_yaml::Value::from(key.clone()), value.clone());
        }
        let mut config_writer = LineWriter::new(std::fs::File::create(&target_homeserver_config)?);
        config_writer
            .write_all(
                &serde_yaml::to_vec(&combined_config)
                    .context("Could not serialize combined homeserver config")?,
            )
            .context("Could not write combined homeserver config")?;
        Ok(())
    }
    pub fn registration_shared_secret_default() -> String {
        "MX_TESTER_REGISTRATION_DEFAULT".to_string()
    }
}

/// The contents of a mx-tester.yaml
#[derive(Debug, Default, Deserialize)]
pub struct Config {
    /// A name for this test.
    pub name: String,

    /// Modules to install in Synapse.
    #[serde(default)]
    pub modules: Vec<ModuleConfig>,

    #[serde(default)]
    /// Values to pass through into the homeserver.yaml for this synapse.
    pub homeserver_config: HomeserverConfig,

    #[serde(default)]
    /// A script to run at the end of the setup phase.
    pub up: Option<Script>,

    #[serde(default)]
    /// The testing script to run.
    pub run: Option<Script>,

    #[serde(default)]
    /// A script to run at the start of the teardown phase.
    pub down: Option<DownScript>,

    #[serde(default)]
    /// Configuration for the docker network.
    pub docker_config: DockerConfig,

    #[serde(default)]
    /// Any users to register and make available
    pub users: Vec<User>,
}

/// The result of the test, as seen by `down()`.
pub enum Status {
    /// The test was a success.
    Success,

    /// The test was a failure.
    Failure,

    /// The test was not executed at all, we just ran `mx-tester down`.
    Manual,
}

pub enum SynapseVersion {
    /// The latest version of Synapse released on https://hub.docker.com/r/matrixdotorg/synapse/
    ReleasedDockerImage,
    // FIXME: Allow using a version of Synapse that lives in a local directory
    // (this will be sufficient to also implement pulling from github develop)
}
impl SynapseVersion {
    pub fn tag(&self) -> Cow<'static, OsStr> {
        let tag: &'static OsStr = PATCHED_IMAGE_DOCKER_TAG.as_ref();
        tag.into()
    }
}

/// The configuration for a single synapse container.
/// We copy some fields from the main config for simplicity.
#[derive(Debug)]
pub struct ContainerConfig {
    /// The docker port configuration. e.g. `9999:9999`.
    pub port_mapping: String,
    /// The hostname of the synapse container.
    pub hostname: String,
    /// The name of a docker network to place the container in.
    pub docker_network: Option<String>,
    /// The URL that can be used to access this container from the host.
    pub public_url: String,
    /// The name of the synapse server.
    pub server_name: String,
}

impl TryFrom<&Config> for ContainerConfig {
    type Error = Error;
    fn try_from(config: &Config) -> Result<ContainerConfig, Error> {
        let homeserver_config = &config.homeserver_config;
        Ok(ContainerConfig {
            docker_network: config.docker_config.docker_network.clone(),
            port_mapping: config.docker_config.port_mapping.clone(),
            hostname: config.docker_config.hostname.clone(),
            public_url: homeserver_config.public_baseurl.clone(),
            server_name: homeserver_config.server_name.clone(),
        })
    }
}

#[derive(Debug, Deserialize)]
#[serde(transparent)]
pub struct Script {
    /// The lines of the script.
    ///
    /// Passed without change to `std::process::Command`.
    ///
    /// To communicate with the script, clients should use
    /// an exchange file.
    lines: Vec<String>,
}
impl Script {
    /// Substitute anything that looks like the use of environment variable such as $A or ${B} using the entries to the hash map provided.
    /// If there is no matching table entry for $FOO then it will remain as $FOO in the script line.
    fn substitute_env_vars<'a>(
        line: &'a str,
        env: &HashMap<&'static OsStr, OsString>,
    ) -> Cow<'a, str> {
        shellexpand::env_with_context_no_errors(line, |str| match env.get(OsStr::new(str)) {
            Some(value) => value.to_str(),
            _ => None,
        })
    }
    /// Parse a line of script into a std::process::Command and its arguments.
    /// Returns None when there is no command token in the line (e.g. just whitespace).
    fn parse_command(&self, line: &str) -> Option<std::process::Command> {
        let tokens = comma::parse_command(line)?;
        let mut token_stream = tokens.iter();
        let mut command = std::process::Command::new(OsString::from(token_stream.next()?));
        for token in token_stream {
            command.arg(token);
        }
        Some(command)
    }
    pub fn run(&self, env: &HashMap<&'static OsStr, OsString>) -> Result<(), Error> {
        debug!("Running with environment variables {:#?}", env);
        for line in &self.lines {
            let line = Script::substitute_env_vars(line, env);
            let mut command = match self.parse_command(&line) {
                Some(command) => command,
                None => {
                    warn!("Skipping empty line in script {:?}", self.lines);
                    continue;
                }
            };
            let status = command.envs(env).spawn()?.wait()?;
            if !status.success() {
                return Err(anyhow!(
                    "Error running command `{line}`: {status}",
                    line = line,
                    status = status
                ));
            }
        }
        Ok(())
    }
}

/// A script for `build`.
#[derive(Debug, Deserialize)]
pub struct ModuleConfig {
    /// The name of the module.
    ///
    /// This name is used to create a subdirectory.
    name: String,

    /// A script to build and copy the module in the directory
    /// specified by environment variable `MX_TEST_MODULE_DIR`.
    build: Script,
}

/// A script for `down`.
#[derive(Debug, Deserialize)]
pub struct DownScript {
    /// Code to run in case the test is a success.
    success: Option<Script>,

    /// Code to run in case the test is a failure.
    failure: Option<Script>,

    /// Code to run regardless of the result of the test.
    ///
    /// Executed after `success` or `failure`.
    finally: Option<Script>,
}

/// Create a map containing the environment variables that are common
/// to all scripts.
///
/// Callers may add additional variables that are specific to a given
/// script step.
fn shared_env_variables() -> Result<HashMap<&'static OsStr, OsString>, Error> {
    let synapse_root = synapse_root();
    let script_tmpdir = std::env::temp_dir().join("mx-tester").join("scripts");
    std::fs::create_dir_all(&script_tmpdir)
        .with_context(|| format!("Could not create directory {:?}", script_tmpdir,))?;
    let curdir = std::env::current_dir()?;
    let mut env: HashMap<&'static OsStr, _> = HashMap::new();
    env.insert(&*MX_TEST_SYNAPSE_DIR, synapse_root.as_os_str().into());
    env.insert(&*MX_TEST_SCRIPT_TMPDIR, script_tmpdir.as_os_str().into());
    env.insert(&*MX_TEST_CWD, curdir.as_os_str().into());
    Ok(env)
}

fn synapse_root() -> PathBuf {
    std::env::temp_dir().join("mx-tester").join("synapse")
}

/// Rebuild the Synapse image with modules.
pub fn build(config: &[ModuleConfig], version: &SynapseVersion) -> Result<(), Error> {
    // This will break (on purpose) once we extend `SynapseVersion`.
    let SynapseVersion::ReleasedDockerImage = *version;
    let synapse_root = synapse_root();
    std::fs::create_dir_all(&synapse_root)
        .with_context(|| format!("Could not create directory {:?}", synapse_root,))?;
    // Build modules
    let mut env = shared_env_variables()?;
    for module in config {
        let path = synapse_root.join(&module.name);
        env.insert(&*MX_TEST_MODULE_DIR, path.as_os_str().into());
        debug!(
            "Calling build script for module {} with MX_TEST_DIR={:?}",
            &module.name, path
        );
        module
            .build
            .run(&env)
            .context("Error running build script")?;
        debug!("Completed one module.");
    }

    // Prepare Dockerfile including modules.
    let dockerfile_content = format!("
# A custom Dockerfile to rebuild synapse from the official release + plugins

FROM matrixdotorg/synapse:latest

# We need gcc to build pyahocorasick
RUN apt-get update --quiet && apt-get install gcc --yes --quiet

# Show the Synapse version, to aid with debugging.
RUN pip show matrix-synapse

# Copy and install custom modules.
RUN mkdir /mx-tester
{copy}

VOLUME [\"/data\"]

EXPOSE 8008/tcp 8009/tcp 8448/tcp
",
    copy = config.iter()
        // FIXME: We probably want to test what happens with weird characters. Perhaps we'll need to somehow escape module.
        .map(|module| format!("COPY {module} /mx-tester/{module}\nRUN /usr/local/bin/python -m pip install /mx-tester/{module}", module=module.name))
        .format("\n")
);
    debug!("dockerfile {}", dockerfile_content);

    let docker_dir_path = std::env::temp_dir().join("mx-tester").join("docker");
    std::fs::create_dir_all(&docker_dir_path)
        .with_context(|| format!("Could not create directory {:?}", docker_dir_path,))?;
    let dockerfile_path = docker_dir_path.join("Dockerfile");
    std::fs::write(&dockerfile_path, dockerfile_content)
        .with_context(|| format!("Could not write file {:?}", dockerfile_path,))?;

    debug!("Building image with tag {:?}", version.tag());
    if std::process::Command::new("docker")
        .arg("build")
        .args(["--pull", "--no-cache"])
        .arg("-t")
        .arg(version.tag())
        .arg("-f")
        .arg(&dockerfile_path)
        .arg(&synapse_root)
        .spawn()
        .context("Error launching `docker build`")?
        .wait()
        .context("Error waiting for `docker build`")?
        .success()
        .not()
    {
        return Err(anyhow!("Command `docker build` indicated an error"));
    }

    Ok(())
}

/// Generate the data directory and default synapse configuration.
fn generate(
    synapse_data_directory: &Path,
    container_config: &ContainerConfig,
) -> Result<(), Error> {
    // FIXME: I think we're creating tonnes of unnamed garbage containers each time we run this.
    let mut command = std::process::Command::new("docker");
    command
        .arg("run")
        .arg("-e")
        // FIXME: Use server name from config.
        .arg(format!(
            "SYNAPSE_SERVER_NAME={}",
            container_config.server_name
        ))
        .arg("-e")
        .arg("SYNAPSE_REPORT_STATS=no")
        .arg("-e")
        .arg("SYNAPSE_CONFIG_DIR=/data");
    // Ensure that the config files and media can be deleted by the user
    // who launched the program by giving synapse the same uid/gid.
    #[cfg(unix)]
    command
        .arg("-e")
        .arg(format!("UID={}", nix::unistd::getuid()))
        .arg("-e")
        .arg(format!("GID={}", nix::unistd::getegid()));
    if command
        .arg("-p")
        .arg(&container_config.port_mapping)
        .arg("-v")
        .arg(format!(
            "{}:/data",
            &synapse_data_directory
                .to_str()
                .context("Invalid synapse data directory")?
        ))
        .arg(&*PATCHED_IMAGE_DOCKER_TAG)
        .arg("generate")
        .spawn()
        .context("Error launching `docker run` to generate synapse configuration")?
        .wait()
        .context("Error waiting for `docker run` to generate synapse configuration")?
        .success()
        .not()
    {
        return Err(anyhow!(
            "Command `docker run` to generate synapse configuration indicated an error"
        ));
    }
    Ok(())
}

/// Ensures that a container is running for the synapse image.
/// If a network is provided then the container will be given the hostname `synapse`.
fn up_image(
    synapse_data_directory: &Path,
    container_config: &ContainerConfig,
    create_new_container: bool,
) -> Result<(), Error> {
    let container_name = "mx-tester_synapse";
    let container_up = is_container_up(container_name)?;
    if container_up && create_new_container {
        container_stop(container_name).context("Error stopping container")?;
    } else if container_up {
        return Ok(());
    }
    if is_container_built(container_name)? {
        container_rm(container_name).context("Error removing container")?;
    }
    let mut command = std::process::Command::new("docker");
    command.arg("run");
    if let Some(ref network) = container_config.docker_network {
        command
            .arg("--network")
            .arg(network)
            .arg("--hostname")
            .arg(&container_config.hostname);
    }
    // Ensure that the config files and media can be deleted by the user
    // who launched the program by giving synapse the same uid/gid.
    #[cfg(unix)]
    command
        .arg("-e")
        .arg(format!("UID={}", nix::unistd::getuid()))
        .arg("-e")
        .arg(format!("GID={}", nix::unistd::getegid()));
    command
        .arg("--detach")
        .arg("--name")
        .arg("mx-tester_synapse")
        .arg("-p")
        .arg(&container_config.port_mapping)
        .arg("-v")
        .arg(format!(
            "{}:{}",
            &synapse_data_directory
                .to_str()
                .context("Invalid synapse data directory")?,
            "/data"
        ))
        .arg(&*PATCHED_IMAGE_DOCKER_TAG);
    if command
        .spawn()
        .context("Error launching `docker run` to create directories")?
        .wait()
        .context("Error waiting for `docker run` to create directories")?
        .success()
        .not()
    {
        return Err(anyhow!("Command `docker run` indicated an error"));
    }
    Ok(())
}

/// Check whether the named container is currently up.
fn is_container_up(container_name: &str) -> Result<bool, Error> {
    let mut command = std::process::Command::new("docker");
    command
        .arg("container")
        .arg("ps")
        .arg("--no-trunc")
        .arg("--filter")
        .arg(format!("name={}", container_name));
    let output = command.output()?;
    debug!(
        "is_container_up name={} output: {:?}",
        container_name, output
    );
    let all_output = String::from_utf8(output.stdout)
        .context("The output of `docker container ps` is invalid utf8")?;
    Ok(all_output.contains(container_name))
}

/// Check whether a container with this name has been built already.
fn is_container_built(container_name: &str) -> Result<bool, Error> {
    let mut command = std::process::Command::new("docker");
    command
        .arg("container")
        .arg("ls")
        .arg("-a")
        .arg("--no-trunc")
        .arg("--filter")
        .arg(format!("name={}", container_name));
    let output = command.output()?;
    debug!(
        "is_container_built name={} output: {:?}",
        container_name, output
    );
    let all_output = String::from_utf8(output.stdout)
        .context("The output of `docker container ls` is invalid utf8")?;
    Ok(all_output.contains(container_name))
}

/// Remove the named container.
fn container_rm(container_name: &str) -> Result<(), Error> {
    let mut command = std::process::Command::new("docker");
    if command
        .arg("container")
        .arg("rm")
        .arg(container_name)
        .spawn()
        .context("Error launching `docker container rm`")?
        .wait()
        .context("Error waiting for `docker container rm`")?
        .success()
        .not()
    {
        return Err(anyhow!("Command `docker container rm` indicated an error"));
    }
    Ok(())
}

/// Create a docker network and give a name.
fn create_network(network_name: &str) -> Result<(), Error> {
    let mut command = std::process::Command::new("docker");
    command
        .arg("network")
        .arg("create")
        .arg(network_name)
        .output()?;
    Ok(())
}

/// Check if the named docker network has previously been created.
fn is_network_created(network_name: &str) -> Result<bool, Error> {
    let mut command = std::process::Command::new("docker");
    command
        .arg("network")
        .arg("ls")
        .arg("--no-trunc")
        .arg("--filter")
        .arg(format!("name={}", network_name));
    let output = command.output()?;
    debug!(
        "is_network_created name={} output: {:?}",
        network_name, output
    );
    let all_output = String::from_utf8(output.stdout)
        .context("The output of `docker network ls` is invalid utf")?;
    Ok(all_output.contains(network_name))
}

/// Ensure the named docker network exists.
fn ensure_network_exists(network_name: &str) -> Result<bool, Error> {
    let network_exists = is_network_created(network_name)?;
    if !network_exists {
        create_network(network_name)?;
    }
    Ok(network_exists)
}

/// Bring things up. Returns any environment variables to pass to the run script.
pub async fn up(
    version: &SynapseVersion,
    config: &Config,
    container_config: &ContainerConfig,
    homeserver_config: &HomeserverConfig,
) -> Result<(), Error> {
    // This will break (on purpose) once we extend `SynapseVersion`.
    let SynapseVersion::ReleasedDockerImage = *version;

    let synapse_data_directory = synapse_root().join("data");
    std::fs::create_dir_all(&synapse_data_directory)
        .with_context(|| format!("Cannot create directory {:?}", synapse_data_directory))?;
    if let Some(ref docker_network) = container_config.docker_network {
        ensure_network_exists(docker_network)?;
    }
    debug!("generating synapse data");
    generate(&synapse_data_directory, container_config)
        .context("Error while generating synapse data")?;
    debug!("done generating");
    // Apply config from mx-tester.yml to the homeserver.yaml that was just made
    homeserver_config
        .update_homeserver_config(&synapse_data_directory.join("homeserver.yaml"))
        .context("Error updating homeserver config")?;
    // FIXME: Allow configuration of recreating container if the image has been rebuilt.
    up_image(&synapse_data_directory, container_config, false)
        .context("Error bringing up image")?;

    for user in &config.users {
        ensure_user_exists(
            &config.homeserver_config.public_baseurl,
            &config.homeserver_config.registration_shared_secret,
            user,
        )
        .await
        .with_context(|| anyhow!("Could not setup user {}", user.localname))?;
    }
    if let Some(ref script) = config.up {
        let env = shared_env_variables()?;
        script.run(&env).context("Error running `up` script")?;
    }
    Ok(())
}

/// Stop a container.
pub fn container_stop(container_name: &str) -> Result<(), Error> {
    if std::process::Command::new("docker")
        .arg("stop")
        .arg(container_name)
        .stdout(std::process::Stdio::inherit())
        .stderr(std::process::Stdio::inherit())
        .spawn()
        .context("Error launching `docker stop`")?
        .wait()
        .context("Error waiting for `docker stop`")?
        .success()
        .not()
    {
        return Err(anyhow!("Command `docker stop` indicated an error"));
    }
    Ok(())
}

/// Bring things down.
pub fn down(
    version: &SynapseVersion,
    script: &Option<DownScript>,
    status: Status,
) -> Result<(), Error> {
    // This will break (on purpose) once we extend `SynapseVersion`.
    let SynapseVersion::ReleasedDockerImage = *version;

    if let Some(ref down_script) = *script {
        let env = shared_env_variables()?;
        // First run on_failure/on_success.
        // Store errors for later.
        let result = match (status, down_script) {
            (
                Status::Failure,
                DownScript {
                    failure: Some(ref on_failure),
                    ..
                },
            ) => on_failure
                .run(&env)
                .context("Error while running script `down/failure`"),
            (
                Status::Success,
                DownScript {
                    success: Some(ref on_success),
                    ..
                },
            ) => on_success
                .run(&env)
                .context("Error while running script `down/success`"),
            _ => Ok(()),
        };
        // Then run on_always.
        if let Some(ref on_always) = down_script.finally {
            on_always
                .run(&env)
                .context("Error while running script `down/finally`")?;
        }
        // Report any error from `on_failure` or `on_success`.
        result?
    }
    debug!("Taking down synapse.");
    container_stop("mx-tester_synapse").context("Error stopping container")?;
    Ok(())
}

/// Run the testing script.
pub fn run(config: &Config) -> Result<(), Error> {
    if let Some(ref code) = config.run {
        let env = shared_env_variables()?;
        code.run(&env).context("Error running `run` script")?;
    }
    Ok(())
}