// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

#![allow(clippy::len_zero)]

use crate::{
    asan::{add_asan_log_env, check_asan_path, check_asan_string},
    env::{get_path_with_directory, update_path, LD_LIBRARY_PATH, PATH},
    expand::Expand,
    machine_id::MachineIdentity,
    process::run_cmd,
};
use anyhow::{Context, Error, Result};
#[cfg(target_os = "linux")]
use nix::sys::signal::{kill, Signal};
use stacktrace_parser::CrashLog;
#[cfg(any(target_os = "linux", target_family = "windows"))]
use stacktrace_parser::StackEntry;
use std::ffi::OsStr;
#[cfg(target_os = "linux")]
use std::process::Stdio;
use std::{collections::HashMap, path::Path, time::Duration};
use tempfile::tempdir;

const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5);

pub struct Tester<'a> {
    setup_dir: &'a Path,
    extra_setup_dir: Option<&'a Path>,
    exe_path: &'a Path,
    arguments: &'a [String],
    environ: &'a HashMap<String, String>,
    timeout: Duration,
    check_asan_log: bool,
    check_asan_stderr: bool,
    check_debugger: bool,
    check_retry_count: u64,
    add_setup_to_ld_library_path: bool,
    add_setup_to_path: bool,
    machine_identity: MachineIdentity,
}

#[derive(Debug)]
pub struct Crash {
    pub call_stack: Vec<String>,
    pub crash_type: String,
    pub crash_site: String,
}

#[derive(Debug)]
pub struct TestResult {
    pub crash_log: Option<CrashLog>,
    pub error: Option<Error>,
}

impl<'a> Tester<'a> {
    pub fn new(
        setup_dir: &'a Path,
        extra_setup_dir: Option<&'a Path>,
        exe_path: &'a Path,
        arguments: &'a [String],
        environ: &'a HashMap<String, String>,
        machine_identity: MachineIdentity,
    ) -> Self {
        Self {
            setup_dir,
            extra_setup_dir,
            exe_path,
            arguments,
            environ,
            timeout: DEFAULT_TIMEOUT,
            check_asan_log: false,
            check_asan_stderr: false,
            check_debugger: false,
            check_retry_count: 0,
            add_setup_to_ld_library_path: false,
            add_setup_to_path: false,
            machine_identity,
        }
    }

    pub fn timeout(self, value: u64) -> Self {
        Self {
            timeout: Duration::from_secs(value),
            ..self
        }
    }

    pub fn check_asan_log(self, value: bool) -> Self {
        Self {
            check_asan_log: value,
            ..self
        }
    }

    pub fn check_asan_stderr(self, value: bool) -> Self {
        Self {
            check_asan_stderr: value,
            ..self
        }
    }

    pub fn check_debugger(self, value: bool) -> Self {
        Self {
            check_debugger: value,
            ..self
        }
    }

    pub fn check_retry_count(self, value: u64) -> Self {
        Self {
            check_retry_count: value,
            ..self
        }
    }

    pub fn add_setup_to_ld_library_path(self, value: bool) -> Self {
        Self {
            add_setup_to_ld_library_path: value,
            ..self
        }
    }

    pub fn add_setup_to_path(self, value: bool) -> Self {
        Self {
            add_setup_to_path: value,
            ..self
        }
    }

    pub fn set_optional<T>(self, value: Option<T>, setter: impl FnOnce(Self, T) -> Self) -> Self {
        if let Some(value) = value {
            setter(self, value)
        } else {
            self
        }
    }

