use cargo_lambda_metadata::{cargo::binary_targets, fs::rename};
use cargo_zigbuild::Build as ZigBuild;
use clap::{Args, ValueHint};
use miette::{IntoDiagnostic, Result, WrapErr};
use object::{read::File as ObjectFile, Architecture, Object};
use sha2::{Digest, Sha256};
use std::{
    env,
    fs::{create_dir_all, read, File},
    io::Write,
    path::{Path, PathBuf},
    str::FromStr,
};
use strum_macros::EnumString;
use target_arch::TargetArch;
use tracing::{debug, warn};
use zip::{write::FileOptions, ZipWriter};

mod toolchain;
mod zig;

#[derive(Args, Clone, Debug)]
#[clap(name = "build")]
pub struct Build {
    /// The format to produce the compile Lambda into, acceptable values are [Binary, Zip]
    #[clap(long, default_value_t = OutputFormat::Binary)]
    output_format: OutputFormat,

    /// Directory where the final lambda binaries will be located
    #[clap(short, long, value_hint = ValueHint::DirPath)]
    lambda_dir: Option<PathBuf>,

    /// Shortcut for --target aarch64-unknown-linux-gnu
    #[clap(long)]
    arm64: bool,

    /// Whether the code that you're building is a Lambda Extension
    #[clap(long)]
    extension: bool,

    /// Put a bootstrap file in the root of the lambda directory.
    /// Use the name of the compiled binary to choose which file to move.
    #[clap(long)]
    flatten: Option<String>,

    /// Disable Zig as the linker.
    /// This option is only allowed when you're building from Linux.
    /// This option can help you build native libraries when
    /// the Zig linker is failing because it cannot find them correctly.
    ///
    /// If you use this option, the GLIBC version is set to 2.26
    /// automatically to match the version supported by AWS Lambda.
    /// https://aws.amazon.com/amazon-linux-2/faqs/
    ///
    /// Example:
    /// cargo lambda build --release --disable-zig-linker
    #[clap(long)]
    disable_zig_linker: bool,

    #[clap(flatten)]
    build: ZigBuild,
}

pub use cargo_zigbuild::Zig;

mod target_arch;

#[derive(Clone, Debug, strum_macros::Display, EnumString)]
#[strum(ascii_case_insensitive)]
enum OutputFormat {
    Binary,
    Zip,
}

impl Build {
    #[tracing::instrument(skip(self), target = "cargo_lambda")]
    pub async fn run(&mut self) -> Result<()> {
        tracing::trace!(options = ?self, "building project");

        let rustc_meta = rustc_version::version_meta().into_diagnostic()?;
        let host_target = &rustc_meta.host;
        let release_channel = &rustc_meta.channel;
        let compatible_host_linker = TargetArch::compatible_host_linker(host_target);

        if self.arm64 && !self.build.target.is_empty() {
            return Err(miette::miette!(
                "invalid options: --arm and --target cannot be specified at the same time"
            ));
        }

        let mut target_arch = if self.arm64 {
            TargetArch::arm64()
        } else {
            let build_target = self.build.target.get(0);
            match build_target {
                Some(target) => TargetArch::from_str(target)?,
                // No explicit target, but build host same as target host
                None if compatible_host_linker => {
                    // Set the target explicitly, so it's easier to find the binaries later
                    TargetArch::from_str(host_target)?
                }
                // No explicit target, and build host not compatible with Lambda hosts
                None => TargetArch::x86_64(),
            }
        };

        if self.disable_zig_linker {
            if compatible_host_linker {
                target_arch.set_al2_glibc_version();
                self.build.disable_zig_linker = self.disable_zig_linker;
            } else {
                return Err(miette::miette!(
                    "invalid options: --disable-zig-linker is only allowed on Linux"
                ));
            }
        }

        self.build.target = vec![target_arch.to_string()];
        let rustc_target_without_glibc_version = target_arch.rustc_target_without_glibc_version();

        #[cfg(windows)]
        self.force_windows_release_profile();

        let profile = match self.build.profile.as_deref() {
            Some("dev" | "test") => "debug",
            Some("release" | "bench") => "release",
            Some(profile) => profile,
            None if self.build.release => "release",
            None => "debug",
        };

        // confirm that target component is included in host toolchain, or add
        // it with `rustup` otherwise.
        toolchain::check_target_component_with_rustc_meta(
            &rustc_target_without_glibc_version,
            host_target,
            release_channel,
        )
        .await?;

        let manifest_path = self
            .build
            .manifest_path
            .as_deref()
            .unwrap_or_else(|| Path::new("Cargo.toml"));
        let binaries = binary_targets(manifest_path)?;
        debug!(binaries = ?binaries, "found new target binaries to build");

        if !self.build.bin.is_empty() {
            for name in &self.build.bin {
                if !binaries.contains(name) {
                    return Err(miette::miette!(
                        "binary target is missing from this project: {}",
                        name
                    ));
                }
            }
        }

        if !self.build.disable_zig_linker {
            zig::check_installation().await?;
        }

        let mut cmd = self
            .build
            .build_command()
            .map_err(|e| miette::miette!("{}", e))?;

        if self.build.release {
            let mut rust_flags = env::var("RUSTFLAGS").unwrap_or_default();
            if !rust_flags.contains("-C strip=") {
                if !rust_flags.is_empty() {
                    rust_flags += " ";
                }
                rust_flags += "-C strip=symbols";
            }
            if !rust_flags.contains("-C target-cpu=") {
                if !rust_flags.is_empty() {
                    rust_flags += " ";
                }
                let target_cpu = target_arch.target_cpu();
                rust_flags += "-C target-cpu=";
                rust_flags += target_cpu.as_str();
            }

            debug!(rust_flags = ?rust_flags, "release RUSTFLAGS");
            cmd.env("RUSTFLAGS", rust_flags);
        }

        let mut child = cmd
            .spawn()
            .into_diagnostic()
            .wrap_err("Failed to run cargo build")?;
        let status = child
            .wait()
            .into_diagnostic()
            .wrap_err("Failed to wait on cargo build process")?;
        if !status.success() {
            std::process::exit(status.code().unwrap_or(1));
        }

        let target_dir = Path::new("target");
        let lambda_dir = if let Some(dir) = &self.lambda_dir {
            dir.clone()
        } else {
            target_dir.join("lambda")
        };

        let base = target_dir
            .join(rustc_target_without_glibc_version)
            .join(profile);

        let mut found_binaries = false;
        for name in &binaries {
            let binary = base.join(name);
            debug!(binary = ?binary, exists = binary.exists(), "checking function binary");

            if binary.exists() {
                found_binaries = true;

                let bootstrap_dir = if self.extension {
                    lambda_dir.join("extensions")
                } else {
                    match self.flatten {
                        Some(ref n) if n == name => lambda_dir.clone(),
                        _ => lambda_dir.join(name),
                    }
                };
                create_dir_all(&bootstrap_dir).into_diagnostic()?;

                let bin_name = if self.extension {
                    name.as_str()
                } else {
                    "bootstrap"
                };

                match self.output_format {
                    OutputFormat::Binary => {
                        rename(binary, bootstrap_dir.join(bin_name)).into_diagnostic()?;
                    }
                    OutputFormat::Zip => {
                        let parent = if self.extension {
                            Some("extensions")
                        } else {
                            None
                        };
                        zip_binary(bin_name, binary, bootstrap_dir, parent)?;
                    }
                }
            }
        }
        if !found_binaries {
            warn!("no binaries found in target after build, try using the --bin or --package options to build specific binaries");
        }

        Ok(())
    }

