use std::{
    fs::File,
    io::{Read, Write},
    path::Path,
};

use crate::{
    command::{ArchiveSource, BundleSpec},
    error::{Error, Result},
    esc::{
        MaterializedInput, materialize_path, materialize_reader, prepare_payload,
        write_payload_body,
    },
    hash::shared_bundle_id,
    shell::{write_export_line, write_set_argv_line},
};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ArchiveKind {
    Tar,
    TarGzip,
    TarBzip2,
    TarXz,
    TarZstd,
}

impl ArchiveKind {
    fn shell_extract(self, source_var: &str, work_var: &str, index: usize) -> String {
        let label = format!("archive {}", index);
        match self {
            Self::Tar => format!(
                "tar -xf \"${{{}}}\" -C \"${{{}}}\" || bdlx_die 'failed to extract {}'\n",
                source_var, work_var, label
            ),
            Self::TarGzip => decompress_extract("gzip", source_var, work_var, index, &label),
            Self::TarBzip2 => decompress_extract("bzip2", source_var, work_var, index, &label),
            Self::TarXz => decompress_extract("xz", source_var, work_var, index, &label),
            Self::TarZstd => decompress_extract("zstd", source_var, work_var, index, &label),
        }
    }
}

#[derive(Debug)]
struct PreparedArchive {
    materialized: MaterializedInput,
    payload: crate::esc::PreparedPayload,
    kind: ArchiveKind,
}

/// Writes a self-extracting POSIX `sh` bundle script.
///
/// Archive sources are read while this function runs and embedded into the
/// generated script. The generated script extracts archives in the order listed
/// in `spec.archives`, exports `spec.env`, then runs `spec.command`.
pub fn write_bundle<W: Write, R: Read + ?Sized>(
    spec: BundleSpec<'_, R>,
    writer: &mut W,
) -> Result<()> {
    let BundleSpec {
        cleanup,
        shared,
        archives,
        env,
        command,
    } = spec;

    let mut archives = materialize_archive_sources(archives)?;
    let archive_hashes = archives
        .iter()
        .map(|archive| archive.materialized.sha256.clone())
        .collect::<Vec<_>>();
    let render_spec: BundleSpec<'_, R> = BundleSpec {
        cleanup,
        shared,
        archives: Vec::new(),
        env,
        command,
    };

    write_prepared_bundle(&render_spec, &mut archives, &archive_hashes, writer)
}

fn write_prepared_bundle<W: Write, R: Read + ?Sized>(
    spec: &BundleSpec<'_, R>,
    archives: &mut [PreparedArchive],
    archive_hashes: &[String],
    writer: &mut W,
) -> Result<()> {
    let shared_id = if spec.shared {
        Some(shared_bundle_id(spec, archive_hashes))
    } else {
        None
    };

    writer
        .write_all(bundle_header().as_bytes())
        .map_err(|source| Error::io("failed to write bundle script header", source))?;
    write_runtime(writer, spec, shared_id.as_deref())?;

    for (index, archive) in archives.iter_mut().enumerate() {
        write_payload_function(writer, index + 1, &mut archive.payload)?;
        write_archive_function(writer, index + 1, archive.kind)?;
    }

    write_main(writer, spec, archives.len())?;
    Ok(())
}

fn materialize_archive_sources<R: Read + ?Sized>(
    sources: Vec<ArchiveSource<'_, R>>,
) -> Result<Vec<PreparedArchive>> {
    let mut archives = Vec::new();

    for source in sources {
        let mut materialized = match source {
            ArchiveSource::Path(path) => materialize_path(&path)?,
            ArchiveSource::Reader { label, reader } => materialize_reader(reader, label)?,
        };
        let kind = sniff_archive_kind(materialized.temp.path(), &materialized.label)?;
        let payload = prepare_payload(&mut materialized)?;
        archives.push(PreparedArchive {
            materialized,
            payload,
            kind,
        });
    }

    Ok(archives)
}

fn bundle_header() -> String {
    format!(
        "#!/bin/sh\nset -eu\n\n# Generated by bdlx {}\n# Runtime flags: --extract --no-rm --no-run [--]\n\n",
        crate::version()
    )
}