    #[cfg(target_family = "windows")]
    async fn test_input_debugger(
        &self,
        argv: &[impl AsRef<OsStr>],
        env: &HashMap<String, String>,
    ) -> Result<Option<CrashLog>> {
        const IGNORE_FIRST_CHANCE_EXCEPTIONS: bool = true;
        let report = input_tester::crash_detector::test_process(
            self.exe_path,
            argv,
            env,
            self.timeout,
            IGNORE_FIRST_CHANCE_EXCEPTIONS,
        )?;

        let crash = if let Some(exception) = report.exceptions.last() {
            let call_stack: Vec<_> = exception
                .stack_frames
                .iter()
                .map(|f| match &f {
                    debugger::stack::DebugStackFrame::CorruptFrame => StackEntry {
                        line: f.to_string(),
                        ..Default::default()
                    },
                    debugger::stack::DebugStackFrame::Frame {
                        module_name,
                        module_offset,
                        symbol,
                        file_info,
                    } => StackEntry {
                        line: f.to_string(),
                        function_name: symbol.as_ref().map(|x| x.symbol().to_owned()),
                        function_offset: symbol.as_ref().map(|x| x.displacement()),
                        address: None,
                        module_offset: Some(*module_offset),
                        module_path: Some(module_name.to_owned()),
                        source_file_line: file_info.as_ref().map(|x| x.line.into()),
                        source_file_name: file_info.as_ref().and_then(|x| {
                            x.file.rsplit_terminator('\\').next().map(|x| x.to_owned())
                        }),
                        source_file_path: file_info.as_ref().map(|x| x.file.to_string()),
                    },
                })
                .collect();

            let fault_type = exception.description.to_string();
            let sanitizer = fault_type.to_string();

            Some(CrashLog::new(
                None, None, sanitizer, fault_type, None, None, call_stack,
            )?)
        } else {
            None
        };

        Ok(crash)
    }

    #[cfg(target_os = "linux")]
    async fn test_input_debugger(
        &self,
        args: &[impl AsRef<OsStr>],
        env: &HashMap<String, String>,
    ) -> Result<Option<CrashLog>> {
        let mut cmd = std::process::Command::new(self.exe_path);
        cmd.args(args).stdin(Stdio::null());
        cmd.envs(env);

        let (sender, receiver) = tokio::sync::oneshot::channel();

        // Create two async tasks: one off-thread task for the blocking triage run,
        // and one task that will kill the triage target if we time out.
        let triage = tokio::task::spawn_blocking(move || {
            // Spawn a triage run, but stop it before execing.
            //
            // This calls a blocking `wait()` internally, on the forked child.
            let triage = crate::triage::TriageCommand::new(cmd)?;

            // Share the new child ID with main thread.
            let Ok(()) = sender.send(triage.pid()) else { bail!("unable to send PID") };

            // The target run is blocking, and may hang.
            triage.run()
        });

        // Save the new process ID of the spawned triage target, so we can try to kill
        // the (possibly hung) target out-of-band, if we time out.
        let target_pid = match receiver.await {
            Ok(pid) => pid,
            Err(e) => {
                if triage.is_finished() {
                    bail!("triage run failed: {:?}", triage.await.unwrap().err());
                } else {
                    bail!("unable to receive PID: {}", e);
                }
            }
        };

        let timeout = tokio::time::timeout(self.timeout, triage).await;
        if timeout.is_err() {
            // Yes. Try to kill the target process, if hung.
            kill(target_pid, Signal::SIGKILL)?;
            bail!("process timed out");
        }

        let report = timeout???;

        let Some(crash) = report.crashes.last() else {
            return Ok(None);
        };

        let crash_thread = crash
            .threads
            .get(&crash.tid.as_raw())
            .ok_or_else(|| anyhow!("no thread info for crash thread ID = {}", crash.tid))?;

        let call_stack: Vec<_> = crash_thread
            .callstack
            .iter()
            .enumerate()
            .map(|(idx, frame)| StackEntry {
                line: format!("#{idx} {frame}"),
                address: Some(frame.addr.0),
                function_name: frame.function.as_ref().map(|x| x.name.clone()),
                function_offset: frame.function.as_ref().map(|x| x.offset),
                module_path: frame.module.as_ref().map(|x| x.name.clone()),
                module_offset: frame.module.as_ref().map(|x| x.offset),
                source_file_name: None,
                source_file_line: None,
                source_file_path: None,
            })
            .collect();

        let crash_type = crash.signal.to_string();
        let sanitizer = crash_type.clone();
        let fault_type = crash_type;

        Ok(Some(CrashLog::new(
            None, None, sanitizer, fault_type, None, None, call_stack,
        )?))
    }