    #[cfg(windows)]
    fn force_windows_release_profile(&mut self) {
        if !self.build.release {
            tracing::info!("Changing profile to release mode. Cargo-lambda doesn't support building on debug mode on Windows");
            self.build.release = true;
            self.build.profile = Some("release".to_string());
        }
    }
}

pub struct BinaryArchive {
    pub architecture: String,
    pub sha256: String,
    pub path: PathBuf,
}

/// Search for the bootstrap file for a function inside the target directory.
/// If the binary file exists, it creates the zip archive and extracts its architectury by reading the binary.
pub fn find_binary_archive<P: AsRef<Path>>(
    name: &str,
    base_dir: &Option<P>,
    is_extension: bool,
) -> Result<BinaryArchive> {
    let target_dir = Path::new("target");
    let (dir_name, binary_name, parent) = if is_extension {
        ("extensions", name, Some("extensions"))
    } else {
        (name, "bootstrap", None)
    };

    let bootstrap_dir = if let Some(dir) = base_dir {
        dir.as_ref().join(dir_name)
    } else {
        target_dir.join("lambda").join(dir_name)
    };

    let binary_path = bootstrap_dir.join(binary_name);
    if !binary_path.exists() {
        let build_cmd = if is_extension {
            "build --extension"
        } else {
            "build"
        };
        return Err(miette::miette!(
            "binary file for {} not found, use `cargo lambda {}` to create it",
            name,
            build_cmd
        ));
    }

    zip_binary(binary_name, binary_path, bootstrap_dir, parent)
}

/// Create a zip file from a function binary.
/// The binary inside the zip file is called `bootstrap` for function binaries.
/// The binary inside the zip file is called by its name, and put inside the `extensions`
/// directory, for extension binaries.
pub fn zip_binary<BP: AsRef<Path>, DD: AsRef<Path>>(
    name: &str,
    binary_path: BP,
    destination_directory: DD,
    parent: Option<&str>,
) -> Result<BinaryArchive> {
    let path = binary_path.as_ref();
    let dir = destination_directory.as_ref();
    let zipped = dir.join(format!("{}.zip", name));

    let zipped_binary = File::create(&zipped).into_diagnostic()?;
    let binary_data = read(path).into_diagnostic()?;
    let binary_perm = binary_permissions(path)?;
    let binary_data = &*binary_data;
    let object = ObjectFile::parse(binary_data).into_diagnostic()?;

    let arch = match object.architecture() {
        Architecture::Aarch64 => "arm64",
        Architecture::X86_64 => "x86_64",
        other => return Err(miette::miette!("invalid binary architecture: {:?}", other)),
    };

    let mut hasher = Sha256::new();
    hasher.update(binary_data);
    let sha256 = format!("{:X}", hasher.finalize());

    let mut zip = ZipWriter::new(zipped_binary);
    let file_name = if let Some(parent) = parent {
        zip.add_directory(parent, FileOptions::default())
            .into_diagnostic()?;
        Path::new(parent).join(name)
    } else {
        PathBuf::from(name)
    };

    zip.start_file(
        file_name.to_str().expect("failed to convert file path"),
        FileOptions::default().unix_permissions(binary_perm),
    )
    .into_diagnostic()?;
    zip.write_all(binary_data).into_diagnostic()?;
    zip.finish().into_diagnostic()?;

    Ok(BinaryArchive {
        architecture: arch.into(),
        path: zipped,
        sha256,
    })
}

#[cfg(unix)]
fn binary_permissions(path: &Path) -> Result<u32> {
    use std::os::unix::prelude::PermissionsExt;
    let meta = std::fs::metadata(path).into_diagnostic()?;
    Ok(meta.permissions().mode())
}

#[cfg(not(unix))]
fn binary_permissions(_path: &Path) -> Result<u32> {
    Ok(0o755)
}