// Copyright 2016 Mozilla Foundation
//
// 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.

use atty::Stream;
use byteorder::{BigEndian, ByteOrder};
use crate::client::{connect_to_server, connect_with_retry, ServerConnection};
use crate::cmdline::{Command, StatsFormat};
use crate::compiler::ColorMode;
use crate::config::Config;
use futures::Future;
use crate::jobserver::Client;
use log::Level::Trace;
use crate::mock_command::{CommandChild, CommandCreatorSync, ProcessCommandCreator, RunCommand};
use crate::protocol::{Compile, CompileFinished, CompileResponse, Request, Response};
use crate::server::{self, ServerInfo, ServerStartup};
use std::env;
use std::ffi::{OsStr, OsString};
use std::fs::{File, OpenOptions};
use std::io::{self, Write};
#[cfg(unix)]
use std::os::unix::process::ExitStatusExt;
use std::path::Path;
use std::process;
use strip_ansi_escapes::Writer;
use tokio::runtime::current_thread::Runtime;
use tokio_io::io::read_exact;
use tokio_io::AsyncRead;
use tokio_timer::Timeout;
use crate::util::daemonize;
use which::which_in;

use crate::errors::*;

/// The default sccache server port.
pub const DEFAULT_PORT: u16 = 4226;

/// The number of milliseconds to wait for server startup.
const SERVER_STARTUP_TIMEOUT_MS: u32 = 5000;

/// Get the port on which the server should listen.
fn get_port() -> u16 {
    env::var("SCCACHE_SERVER_PORT")
        .ok()
        .and_then(|s| s.parse().ok())
        .unwrap_or(DEFAULT_PORT)
}

fn read_server_startup_status<R: AsyncRead>(
    server: R,
) -> impl Future<Item = ServerStartup, Error = Error> {
    // This is an async equivalent of ServerConnection::read_one_response
    read_exact(server, [0u8; 4])
        .map_err(Error::from)
        .and_then(|(server, bytes)| {
            let len = BigEndian::read_u32(&bytes);
            let data = vec![0; len as usize];
            read_exact(server, data)
                .map_err(Error::from)
                .and_then(|(_server, data)| Ok(bincode::deserialize(&data)?))
        })
}

/// Re-execute the current executable as a background server, and wait
/// for it to start up.
#[cfg(not(windows))]
fn run_server_process() -> Result<ServerStartup> {
    use futures::Stream;
    use std::time::Duration;
    use tempdir::TempDir;

    trace!("run_server_process");
    let tempdir = TempDir::new("sccache")?;
    let socket_path = tempdir.path().join("sock");
    let mut runtime = Runtime::new()?;
    let listener = tokio_uds::UnixListener::bind(&socket_path)?;
    let exe_path = env::current_exe()?;
    let _child = process::Command::new(exe_path)
        .env("SCCACHE_START_SERVER", "1")
        .env("SCCACHE_STARTUP_NOTIFY", &socket_path)
        .env("RUST_BACKTRACE", "1")
        .spawn()?;

    let startup = listener.incoming().into_future().map_err(|e| e.0);
    let startup = startup.map_err(Error::from).and_then(|(socket, _rest)| {
        let socket = socket.unwrap(); // incoming() never returns None
        read_server_startup_status(socket)
    });

    let timeout = Duration::from_millis(SERVER_STARTUP_TIMEOUT_MS.into());
    let timeout = Timeout::new(startup, timeout).or_else(|err| {
        if err.is_elapsed() {
            Ok(ServerStartup::TimedOut)
        } else if err.is_inner() {
            Err(err.into_inner().unwrap().into())
        } else {
            Err(err.into_timer().unwrap().into())
        }
    });
    runtime.block_on(timeout)
}

#[cfg(not(windows))]
fn redirect_stderr(f: File) -> Result<()> {
    use libc::dup2;
    use std::os::unix::io::IntoRawFd;
    // Ignore errors here.
    unsafe {
        dup2(f.into_raw_fd(), 2);
    }
    Ok(())
}

#[cfg(windows)]
fn redirect_stderr(f: File) -> Result<()> {
    use std::os::windows::io::IntoRawHandle;
    use winapi::um::processenv::SetStdHandle;
    use winapi::um::winbase::STD_ERROR_HANDLE;
    // Ignore errors here.
    unsafe {
        SetStdHandle(STD_ERROR_HANDLE, f.into_raw_handle());
    }
    Ok(())
}