fn write_runtime<W: Write, R: Read + ?Sized>(
    writer: &mut W,
    spec: &BundleSpec<'_, R>,
    shared_id: Option<&str>,
) -> Result<()> {
    writer
        .write_all(
            b"bdlx_die() {\n  printf '%s\\n' \"$*\" >&2\n  exit 1\n}\n\nbdlx_cleanup_policy='",
        )
        .map_err(|source| Error::io("failed to write runtime shell", source))?;
    writer
        .write_all(spec.cleanup.as_shell_value().as_bytes())
        .map_err(|source| Error::io("failed to write runtime shell", source))?;
    writer
        .write_all(
            b"'\n\
bdlx_keep=0\n\
bdlx_no_run=0\n\
bdlx_print_path=0\n\
bdlx_work=''\n\
\n\
bdlx_cleanup() {\n\
  bdlx_status=$?\n\
  trap - EXIT HUP INT TERM\n\
  if [ -n \"${bdlx_work}\" ] && [ \"$bdlx_keep\" -eq 0 ]; then\n\
    case \"$bdlx_cleanup_policy\" in\n\
      success) [ \"$bdlx_status\" -eq 0 ] && rm -rf \"$bdlx_work\" ;;\n\
      failure) [ \"$bdlx_status\" -ne 0 ] && rm -rf \"$bdlx_work\" ;;\n\
      always) rm -rf \"$bdlx_work\" ;;\n\
      never) : ;;\n\
      *) bdlx_die 'invalid cleanup policy' ;;\n\
    esac\n\
  fi\n\
  exit \"$bdlx_status\"\n\
}\n\
\n\
trap 'exit 129' HUP\n\
trap 'exit 130' INT\n\
trap 'exit 143' TERM\n\
trap 'bdlx_cleanup' EXIT\n\
\n\
bdlx_make_fresh_dir() {\n\
  bdlx_tmp_root=${TMPDIR:-/tmp}\n\
  bdlx_i=0\n\
  while [ \"$bdlx_i\" -lt 1000 ]; do\n\
    bdlx_candidate=$bdlx_tmp_root/bdlx.$$.$bdlx_i\n\
    if (umask 077 && mkdir \"$bdlx_candidate\") 2>/dev/null; then\n\
      printf '%s\\n' \"$bdlx_candidate\"\n\
      return 0\n\
    fi\n\
    bdlx_i=$((bdlx_i + 1))\n\
  done\n\
  return 1\n\
}\n\
\n",
        )
        .map_err(|source| Error::io("failed to write runtime shell", source))?;

    if let Some(shared_id) = shared_id {
        write!(
            writer,
            "bdlx_shared_id='{}'\n\
bdlx_prepare_shared_dir() {{\n\
  bdlx_tmp_root=${{TMPDIR:-/tmp}}\n\
  bdlx_final=$bdlx_tmp_root/bdlx.shared.$bdlx_shared_id\n\
  if [ -f \"$bdlx_final/.bdlx-ready\" ]; then\n\
    printf '%s\\n' \"$bdlx_final\"\n\
    return 0\n\
  fi\n\
  if [ -e \"$bdlx_final\" ]; then\n\
    bdlx_die \"shared directory exists but is not ready: $bdlx_final\"\n\
  fi\n\
  bdlx_i=0\n\
  while [ \"$bdlx_i\" -lt 1000 ]; do\n\
    bdlx_stage=$bdlx_tmp_root/bdlx.shared.$bdlx_shared_id.stage.$$.$bdlx_i\n\
    if (umask 077 && mkdir \"$bdlx_stage\") 2>/dev/null; then\n\
      printf '%s\\n' \"$bdlx_stage\"\n\
      return 0\n\
    fi\n\
    bdlx_i=$((bdlx_i + 1))\n\
  done\n\
  return 1\n\
}}\n\n",
            shared_id
        )
        .map_err(|source| Error::io("failed to write runtime shell", source))?;
    }

    Ok(())
}

