//! Utility functions for meka

use nix::libc::{STDERR_FILENO, STDOUT_FILENO};
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::env;
use std::fs::{File, read_link, read_to_string};
use std::io::{Read, Seek, SeekFrom};
use std::os::fd::IntoRawFd;
use std::os::fd::OwnedFd;
use std::os::unix::io::FromRawFd;
use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime};

use crate::process::ProcessOptions;

/// Default meka output file
pub const MEKA_FILE: &str = "meka_recording.json";

/// Read the contents of a memfd file descriptor and return it as a String
pub(crate) fn read_memfd(fd: OwnedFd) -> Option<String> {
    let mut file = unsafe { File::from_raw_fd(fd.into_raw_fd()) };
    file.seek(SeekFrom::Start(0)).ok()?;
    let mut buffer = Vec::new();
    file.read_to_end(&mut buffer).ok()?;
    std::mem::forget(file); // Prevent fd being closed
    Some(String::from_utf8_lossy(&buffer).to_string())
}

/// Find the full path of a binary in the given PATH variable
pub fn find_binary<P>(path_var: &[PathBuf], exe_name: P) -> Option<PathBuf>
where
    P: AsRef<Path>,
{
    let exe_name = exe_name.as_ref();
    path_var
        .iter()
        .filter_map(|dir| {
            let full_path = dir.join(exe_name);
            if full_path.is_file() {
                Some(full_path)
            } else {
                None
            }
        })
        .next()
}

/// Proc struct
pub struct Proc {
    /// pid of the proc
    _pid: i32,
}

impl Proc {
    /// create a new Proc instance
    pub fn new(pid: i32) -> Self {
        Self { _pid: pid }
    }

    /// Get the pid
    pub fn pid(&self) -> i32 {
        self._pid
    }

    /// Get the proc path
    fn get_proc_path(&self) -> PathBuf {
        PathBuf::from(format!("/proc/{}", self.pid()))
    }

    /// Get the parent PID of a given process ID
    pub fn parent_pid(&self) -> Option<i32> {
        // Read the status file for the given PID
        let path = self.get_proc_path().join("status");
        let content = read_to_string(path).ok()?;

        // Look for the line starting with "PPid:"
        for line in content.lines() {
            if let Some(rest) = line.strip_prefix("PPid:") {
                return rest.trim().parse().ok();
            }
        }
        None
    }

    /// get the environ of this proc
    pub fn environ(&self) -> Option<BTreeMap<String, String>> {
        let path = self.get_proc_path().join("environ");
        let content = read_to_string(&path).ok()?;
        let env = content
            .split('\0')
            .filter(|s| !s.is_empty()) // remove trailing empty from last \0
            .filter_map(|v| {
                let mut parts = v.splitn(2, '='); // split into 2 parts only
                match (parts.next(), parts.next()) {
                    (Some(k), Some(val)) => Some((k.to_string(), val.to_string())),
                    _ => None,
                }
            })
            .collect();

        Some(env)
    }

    /// get the cwd of this proc
    pub fn cwd(&self) -> Option<PathBuf> {
        let path = self.get_proc_path().join("cwd");
        let content = read_link(&path).ok()?;
        Some(content)
    }

    /// get the cmdline of this proc
    pub fn cmdline(&self) -> Option<Vec<String>> {
        let path = self.get_proc_path().join("cmdline");
        let content = read_to_string(&path).ok()?;
        let mut parts: Vec<String> = content.split("\0").map(|s| s.to_string()).collect();
        if parts.last().is_some_and(|s| s.is_empty()) {
            parts.pop();
        }
        Some(parts)
    }

    /// get the status of this proc
    pub fn status(&self) -> Option<Vec<String>> {
        let path = self.get_proc_path().join("status");
        let content = std::fs::read_to_string(&path).ok()?;
        let res = content.split("\n").map(|s| s.to_string()).collect();
        Some(res)
    }
}

/// Correctly split command line arguments, handling multiple spaces
pub fn correct_args(args: Vec<String>) -> Vec<String> {
    args.join(" ")
        .split(' ')
        .map(|s| s.to_string())
        .filter(|s| !s.is_empty())
        .collect()
}

/// Get the PATH variable
pub fn get_path_var() -> Vec<PathBuf> {
    let current_dir = env::current_dir().unwrap_or_else(|_| PathBuf::from("."));

    env::split_paths(&env::var("PATH").unwrap_or_default())
        .chain(std::iter::once(current_dir))
        .collect()
}

/// Parsed command
#[derive(Debug)]
pub struct ParsedCommand {
    /// Env variable
    pub envs: HashMap<String, String>,
    /// Executable path
    pub executable: PathBuf,
    /// Args
    pub args: Vec<String>,
}