/// If `SCCACHE_ERROR_LOG` is set, redirect stderr to it.
fn redirect_error_log() -> Result<()> {
    let name = match env::var("SCCACHE_ERROR_LOG") {
        Ok(filename) => filename,
        _ => return Ok(()),
    };
    let f = OpenOptions::new().create(true).append(true).open(name)?;
    redirect_stderr(f)
}

/// Re-execute the current executable as a background server.
#[cfg(windows)]
fn run_server_process() -> Result<ServerStartup> {
    use futures::future;
    use std::mem;
    use std::os::windows::ffi::OsStrExt;
    use std::ptr;
    use std::time::Duration;
    use tokio::reactor::Handle;
    use tokio_named_pipes::NamedPipe;
    use uuid::Uuid;
    use winapi::shared::minwindef::{DWORD, FALSE, LPVOID, TRUE};
    use winapi::um::processthreadsapi::{CreateProcessW, PROCESS_INFORMATION, STARTUPINFOW};
    use winapi::um::winbase::{
        CREATE_NEW_PROCESS_GROUP, CREATE_UNICODE_ENVIRONMENT, DETACHED_PROCESS,
    };
    use winapi::um::handleapi::CloseHandle;

    trace!("run_server_process");

    // Create a mini event loop and register our named pipe server
    let mut runtime = Runtime::new()?;
    let pipe_name = format!(r"\\.\pipe\{}", Uuid::new_v4().to_simple_ref());
    let server = runtime.block_on(future::lazy(|| {
        NamedPipe::new(&pipe_name, &Handle::current())
    }))?;

    // Connect a client to our server, and we'll wait below if it's still in
    // progress.

    match server.connect() {
        Ok(()) => {}
        Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {}
        Err(e) => return Err(e.into()),
    }

    // Spawn a server which should come back and connect to us
    let exe_path = env::current_exe()?;
    let mut exe = OsStr::new(&exe_path)
        .encode_wide()
        .chain(Some(0u16))
        .collect::<Vec<u16>>();
    let mut envp = {
        let mut v = vec![];
        let extra_vars = vec![
            (OsString::from("SCCACHE_START_SERVER"), OsString::from("1")),
            (
                OsString::from("SCCACHE_STARTUP_NOTIFY"),
                OsString::from(&pipe_name),
            ),
            (OsString::from("RUST_BACKTRACE"), OsString::from("1")),
        ];
        for (key, val) in env::vars_os().chain(extra_vars) {
            v.extend(
                key.encode_wide()
                    .chain(Some('=' as u16))
                    .chain(val.encode_wide())
                    .chain(Some(0)),
            );
        }
        v.push(0);
        v
    };

    // TODO: Expose `bInheritHandles` argument of `CreateProcessW` through the
    //       standard library's `Command` type and then use that instead.
    let mut pi = PROCESS_INFORMATION {
        hProcess: ptr::null_mut(),
        hThread: ptr::null_mut(),
        dwProcessId: 0,
        dwThreadId: 0,
    };
    let mut si: STARTUPINFOW = unsafe { mem::zeroed() };
    si.cb = mem::size_of::<STARTUPINFOW>() as DWORD;
    if unsafe {
        CreateProcessW(
            exe.as_mut_ptr(),
            ptr::null_mut(),
            ptr::null_mut(),
            ptr::null_mut(),
            FALSE,
            CREATE_UNICODE_ENVIRONMENT | DETACHED_PROCESS | CREATE_NEW_PROCESS_GROUP,
            envp.as_mut_ptr() as LPVOID,
            ptr::null(),
            &mut si,
            &mut pi,
        ) == TRUE
    } {
        unsafe {
            CloseHandle(pi.hProcess);
            CloseHandle(pi.hThread);
        }
    } else {
        return Err(io::Error::last_os_error().into());
    }

    let result = read_server_startup_status(server);

    let timeout = Duration::from_millis(SERVER_STARTUP_TIMEOUT_MS.into());
    let timeout = Timeout::new(result, timeout).or_else(|err| {
        if err.is_elapsed() {
            Ok(ServerStartup::TimedOut)
        } else if err.is_inner() {
            Err(err.into_inner().unwrap().into())
        } else {
            Err(err.into_timer().unwrap().into())
        }
    });
    runtime.block_on(timeout)
}

/// Attempt to connect to an sccache server listening on `port`, or start one if no server is running.
fn connect_or_start_server(port: u16) -> Result<ServerConnection> {
    trace!("connect_or_start_server({})", port);
    match connect_to_server(port) {
        Ok(server) => Ok(server),
        Err(ref e)
            if e.kind() == io::ErrorKind::ConnectionRefused
                || e.kind() == io::ErrorKind::TimedOut =>
        {
            // If the connection was refused we probably need to start
            // the server.
            //TODO: check startup value!
            let _startup = run_server_process()?;
            let server = connect_with_retry(port)?;
            Ok(server)
        }
        Err(e) => Err(e.into()),
    }
}