fn write_payload_function<W: Write>(
    writer: &mut W,
    index: usize,
    payload: &mut crate::esc::PreparedPayload,
) -> Result<()> {
    write!(writer, "bdlx_payload_{}() {{\n  ", index)
        .map_err(|source| Error::io("failed to write payload function", source))?;
    let mut fragment = String::new();
    payload
        .decoder
        .write_decoder_fragment(payload.original_len, &mut fragment)
        .map_err(|_| Error::Internal("failed to render shell decoder fragment".to_owned()))?;
    writer
        .write_all(fragment.as_bytes())
        .map_err(|source| Error::io("failed to write payload function", source))?;
    writeln!(writer, " <<'{}'", payload.delimiter)
        .map_err(|source| Error::io("failed to write payload function", source))?;
    write_payload_body(payload, writer)?;
    write!(writer, "\n{}\n}}\n\n", payload.delimiter)
        .map_err(|source| Error::io("failed to finish payload function", source))?;
    Ok(())
}

fn write_archive_function<W: Write>(writer: &mut W, index: usize, kind: ArchiveKind) -> Result<()> {
    let archive_var = format!("bdlx_archive_{}", index);
    write!(
        writer,
        "bdlx_unpack_{}() {{\n  {}=\"$bdlx_work/.bdlx-archive-{}\"\n  bdlx_payload_{} > \"${{{}}}\" || bdlx_die 'failed to materialize archive {}'\n  ",
        index,
        archive_var,
        index,
        index,
        archive_var,
        index
    )
    .map_err(|source| Error::io("failed to write archive function", source))?;
    writer
        .write_all(
            kind.shell_extract(&archive_var, "bdlx_work", index)
                .as_bytes(),
        )
        .map_err(|source| Error::io("failed to write archive function", source))?;
    writeln!(writer, "  rm -f \"${{{}}}\"\n}}\n", archive_var)
        .map_err(|source| Error::io("failed to write archive function", source))?;
    Ok(())
}

