//! This example allows loading a single zip file, you can then list all
//! contained files and extract the compressed data from a single file within

use bytes::Buf;
use clap::Parser;
use std::{
    fs::File,
    io::{Read, Write},
    path::PathBuf,
};
use zip_core::{raw::parse::Parse, signature::Signature, structs::CompressionMethod};

/// Extract Single file from zip
#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Args {
    /// file to read
    #[arg(short, long)]
    input: PathBuf,

    /// file to write to (if not specified, output to tty)
    #[arg(short, long)]
    output: Option<PathBuf>,

    /// file to select for export
    #[arg(long)]
    selected_file: Option<String>,

    #[arg(long, default_value_t = true)]
    raw: bool,
}

pub fn main() {
    let args = Args::parse();
    println!("Loading whole file to memory");
    let mut file = File::open(&args.input).expect("Problem opening file");
    let mut buf = Vec::new();
    file.read_to_end(&mut buf).expect("Problem reading file to end");
    println!("File read successful");

    println!("Find EndOfCentralDirectory");
    const EXPECTED_POSITION_OFFSET: usize = 10_000;
    const END_SIGNATURE: [u8; 4] =
        zip_core::raw::EndOfCentralDirectoryFixed::END_OF_CENTRAL_DIR_SIGNATURE.to_le_bytes();
    let lower_start = buf.len().saturating_sub(EXPECTED_POSITION_OFFSET);
    // first check the last 10000 bytes, thats usually quick, if not found, do a
    // full file scan
    let end_of_central_directory_pos = zip_core::raw::parse::find_next_signature(&buf[lower_start..], END_SIGNATURE)
        .map(|v| v + lower_start)
        .or_else(|| zip_core::raw::parse::find_next_signature(&buf, END_SIGNATURE))
        .expect("No End Of Central Directory Found, is this a zip file?");
    println!("EndOfCentralDirectory found at byte: {end_of_central_directory_pos}");

    let mut eoc_buf = &buf[end_of_central_directory_pos..];
    let end_of_central_directory = zip_core::raw::EndOfCentralDirectory::from_buf(&mut eoc_buf)
        .expect("It seems like EndOfCentralDirectory is invalid");
    if !end_of_central_directory.is_valid_signature() {
        println!("WARN, EndOfCentralDirectory has wrong signature");
    }
    println!("EndOfCentralDirectory valid and parsed");

    let eoc_f = &end_of_central_directory.fixed;
    if eoc_f.total_number_of_entries_in_the_central_directory
        != eoc_f.total_number_of_entries_in_the_central_directory_on_this_disk
        || 0 != eoc_f.number_of_the_disk_with_the_start_of_the_central_directory
        || 0 != eoc_f.number_of_this_disk
    {
        println!("WARN: it seems like you have a multi-file ZIP file, this is not supported");
    }

    println!("parsing CentralDirectrories");
    let mut cd_buf = &buf[eoc_f.offset_of_start_of_central_directory_with_respect_to_the_starting_disk_number as usize
        ..eoc_f.offset_of_start_of_central_directory_with_respect_to_the_starting_disk_number as usize
            + eoc_f.size_of_the_central_directory as usize];
    println!("CentralDirectrories parsed");

    let mut cds = Vec::new();
    while let Ok(cd) = zip_core::raw::CentralDirectoryHeader::from_buf(&mut cd_buf) {
        let is_file = cd.fixed.general_purpose_bit_flag != 0 && cd.fixed.uncompressed_size != 0;
        if !cd.is_valid_signature() {
            println!("WARN, CentralDirectory has wrong signature");
        }
        if is_file {
            cds.push(cd);
        }
    }

    let mut selected = None;
    if let Some(try_select) = args.selected_file {
        selected = cds.iter().position(|cd| {
            let filename = std::str::from_utf8(&cd.file_name).expect("non UTF-8 filename found");
            filename == try_select
        });
        if selected.is_none() {
            println!("WARN: you selected a file, but that couldn't be found, please select");
        }
    }

    if selected.is_none() {
        println!("\nselect file to extract:");
        for (i, cd) in cds.iter().enumerate() {
            let filename = std::str::from_utf8(&cd.file_name).expect("non UTF-8 filename found");
            println!("{i}: {filename}");
        }
        println!("Please select number to extract:");
        let mut user_sel: String = String::new();
        let _ = std::io::stdin()
            .read_line(&mut user_sel)
            .expect("Did not enter a correct string");
        selected = Some(
            user_sel
                .trim()
                .parse()
                .expect("You didn't add a number :/ Shame on you"),
        );
    }

    let selected = selected.unwrap();
    let selected_cd = cds[selected].clone();

    let filename = std::str::from_utf8(&selected_cd.file_name).expect("non UTF-8 filename found");
    let compressed_size = selected_cd.fixed.compressed_size as usize;
    let compression_method = selected_cd.fixed.compression_method;
    let compression_method = CompressionMethod::try_from(compression_method).expect("Unknown compression method used");
    println!("Trying to extract file #{selected}: {filename}");

    let mut lfh_buf = &buf[selected_cd.fixed.relative_offset_of_local_header as usize..];

    let lfh = zip_core::raw::LocalFileHeader::from_buf(&mut lfh_buf).expect("Error reading LocalFileHeader");
    if !lfh.is_valid_signature() {
        println!("WARN: LocalFileHeader has an invalid Signature");
    }
    println!("LocalFileHeader valid");

    let compressed_data = Buf::take(lfh_buf, compressed_size);
    let mut show_data = if args.raw {
        compressed_data
    } else {
        match compression_method {
            CompressionMethod::Stored => compressed_data,
            other => {
                panic!("Sorry, we don't support the format {other:?} yet");
            },
        }
    };
    println!("File extracted");

    if let Some(outfile) = args.output {
        let mut file = File::create(outfile).expect("Error creating outfile");
        let bytes = show_data.copy_to_bytes(show_data.remaining());
        file.write_all(&bytes).unwrap();
    } else {
        let bytes = show_data.copy_to_bytes(show_data.remaining());
        std::io::stdout().write_all(&bytes).expect("Couldn't write to stdout");
    }
}