use std::path::{Path, PathBuf};
use std::{env, fs, io};

use anyhow::{anyhow, Context};
use clap::Parser;

#[derive(Parser)]
#[clap(author, version, about)]
struct Args {
    /// The environment.d directory to load
    ///
    /// If this is not specified on the command line, $XDG_CONFIG_HOME/environment.d or ~/.config/environment.d will be used.
    #[clap(verbatim_doc_comment)]
    pub directory: Option<PathBuf>,

    /// The format to write the output in
    ///
    /// Supported formats:
    ///   - normal: VARIABLE=value (with shell escapes)
    ///   - export: export VARIABLE=value (with shell escapes)
    ///   - json: {"VARIABLE":"value"} (only if `json` feature enabled during compilation)
    ///   - hyprland: env=VARIABLE,value
    ///   - kdl: VARIABLE "value"
    #[clap(
        short,
        long,
        value_enum,
        default_value = "normal",
        verbatim_doc_comment
    )]
    pub format: Format,

    /// Deprecated alias for --format=export
    #[deprecated]
    #[clap(short, long, conflicts_with = "format")]
    pub export: bool,

    /// Inject output into a file instead of writing it to stdout
    ///
    /// Where the environment variables will be injected is determined using delimiter comments, as described below. If the delimiter comments are not present, the command will fail.
    ///
    /// The delimiter comments should look something like the following:
    ///
    ///     # degenverate: begin injected vars
    ///     SOME_VARIABLE=value
    ///     # degenverate: end injected vars
    ///
    /// Since comment syntax can vary across configuration file formats, any type of comment may be used as long as the line contains 'degenverate: begin/end injected vars'.
    #[cfg(feature = "inject")]
    #[clap(short = 'j', long, verbatim_doc_comment)]
    pub inject: Option<PathBuf>,

    /// The environment variables to exclude from the output
    #[clap(short = 'x', long)]
    pub exclude: Vec<String>,
}

#[derive(Debug, Clone, Copy, clap::ValueEnum)]
enum Format {
    Normal,
    Export,
    #[cfg(feature = "json")]
    Json,
    Hyprland,
    Kdl,
}

#[cfg(feature = "preserve_order")]
type EnvironmentMap = indexmap::IndexMap<String, String>;
#[cfg(not(feature = "preserve_order"))]
type EnvironmentMap = std::collections::BTreeMap<String, String>;

fn read_environment_d(directory: &Path) -> anyhow::Result<EnvironmentMap> {
    let mut entries = fs::read_dir(directory)?
        .map(|res| res.map(|e| e.path()))
        .collect::<Result<Vec<_>, io::Error>>()?;

    // The order in which `read_dir` returns entries is not guaranteed. If reproducible
    // ordering is required the entries should be explicitly sorted.
    entries.sort();

    let mut variables = EnvironmentMap::new();

    for path in entries {
        for item in dotenvy::from_path_iter(&path)? {
            let (key, val) = item
                .with_context(|| format!("failed to load environment file {}", path.display()))?;

            if !is_valid_env_var(&key) {
                eprintln!("degenverate: '{key}' is not a valid environment variable name");
            }

            // This allows future variables to reference this one.
            unsafe {
                env::set_var(&key, &val);
            }

            variables.insert(key, val);
        }
    }

    Ok(variables)
}

fn write_variables(
    mut writer: impl io::Write,
    format: Format,
    variables: &EnvironmentMap,
) -> io::Result<()> {
    match format {
        Format::Normal => {
            for (key, val) in variables {
                let escaped_val = shell_escape::escape(val.into());
                writeln!(writer, "{key}={escaped_val}")?;
            }
        }
        Format::Export => {
            for (key, val) in variables {
                let escaped_val = shell_escape::escape(val.into());
                writeln!(writer, "export {key}={escaped_val}")?;
            }
        }
        #[cfg(feature = "json")]
        Format::Json => {
            serde_json::to_writer(writer, variables)?;
        }
        Format::Hyprland => {
            for (key, val) in variables {
                writeln!(writer, "env={key},{val}")?;
            }
        }
        Format::Kdl => {
            for (key, val) in variables {
                let escaped_val = val.replace("\\", "\\\\").replace("\"", "\\\"");
                writeln!(writer, "{key} \"{escaped_val}\"")?;
            }
        }
    };
    Ok(())
}

#[cfg(feature = "inject")]
fn inject_variables_to_file(
    file_path: &Path,
    format: Format,
    variables: &EnvironmentMap,
) -> anyhow::Result<()> {
    use std::io::{Read, Seek, Write};

    let mut file = fs::OpenOptions::new()
        .read(true)
        .write(true)
        .open(file_path)
        .with_context(|| format!("failed to open file: {}", file_path.display()))?;

    let mut file_content = Vec::new();
    file.read_to_end(&mut file_content)?;

    let begin_index: usize =
        memchr::memmem::find(&file_content, b"degenverate: begin injected vars")
            .and_then(|i| memchr::memchr(b'\n', &file_content[i..]).map(|j| j + i + 1))
            .context("file does not contain injection delimiter comments")?;
    let end_index: usize = memchr::memmem::find(
        &file_content[begin_index..],
        b"degenverate: end injected vars",
    )
    .map(|i| i + begin_index)
    .and_then(|i| {
        memchr::memrchr(b'\n', &file_content[(begin_index - 1)..i]).map(|i| i + begin_index - 1)
    })
    .context("file does not contain injection delimiter comments")?;

    file.set_len(0)?;
    file.seek(io::SeekFrom::Start(0))?;

    file.write_all(&file_content[..begin_index])?;
    write_variables(&mut file, format, variables)?;
    file.write_all(&file_content[end_index..])?;

    eprintln!(
        "successfully injected {} environment variables into {}",
        variables.len(),
        file_path.display()
    );

    Ok(())
}

fn main() -> anyhow::Result<()> {
    let args = Args::parse();

    let directory = args
        .directory
        .or_else(default_environment_d)
        .ok_or_else(|| anyhow!("could not determine path to environment.d directory"))?;

    #[allow(deprecated)]
    let format = if args.export {
        Format::Export
    } else {
        args.format
    };

    let mut variables = read_environment_d(&directory)?;
    for var in args.exclude {
        variables.remove(&var);
    }

    #[cfg(feature = "inject")]
    if let Some(file_path) = args.inject {
        return inject_variables_to_file(&file_path, args.format, &variables);
    }

    write_variables(io::stdout(), format, &variables)?;

    Ok(())
}

fn default_environment_d() -> Option<PathBuf> {
    let mut path = dirs::config_dir()?;
    path.push("environment.d");
    Some(path)
}

fn is_valid_env_var(key: &str) -> bool {
    let mut chars_iter = key.chars();
    chars_iter
        .next()
        .map_or(false, |ch| ch.is_ascii_alphabetic() || ch == '_')
        && chars_iter.all(|ch| ch.is_ascii_alphanumeric() || ch == '_')
}