fn write_main<W: Write, R: Read + ?Sized>(
    writer: &mut W,
    spec: &BundleSpec<'_, R>,
    archive_count: usize,
) -> Result<()> {
    writer
        .write_all(
            b"main() {\n\
  while [ \"$#\" -gt 0 ]; do\n\
    case \"$1\" in\n\
      --extract)\n\
        bdlx_no_run=1\n\
        bdlx_keep=1\n\
        bdlx_print_path=1\n\
        shift\n\
        ;;\n\
      --no-rm)\n\
        bdlx_keep=1\n\
        shift\n\
        ;;\n\
      --no-run)\n\
        bdlx_no_run=1\n\
        bdlx_keep=1\n\
        bdlx_print_path=1\n\
        shift\n\
        ;;\n\
      --)\n\
        shift\n\
        break\n\
        ;;\n\
      *)\n\
        break\n\
        ;;\n\
    esac\n\
  done\n",
        )
        .map_err(|source| Error::io("failed to write bundle main", source))?;
    if spec.shared {
        writer
            .write_all(b"  bdlx_work=$(bdlx_prepare_shared_dir) || bdlx_die 'failed to prepare shared directory'\n")
            .map_err(|source| Error::io("failed to write bundle main", source))?;
        writer
            .write_all(b"  if [ -f \"$bdlx_work/.bdlx-ready\" ]; then\n")
            .map_err(|source| Error::io("failed to write bundle main", source))?;
        writer
            .write_all(b"    bdlx_keep=1\n")
            .map_err(|source| Error::io("failed to write bundle main", source))?;
        writer
            .write_all(b"  else\n")
            .map_err(|source| Error::io("failed to write bundle main", source))?;
        for index in 0..archive_count {
            writeln!(writer, "    bdlx_unpack_{}", index + 1)
                .map_err(|source| Error::io("failed to write bundle main", source))?;
        }
        writer
            .write_all(
                b"    : > \"$bdlx_work/.bdlx-ready\" || bdlx_die 'failed to mark shared directory ready'\n\
    bdlx_final=${TMPDIR:-/tmp}/bdlx.shared.$bdlx_shared_id\n\
    if [ \"$bdlx_work\" != \"$bdlx_final\" ]; then\n\
      if mv \"$bdlx_work\" \"$bdlx_final\" 2>/dev/null; then\n\
        bdlx_work=$bdlx_final\n\
      else\n\
        rm -rf \"$bdlx_work\"\n\
        [ -f \"$bdlx_final/.bdlx-ready\" ] || bdlx_die 'failed to publish shared directory'\n\
        bdlx_work=$bdlx_final\n\
      fi\n\
    fi\n",
            )
            .map_err(|source| Error::io("failed to write bundle main", source))?;
        writer
            .write_all(b"    bdlx_keep=1\n")
            .map_err(|source| Error::io("failed to write bundle main", source))?;
        writer
            .write_all(b"  fi\n")
            .map_err(|source| Error::io("failed to write bundle main", source))?;
    } else {
        writer
            .write_all(b"  bdlx_work=$(bdlx_make_fresh_dir) || bdlx_die 'failed to create temp directory'\n")
            .map_err(|source| Error::io("failed to write bundle main", source))?;
        for index in 0..archive_count {
            writeln!(writer, "  bdlx_unpack_{}", index + 1)
                .map_err(|source| Error::io("failed to write bundle main", source))?;
        }
    }

    for env in &spec.env {
        writer
            .write_all(b"  ")
            .map_err(|source| Error::io("failed to write bundle main", source))?;
        write_export_line(writer, &env.name, &env.value)
            .map_err(|source| Error::io("failed to write bundle main", source))?;
    }
    writer
        .write_all(
            b"  cd \"$bdlx_work\" || bdlx_die 'failed to enter work directory'\n\
  if [ \"$bdlx_print_path\" -eq 1 ]; then\n\
    printf '%s\\n' \"$bdlx_work\"\n\
  fi\n\
  if [ \"$bdlx_no_run\" -eq 1 ]; then\n\
    exit 0\n\
  fi\n  ",
        )
        .map_err(|source| Error::io("failed to write bundle main", source))?;
    write_set_argv_line(writer, &spec.command.program, &spec.command.args)
        .map_err(|source| Error::io("failed to write bundle main", source))?;
    writer
        .write_all(b"  \"$@\"\n}\n\nmain \"$@\"\n")
        .map_err(|source| Error::io("failed to write bundle main", source))?;
    Ok(())
}

fn sniff_archive_kind(path: &Path, label: &str) -> Result<ArchiveKind> {
    let mut file = File::open(path)
        .map_err(|source| Error::io_path("failed to open archive", path, source))?;
    let mut header = [0u8; 512];
    let read = file
        .read(&mut header)
        .map_err(|source| Error::io(format!("failed to read {}", label), source))?;
    let header = &header[..read];

    if header.starts_with(&[0x1f, 0x8b]) {
        return Ok(ArchiveKind::TarGzip);
    }
    if header.starts_with(&[0x42, 0x5a, 0x68]) {
        return Ok(ArchiveKind::TarBzip2);
    }
    if header.starts_with(&[0xfd, b'7', b'z', b'X', b'Z', 0x00]) {
        return Ok(ArchiveKind::TarXz);
    }
    if header.starts_with(&[0x28, 0xb5, 0x2f, 0xfd]) {
        return Ok(ArchiveKind::TarZstd);
    }
    if header.len() >= 262 && &header[257..262] == b"ustar" {
        return Ok(ArchiveKind::Tar);
    }

    Err(Error::UnsupportedArchiveFormat {
        path: label.to_owned(),
    })
}

fn decompress_extract(
    tool: &str,
    source_var: &str,
    work_var: &str,
    _index: usize,
    label: &str,
) -> String {
    format!(
        "bdlx_tmp_tar=\"${{{}}}.tar\"\n  {} -dc \"${{{}}}\" > \"$bdlx_tmp_tar\" || bdlx_die 'failed to decompress {}'\n  tar -xf \"$bdlx_tmp_tar\" -C \"${{{}}}\" || bdlx_die 'failed to extract {}'\n  rm -f \"$bdlx_tmp_tar\"\n",
        source_var, tool, source_var, label, work_var, label
    )
}