/// Send a `ZeroStats` request to the server, and return the `ServerInfo` request if successful.
pub fn request_zero_stats(mut conn: ServerConnection) -> Result<ServerInfo> {
    debug!("request_stats");
    let response = conn.request(Request::ZeroStats).chain_err(|| {
        "failed to send zero statistics command to server or failed to receive respone"
    })?;
    if let Response::Stats(stats) = response {
        Ok(stats)
    } else {
        bail!("Unexpected server response!")
    }
}

/// Send a `GetStats` request to the server, and return the `ServerInfo` request if successful.
pub fn request_stats(mut conn: ServerConnection) -> Result<ServerInfo> {
    debug!("request_stats");
    let response = conn
        .request(Request::GetStats)
        .chain_err(|| "Failed to send data to or receive data from server")?;
    if let Response::Stats(stats) = response {
        Ok(stats)
    } else {
        bail!("Unexpected server response!")
    }
}

/// Send a `Shutdown` request to the server, and return the `ServerInfo` contained within the response if successful.
pub fn request_shutdown(mut conn: ServerConnection) -> Result<ServerInfo> {
    debug!("request_shutdown");
    //TODO: better error mapping
    let response = conn
        .request(Request::Shutdown)
        .chain_err(|| "Failed to send data to or receive data from server")?;
    if let Response::ShuttingDown(stats) = response {
        Ok(stats)
    } else {
        bail!("Unexpected server response!")
    }
}

/// Send a `Compile` request to the server, and return the server response if successful.
fn request_compile<W, X, Y>(
    conn: &mut ServerConnection,
    exe: W,
    args: &Vec<X>,
    cwd: Y,
    env_vars: Vec<(OsString, OsString)>,
) -> Result<CompileResponse>
where
    W: AsRef<Path>,
    X: AsRef<OsStr>,
    Y: AsRef<Path>,
{
    let req = Request::Compile(Compile {
        exe: exe.as_ref().to_owned().into(),
        cwd: cwd.as_ref().to_owned().into(),
        args: args.iter().map(|a| a.as_ref().to_owned()).collect(),
        env_vars: env_vars,
    });
    trace!("request_compile: {:?}", req);
    //TODO: better error mapping?
    let response = conn
        .request(req)
        .chain_err(|| "Failed to send data to or receive data from server")?;
    if let Response::Compile(response) = response {
        Ok(response)
    } else {
        bail!("Unexpected response from server")
    }
}

/// Return the signal that caused a process to exit from `status`.
#[cfg(unix)]
#[allow(dead_code)]
fn status_signal(status: process::ExitStatus) -> Option<i32> {
    status.signal()
}

/// Not implemented for non-Unix.
#[cfg(not(unix))]
#[allow(dead_code)]
fn status_signal(_status: process::ExitStatus) -> Option<i32> {
    None
}

/// Handle `response`, the output from running a compile on the server.
/// Return the compiler exit status.
fn handle_compile_finished(
    response: CompileFinished,
    stdout: &mut dyn Write,
    stderr: &mut dyn Write,
) -> Result<i32> {
    trace!("handle_compile_finished");
    fn write_output(
        stream: Stream,
        writer: &mut dyn Write,
        data: &[u8],
        color_mode: ColorMode,
    ) -> Result<()> {
        // rustc uses the `termcolor` crate which explicitly checks for TERM=="dumb", so
        // match that behavior here.
        let dumb_term = env::var("TERM").map(|v| v == "dumb").unwrap_or(false);
        // If the compiler options explicitly requested color output, or if this output stream
        // is a terminal and the compiler options didn't explicitly request non-color output,
        // then write the compiler output directly.
        if color_mode == ColorMode::On
            || (!dumb_term && atty::is(stream) && color_mode != ColorMode::Off)
        {
            writer.write_all(data)?;
        } else {
            // Remove escape codes (and thus colors) while writing.
            let mut writer = Writer::new(writer);
            writer.write_all(data)?;
        }
        Ok(())
    }
    // It might be nice if the server sent stdout/stderr as the process
    // ran, but then it would have to also save them in the cache as
    // interleaved streams to really make it work.
    write_output(
        Stream::Stdout,
        stdout,
        &response.stdout,
        response.color_mode,
    )?;
    write_output(
        Stream::Stderr,
        stderr,
        &response.stderr,
        response.color_mode,
    )?;

    if let Some(ret) = response.retcode {
        trace!("compiler exited with status {}", ret);
        Ok(ret)
    } else if let Some(signal) = response.signal {
        println!("Compiler killed by signal {}", signal);
        Ok(-2)
    } else {
        println!("Missing compiler exit status!");
        Ok(-3)
    }
}