/// Parse a command
/// # Errors
/// Return an error if args is invalid
pub fn parse_cmd<S: AsRef<str>>(
    all_args: &[S],
    path_var: &[PathBuf],
    wrap_shell: bool,
) -> Result<ParsedCommand, String> {
    let mut envs = HashMap::new();
    let args_string = all_args
        .iter()
        .map(|s| s.as_ref().to_string())
        .collect::<Vec<String>>();

    if all_args.is_empty() {
        return Err("all_args should not be empty".to_string());
    }
    if wrap_shell {
        let arg = "sh".to_string();
        let executable = find_binary(path_var, &arg)
            .ok_or_else(|| format!("Binary not found in path: '{}'", arg))?;

        // Collect remaining args including executable itself
        let args = args_string.join(" ");
        let args = ["sh".to_string(), "-c".to_string(), args].to_vec();

        return Ok(ParsedCommand {
            envs,
            executable,
            args,
        });
    }
    let mut args_iter = args_string.iter().cloned();
    while let Some(arg) = args_iter.next() {
        if let Some(eq_pos) = arg.find('=') {
            // Treat as an environment variable
            let (key, value) = arg.split_at(eq_pos);
            let value = &value[1..]; // skip '='
            envs.insert(key.to_string(), value.to_string());
        } else {
            // arg is exec
            let executable = find_binary(path_var, &arg)
                .ok_or_else(|| format!("Binary not found in path: '{}'", arg))?;

            // Collect remaining args including executable itself
            let args = std::iter::once(arg).chain(args_iter).collect();

            return Ok(ParsedCommand {
                envs,
                executable,
                args,
            });
        }
    }

    Err("Cannot parse command: no executable found".to_string())
}

/// Information about a single process
#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
#[serde(default)]
#[non_exhaustive]
pub struct ProgramData {
    /// Shell wrapped
    pub wrap_shell: bool,
    /// Full command line used to start the process
    pub cmdline: Vec<String>,
    /// Current working directory of the process
    pub cwd: PathBuf,
    /// environment variables
    pub environ: Option<BTreeMap<String, String>>,
    /// status
    pub status: Option<Vec<String>>,
    /// Subprograms started by this process
    pub subprograms: BTreeSet<i32>,
    /// Total time the process ran in milliseconds
    pub total_time: Duration,
    /// Parent process ID
    pub parent_pid: i32,
    /// Start time (UNIX timestamp in milliseconds)
    pub start_time: SystemTime,
    /// Process ID
    pub pid: i32,
    /// After exec cmd
    pub after_exec: Vec<Vec<String>>,
}

impl Default for ProgramData {
    fn default() -> Self {
        Self {
            wrap_shell: false,
            cmdline: Vec::new(),
            cwd: PathBuf::new(),
            environ: None,
            status: None,
            subprograms: BTreeSet::new(),
            total_time: Duration::new(0, 0),
            parent_pid: 0,
            start_time: SystemTime::now(),
            pid: 0,
            after_exec: Vec::new(),
        }
    }
}

impl ProgramData {
    /// Create new program data
    pub fn from_proc(proc: &Proc, options: ProcessOptions) -> Self {
        let ProcessOptions {
            capture_environ,
            capture_status,
            wrap_shell,
            ..
        } = options;
        let environ = capture_environ.then(|| proc.environ());
        let status = capture_status.then(|| proc.status());
        Self {
            pid: proc.pid(),
            wrap_shell,
            cmdline: proc.cmdline().unwrap_or_default(),
            cwd: proc.cwd().unwrap_or_default(),
            environ: environ.unwrap_or_default(),
            status: status.unwrap_or_default(),
            ..Default::default()
        }
    }

    /// Set the process as finished and calculate total time
    pub fn set_finished(&mut self) {
        if let Ok(elapsed) = self.start_time.elapsed() {
            self.total_time = elapsed;
        }
        // self.stdout = self.read_stdout().unwrap_or_default();
        // self.stderr = self.read_stderr().unwrap_or_default();
    }

    /// Read a fd
    /// # Errors
    /// Error is cannot access file
    pub fn read_fd(&self, fd_num: i32) -> Result<String, std::io::Error> {
        let path = format!("/proc/{}/fd/{}", self.pid, fd_num);
        let data = std::fs::read(path)?;
        Ok(String::from_utf8_lossy(&data).into_owned())
    }

    /// Read stderr
    /// # Errors
    /// Error is cannot access file
    pub fn read_stderr(&self) -> Result<String, std::io::Error> {
        self.read_fd(STDERR_FILENO)
    }

    /// Read stdout
    /// # Errors
    /// Error is cannot access file
    pub fn read_stdout(&self) -> Result<String, std::io::Error> {
        self.read_fd(STDOUT_FILENO)
    }
}

/// convert to relative pid
pub(crate) fn convert_to_relative(map: BTreeMap<i32, ProgramData>) -> BTreeMap<i32, ProgramData> {
    let mut new_map = BTreeMap::<i32, ProgramData>::new();
    let mut parent_pid = 0;
    let mut get_parent = false;
    for (key, value) in map {
        if !get_parent {
            parent_pid = key - 1;
            get_parent = true;
        }
        let mut new_value = value.clone();
        new_value.subprograms = value.subprograms.iter().map(|v| v - parent_pid).collect();
        new_value.pid = key - parent_pid;
        new_map.insert(new_value.pid, new_value);
    }
    new_map
}