    pub async fn test_input(&self, input_file: impl AsRef<Path>) -> Result<TestResult> {
        let asan_dir = if self.check_asan_log {
            Some(tempdir()?)
        } else {
            None
        };

        let (argv, env) = {
            let expand = Expand::new(&self.machine_identity)
                .machine_id()
                .input_path(input_file)
                .target_exe(self.exe_path)
                .target_options(self.arguments)
                .setup_dir(self.setup_dir)
                .set_optional(self.extra_setup_dir, Expand::extra_setup_dir);

            let argv = expand.evaluate(self.arguments)?;
            let mut env: HashMap<String, String> = HashMap::new();
            for (k, v) in self.environ {
                env.insert(k.clone(), expand.evaluate_value(v)?);
            }

            let setup_dir = &self.setup_dir.to_path_buf();
            if self.add_setup_to_path {
                let new_path = match env.get(PATH) {
                    Some(v) => update_path(v.clone().into(), setup_dir)?,
                    None => get_path_with_directory(PATH, setup_dir)?,
                };

                env.insert(PATH.to_string(), new_path.to_string_lossy().to_string());
            }
            if self.add_setup_to_ld_library_path {
                let new_path = match env.get(LD_LIBRARY_PATH) {
                    Some(v) => update_path(v.clone().into(), setup_dir)?,
                    None => get_path_with_directory(LD_LIBRARY_PATH, setup_dir)?,
                };
                env.insert(
                    LD_LIBRARY_PATH.to_string(),
                    new_path.to_string_lossy().to_string(),
                );
            }

            if let Some(asan_dir) = &asan_dir {
                add_asan_log_env(&mut env, asan_dir.path());
            }

            (argv, env)
        };

        let mut error = None;
        let mut crash_log = None;

        let attempts = 1 + self.check_retry_count;
        for _ in 0..attempts {
            let result = if self.check_debugger {
                match self.test_input_debugger(&argv, &env).await {
                    Ok(crash) => (crash, None, None),
                    Err(error) => (None, Some(error), None),
                }
            } else {
                match run_cmd(self.exe_path, argv.clone(), &env, self.timeout).await {
                    Ok(output) => (None, None, Some(output)),
                    Err(error) => (None, Some(error), None),
                }
            };

            crash_log = result.0;
            error = result.1;
            let output = result.2;

            // order of operations for checking for crashes:
            // 1. if we ran under a debugger, and that caught a crash
            // 2. if we have an ASAN log in our temp directory
            // 3. if we have an ASAN log to STDERR
            if crash_log.is_none() {
                crash_log = if let Some(asan_dir) = &asan_dir {
                    check_asan_path(asan_dir.path())
                        .await
                        .context("parsing ASAN logs failed")?
                } else {
                    None
                };
            }

            if crash_log.is_none() && self.check_asan_stderr {
                if let Some(output) = output {
                    crash_log = check_asan_string(output.stderr)
                        .await
                        .context("parsing STDERR as ASAN failed")?;
                }
            }

            if crash_log.is_some() {
                break;
            }
        }

        Ok(TestResult { crash_log, error })
    }

    pub async fn is_crash(&self, input_file: impl AsRef<Path>) -> Result<bool> {
        let input_file = input_file.as_ref();
        let test_result = self
            .test_input(input_file)
            .await
            .with_context(|| format!("testing input failed: {}", input_file.display()))?;
        Ok(test_result.crash_log.is_some())
    }
}