/// Handle `response`, the response from sending a `Compile` request to the server. Return the compiler exit status.
///
/// If the server returned `CompileStarted`, wait for a `CompileFinished` and
/// print the results.
///
/// If the server returned `UnhandledCompile`, run the compilation command
/// locally using `creator` and return the result.
fn handle_compile_response<T>(
    mut creator: T,
    runtime: &mut Runtime,
    conn: &mut ServerConnection,
    response: CompileResponse,
    exe: &Path,
    cmdline: Vec<OsString>,
    cwd: &Path,
    stdout: &mut dyn Write,
    stderr: &mut dyn Write,
) -> Result<i32>
where
    T: CommandCreatorSync,
{
    match response {
        CompileResponse::CompileStarted => {
            debug!("Server sent CompileStarted");
            // Wait for CompileFinished.
            match conn.read_one_response() {
                Ok(Response::CompileFinished(result)) => {
                    return handle_compile_finished(result, stdout, stderr)
                }
                Ok(_) => bail!("unexpected response from server"),
                Err(Error(ErrorKind::Io(ref e), _)) if e.kind() == io::ErrorKind::UnexpectedEof => {
                    writeln!(
                        io::stderr(),
                        "warning: sccache server looks like it shut down \
                         unexpectedly, compiling locally instead"
                    ).unwrap();
                }
                Err(e) => {
                    return Err(e).chain_err(|| {
                        //TODO: something better here?
                        "error reading compile response from server"
                    });
                }
            }
        }
        CompileResponse::UnsupportedCompiler => {
            debug!("Server sent UnsupportedCompiler");
            bail!("Compiler not supported");
        }
        CompileResponse::UnhandledCompile => {
            debug!("Server sent UnhandledCompile");
        }
    };

    let mut cmd = creator.new_command_sync(exe);
    cmd.args(&cmdline).current_dir(cwd);
    if log_enabled!(Trace) {
        trace!("running command: {:?}", cmd);
    }
    match runtime.block_on(
        cmd.spawn()
            .and_then(|c| c.wait().chain_err(|| "failed to wait for child")),
    ) {
        Ok(status) => {
            Ok(status.code().unwrap_or_else(|| {
                if let Some(sig) = status_signal(status) {
                    println!("Compile terminated by signal {}", sig);
                }
                // Arbitrary.
                2
            }))
        }
        Err(e) => Err(e),
    }
}

/// Send a `Compile` request to the sccache server `conn`, and handle the response.
///
/// The first entry in `cmdline` will be looked up in `path` if it is not
/// an absolute path.
/// See `request_compile` and `handle_compile_response`.
pub fn do_compile<T>(
    creator: T,
    runtime: &mut Runtime,
    mut conn: ServerConnection,
    exe: &Path,
    cmdline: Vec<OsString>,
    cwd: &Path,
    path: Option<OsString>,
    env_vars: Vec<(OsString, OsString)>,
    stdout: &mut dyn Write,
    stderr: &mut dyn Write,
) -> Result<i32>
where
    T: CommandCreatorSync,
{
    trace!("do_compile");
    let exe_path = which_in(exe, path, &cwd)?;
    let res = request_compile(&mut conn, &exe_path, &cmdline, &cwd, env_vars)?;
    handle_compile_response(
        creator, runtime, &mut conn, res, &exe_path, cmdline, cwd, stdout, stderr,
    )
}

