// SPDX-FileCopyrightText: 2023-2026 Andrew Gunnerson
// SPDX-License-Identifier: GPL-3.0-only

use std::{
    io::{self, Seek, SeekFrom, Write},
    ops::Range,
    sync::atomic::AtomicBool,
};

use memchr::memmem;
use rawzip::ZipArchive;
use rayon::iter::{IntoParallelIterator, ParallelIterator};
use thiserror::Error;
use tracing::{Span, debug, debug_span, trace};
use x509_cert::Certificate;

use crate::{
    crypto::{SigningMethod, SigningPrivateKey},
    format::{
        avb::{self, AppendedDescriptorMut, Footer},
        ota,
        zip::{ZipFileHeaderRecordExt, ZipSliceEntriesSafeExt},
    },
    patch::otacert,
    stream::{self, ReadFixedSizeExt, ReadWriteAt, UserPosFile},
    util,
};

#[derive(Debug, Error)]
pub enum Error {
    #[error("Old otacerts.zip not found in image")]
    OldZipNotFound,
    #[error("Image has no vbmeta footer")]
    NoFooter,
    #[error("No hash tree descriptor found in vbmeta header")]
    NoHashTreeDescriptor,
    #[error("{0:?} overflowed integer bounds during calculations")]
    IntOverflow(&'static str),
    #[error("Failed to update AVB header")]
    AvbUpdate(#[source] avb::Error),
    #[error("Failed to generate replacement otacerts zip")]
    OtaCertZip(#[source] otacert::Error),
    #[error("Failed to read image data")]
    ReadData(#[source] io::Error),
    #[error("Failed to write image data")]
    WriteData(#[source] io::Error),
}

type Result<T> = std::result::Result<T, Error>;

/// Find the bounds of a non-zip64 zip starting from the EOCD magic offset.
fn find_zip_bounds(data: &[u8], eocd_offset: usize) -> Option<Range<usize>> {
    let eocd = &data[eocd_offset..];
    if eocd.len() < 22 {
        trace!("Buffer is too small to contain EOCD");
        return None;
    }

    let cd_size = u32::from_le_bytes(eocd[12..16].try_into().unwrap()) as usize;
    let cd_offset = u32::from_le_bytes(eocd[16..20].try_into().unwrap()) as usize;
    let comment_size = usize::from(u16::from_le_bytes(eocd[20..22].try_into().unwrap()));

    let start = eocd_offset.checked_sub(cd_size)?.checked_sub(cd_offset)?;
    let end = eocd_offset.checked_add(22)?.checked_add(comment_size)?;
    if end > data.len() {
        trace!("End of zip is out of bounds");
        return None;
    }

    trace!("Found zip bounds: {:?}", start..end);

    let archive = ZipArchive::from_slice(&data[start..end]).ok()?;
    let mut entries = archive.entries_safe();
    let mut matches = 0;

    while let Some((cd_entry, _)) = entries.next_entry().ok()? {
        let path = cd_entry.file_path_utf8().ok()?;

        if !path.ends_with(".x509.pem") {
            // otacerts.zip files only contain files named this way.
            trace!("Excluded due to invalid name: {path:?}");
            return None;
        }

        matches += 1;
    }

    if matches == 0 {
        // otacerts.zip files contain at least one cert.
        trace!("Zip is empty");
        return None;
    }

    debug!("Found otacerts.zip candidate");

    // There's one or more entries and every one is named *.x509.pem.
    Some(start..end)
}

/// Replace `otacerts.zip` with a new one containing the new certificate, but
/// padded to the same size. If the new zip is too large, the certificate will
/// be modified to remove unnecessary components until it fits. All operations
/// run in parallel where possible.
///
/// Returns two sorted and non-overlapping lists of byte ranges that were
/// modified. The first list are the byte regions within the filesystem data
/// that contained otacerts.zip. The second list is the list of byte regions
/// outside of the filesyste, like the hash tree, FEC data, and AVB metadata.
///
/// If [`Error::OldZipNotFound`] is returned, the output will not have been
/// modified.
#[allow(clippy::type_complexity)]
pub fn patch_system_image(
    raw_file: &(dyn ReadWriteAt + Sync),
    certificate: &Certificate,
    key: &SigningPrivateKey,
    method: SigningMethod,
    cancel_signal: &AtomicBool,
) -> Result<(Vec<Range<u64>>, Vec<Range<u64>>)> {
    // This must be a multiple of normal filesystem block sizes (eg. 4 KiB).
    // This ensures that the block containing otacerts.zip's data won't cross
    // chunk boundaries.
    const CHUNK_SIZE: u64 = 2 * 1024 * 1024;

    let parent_span = Span::current();

    let (mut header, footer, image_size) =
        avb::load_image(UserPosFile::new(raw_file)).map_err(Error::AvbUpdate)?;
    let Some(mut footer) = footer else {
        return Err(Error::NoFooter);
    };
    let AppendedDescriptorMut::HashTree(descriptor) =
        header.appended_descriptor_mut().map_err(Error::AvbUpdate)?
    else {
        return Err(Error::NoHashTreeDescriptor);
    };

    let num_chunks = footer.original_image_size.div_ceil(CHUNK_SIZE);
    trace!("Parallel heuristics search for otacerts.zip with {num_chunks} chunks");

    let modified_ranges = (0..num_chunks)
        .into_par_iter()
        .map(|chunk| -> Result<Vec<Range<u64>>> {
            stream::check_cancel(cancel_signal).map_err(Error::ReadData)?;

            let offset = chunk * CHUNK_SIZE;
            let size = CHUNK_SIZE.min(footer.original_image_size - offset);

            let mut file = UserPosFile::new(raw_file);
            file.seek(SeekFrom::Start(offset))
                .map_err(Error::ReadData)?;
            let buf = file
                .read_vec_exact(size as usize)
                .map_err(Error::ReadData)?;

            let mut ranges = Vec::<Range<u64>>::new();

            for eocd_offset_rel in memmem::find_iter(&buf, ota::ZIP_EOCD_MAGIC) {
                let _span = debug_span!(parent: &parent_span, "otacerts", offset, eocd_offset_rel)
                    .entered();

                let Some(bounds_rel) = find_zip_bounds(&buf, eocd_offset_rel) else {
                    continue;
                };

                let zip_size = bounds_rel.end - bounds_rel.start;
                let new_zip = otacert::create_zip_with_size(certificate, zip_size)
                    .map_err(Error::OtaCertZip)?;

                let bounds = offset + bounds_rel.start as u64..offset + bounds_rel.end as u64;

                stream::check_cancel(cancel_signal).map_err(Error::WriteData)?;

                file.seek(SeekFrom::Start(bounds.start))
                    .map_err(Error::WriteData)?;
                file.write_all(&new_zip).map_err(Error::WriteData)?;

                ranges.push(bounds);
            }

            Ok(ranges)
        })
        .try_reduce(Vec::new, |mut result, item| {
            result.extend(item);
            Ok(result)
        })?;

    if modified_ranges.is_empty() {
        return Err(Error::OldZipNotFound);
    }

    // Only need to update the hash tree and FEC data corresponding to the
    // modified regions.
    let update_ranges = Some(modified_ranges.as_slice());

    descriptor
        .update(raw_file, update_ranges, cancel_signal)
        .map_err(Error::AvbUpdate)?;

    if !header.public_key.is_empty() {
        debug!("Signing system image");
        header.set_algo_for_key(key).map_err(Error::AvbUpdate)?;
        header.sign(key, method).map_err(Error::AvbUpdate)?;
    }

    let file = UserPosFile::new(raw_file);
    avb::write_appended_image(file, &header, &mut footer, Some(image_size))
        .map_err(Error::AvbUpdate)?;

    let AppendedDescriptorMut::HashTree(descriptor) =
        header.appended_descriptor_mut().map_err(Error::AvbUpdate)?
    else {
        return Err(Error::NoHashTreeDescriptor);
    };

    // The hash tree, FEC data, and AVB regions will have been modified.
    let hash_tree_end = descriptor
        .tree_offset
        .checked_add(descriptor.tree_size)
        .ok_or(Error::IntOverflow("hash_tree_end"))?;
    let fec_data_end = descriptor
        .fec_offset
        .checked_add(descriptor.fec_size)
        .ok_or(Error::IntOverflow("fec_data_end"))?;
    let header_end = footer
        .vbmeta_offset
        .checked_add(footer.vbmeta_size)
        .ok_or(Error::IntOverflow("avb_end"))?;
    let footer_start = image_size - Footer::SIZE as u64;

    let other_ranges = util::merge_overlapping(&[
        descriptor.tree_offset..hash_tree_end,
        descriptor.fec_offset..fec_data_end,
        footer.vbmeta_offset..header_end,
        footer_start..image_size,
    ]);

    Ok((modified_ranges, other_ranges))
}