/// Run `cmd` and return the process exit status.
pub fn run_command(cmd: Command) -> Result<i32> {
    // Config isn't required for all commands, but if it's broken then we should flag
    // it early and loudly.
    let config = &Config::load()?;

    match cmd {
        Command::ShowStats(fmt) => {
            trace!("Command::ShowStats({:?})", fmt);
            let srv = connect_or_start_server(get_port())?;
            let stats = request_stats(srv).chain_err(|| "failed to get stats from server")?;
            match fmt {
                StatsFormat::text => stats.print(),
                StatsFormat::json => serde_json::to_writer(&mut io::stdout(), &stats)?,
            }
        }
        Command::InternalStartServer => {
            trace!("Command::InternalStartServer");
            // Can't report failure here, we're already daemonized.
            daemonize()?;
            redirect_error_log()?;
            server::start_server(config, get_port())?;
        }
        Command::StartServer => {
            trace!("Command::StartServer");
            println!("Starting sccache server...");
            let startup = run_server_process().chain_err(|| "failed to start server process")?;
            match startup {
                ServerStartup::Ok { port } => {
                    if port != DEFAULT_PORT {
                        println!("Listening on port {}", port);
                    }
                }
                ServerStartup::TimedOut => bail!("Timed out waiting for server startup"),
                ServerStartup::Err { reason } => bail!("Server startup failed: {}", reason),
            }
        }
        Command::StopServer => {
            trace!("Command::StopServer");
            println!("Stopping sccache server...");
            let server =
                connect_to_server(get_port()).chain_err(|| "couldn't connect to server")?;
            let stats = request_shutdown(server)?;
            stats.print();
        }
        Command::ZeroStats => {
            trace!("Command::ZeroStats");
            let conn = connect_or_start_server(get_port())?;
            let stats = request_zero_stats(conn).chain_err(|| "couldn't zero stats on server")?;
            stats.print();
        }
        #[cfg(feature = "dist-client")]
        Command::DistAuth => {
            use crate::config;
            use crate::dist;
            use url::Url;

            match &config.dist.auth {
                config::DistAuth::Token { .. } => {
                    info!("No authentication needed for type 'token'")
                }
                config::DistAuth::Oauth2CodeGrantPKCE {
                    client_id,
                    auth_url,
                    token_url,
                } => {
                    let cached_config = config::CachedConfig::load()?;

                    let parsed_auth_url = Url::parse(auth_url)
                        .map_err(|_| format!("Failed to parse URL {}", auth_url))?;
                    let token = dist::client_auth::get_token_oauth2_code_grant_pkce(
                        client_id,
                        parsed_auth_url,
                        token_url,
                    )?;

                    cached_config
                        .with_mut(|c| {
                            c.dist.auth_tokens.insert(auth_url.to_owned(), token);
                        }).chain_err(|| "Unable to save auth token")?;
                    println!("Saved token")
                }
                config::DistAuth::Oauth2Implicit {
                    client_id,
                    auth_url,
                } => {
                    let cached_config = config::CachedConfig::load()?;

                    let parsed_auth_url = Url::parse(auth_url)
                        .map_err(|_| format!("Failed to parse URL {}", auth_url))?;
                    let token =
                        dist::client_auth::get_token_oauth2_implicit(client_id, parsed_auth_url)?;

                    cached_config
                        .with_mut(|c| {
                            c.dist.auth_tokens.insert(auth_url.to_owned(), token);
                        }).chain_err(|| "Unable to save auth token")?;
                    println!("Saved token")
                }
            };
        }
        #[cfg(not(feature = "dist-client"))]
        Command::DistAuth => bail!(
            "Distributed compilation not compiled in, please rebuild with the dist-client feature"
        ),
        #[cfg(feature = "dist-client")]
        Command::PackageToolchain(executable, out) => {
            use crate::compiler;
            use futures_cpupool::CpuPool;

            trace!("Command::PackageToolchain({})", executable.display());
            let mut runtime = Runtime::new()?;
            let jobserver = unsafe { Client::new() };
            let creator = ProcessCommandCreator::new(&jobserver);
            let env: Vec<_> = env::vars_os().collect();
            let pool = CpuPool::new(1);
            let out_file = File::create(out)?;

            let compiler = compiler::get_compiler_info(&creator, &executable, &env, &pool);
            let packager = compiler.map(|c| c.get_toolchain_packager());
            let res = packager.and_then(|p| p.write_pkg(out_file));
            runtime.block_on(res)?
        }
        #[cfg(not(feature = "dist-client"))]
        Command::PackageToolchain(_executable, _out) => bail!(
            "Toolchain packaging not compiled in, please rebuild with the dist-client feature"
        ),
        Command::Compile {
            exe,
            cmdline,
            cwd,
            env_vars,
        } => {
            trace!("Command::Compile {{ {:?}, {:?}, {:?} }}", exe, cmdline, cwd);
            let jobserver = unsafe { Client::new() };
            let conn = connect_or_start_server(get_port())?;
            let mut runtime = Runtime::new()?;
            let res = do_compile(
                ProcessCommandCreator::new(&jobserver),
                &mut runtime,
                conn,
                exe.as_ref(),
                cmdline,
                &cwd,
                env::var_os("PATH"),
                env_vars,
                &mut io::stdout(),
                &mut io::stderr(),
            );
            return res.chain_err(|| "failed to execute compile");
        }
    }

    Ok(0)
}
