//! Function that are can be usefull on the python side.

use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyBytesMethods};

use std::collections::{HashMap, HashSet};
use std::fs::File;
use std::io::{Cursor, Seek, SeekFrom};
use std::path::PathBuf;

use crate::{Apk, IdType, Result};
use androscalpel_serializer::{
    DexFileReader, HeaderItem, Serializable, Sleb128, Uleb128, Uleb128p1,
};
use apk_frauder::{end_of_central_directory::EndCentralDirectory, ZipFileReader};

/// Convert an integer to the uleb128 byte encoding
#[pyfunction]
pub fn int_to_uleb128(py: Python, x: u32) -> Result<PyObject> {
    Ok(PyBytes::new(py, &Uleb128(x).serialize_to_vec()?).into())
}

/// Convert an integer to the uleb128p1 byte encoding
#[pyfunction]
pub fn int_to_uleb128p1(py: Python, x: u32) -> Result<PyObject> {
    Ok(PyBytes::new(py, &Uleb128p1(x).serialize_to_vec()?).into())
}
/// Convert an integer to the sleb128 byte encoding
#[pyfunction]
pub fn int_to_sleb128(py: Python, x: i32) -> Result<PyObject> {
    Ok(PyBytes::new(py, &Sleb128(x).serialize_to_vec()?).into())
}

/// Decode an uleb128 encoded integer
#[pyfunction]
pub fn uleb128_to_int(b: &[u8]) -> Result<u32> {
    Ok(Uleb128::deserialize_from_slice(b)?.0)
}

/// Decode an uleb128p1 encoded integer
#[pyfunction]
pub fn uleb128p1_to_int(b: &[u8]) -> Result<u32> {
    Ok(Uleb128p1::deserialize_from_slice(b)?.0)
}

/// Decode an sleb128 encoded integer
#[pyfunction]
pub fn sleb128_to_int(b: &[u8]) -> Result<i32> {
    Ok(Sleb128::deserialize_from_slice(b)?.0)
}

// TODO: list_defined_classes, is_dex, is_zip take only &[u8] or file, but should allow to also read from both

/// List all classes defined in a dex file.
#[pyfunction]
pub fn list_defined_classes(dex: &[u8]) -> Result<HashSet<IdType>> {
    let dex = DexFileReader::new(dex)?;
    dex.get_class_defs()
        .iter()
        .map(|cdef| Apk::get_id_type_from_idx(cdef.class_idx as usize, &dex))
        .collect()
}

/// Test if a file is as .dex file an return the dex version if it is, else return None.
#[pyfunction]
pub fn is_dex(file: PathBuf) -> Option<usize> {
    let mut file = match File::open(file) {
        Ok(file) => file,
        Err(_) => return None,
    };
    HeaderItem::deserialize(&mut file)
        .ok()
        .and_then(|header| String::from_utf8(header.magic.version.to_vec()).ok())
        .and_then(|version| version.parse::<usize>().ok())
}

/// Test if a file is a zip file.
#[pyfunction]
pub fn is_zip(file: PathBuf) -> bool {
    let mut file = match File::open(file) {
        Ok(file) => file,
        Err(_) => return false,
    };
    let ecd_off = if let Some(off) = ZipFileReader::get_end_of_central_directory_offset(&mut file) {
        off
    } else {
        return false;
    };
    if file.seek(SeekFrom::Start(ecd_off)).is_err() {
        return false;
    }
    if let Ok(sig) = apk_frauder::Signature::deserialize(&mut file) {
        EndCentralDirectory::SIGNATURE == sig
    } else {
        false
    }
}

/// Replace the dex files a an apk an resigned the apk.
///
/// # Warning
///
/// For now, only jks keystore are allowed.
///
/// The `zipalign` and `apksigner` args allow to use a specific version of the
/// tools instead of the one in the PATH (if it even exist)
///
/// `additionnal_files` is a dict of file to add, modify or remove in the apk.
/// The keys are the file names and the values are `None` to remove the file, or
/// `bytes` for the content of the file.
#[pyfunction]
#[pyo3(signature = (
    apk, 
    dst, 
    dexfiles, 
    keystore, 
    zipalign=None, 
    apksigner=None, 
    additionnal_files=None
))]
pub fn replace_dex(
    apk: PathBuf,
    dst: PathBuf,
    dexfiles: Vec<Bound<'_, PyBytes>>,
    keystore: PathBuf,
    zipalign: Option<PathBuf>,
    apksigner: Option<PathBuf>,
    additionnal_files: Option<HashMap<String, Option<Bound<'_, PyBytes>>>>,
) {
    let mut dexfiles: Vec<_> = dexfiles
        .iter()
        .map(PyBytesMethods::as_bytes)
        .map(Cursor::new)
        .collect();
    let additionnal_files: Option<HashMap<_, _>> = additionnal_files.as_ref().map(|additionnal_files|
        additionnal_files.iter()
            .map(|(k, v)| (
                k.clone(), 
                v.as_ref().map(|bytes| bytes.as_bytes()).map(Cursor::new)
            )).collect()
        );
    apk_frauder::replace_dex(
        apk,
        dst,
        &mut dexfiles,
        keystore,
        zipalign,
        apksigner,
        additionnal_files,
    )
}

/// export the function in a python module
pub(crate) fn export_module(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(int_to_uleb128, m)?)?;
    m.add_function(wrap_pyfunction!(int_to_uleb128p1, m)?)?;
    m.add_function(wrap_pyfunction!(int_to_sleb128, m)?)?;
    m.add_function(wrap_pyfunction!(uleb128_to_int, m)?)?;
    m.add_function(wrap_pyfunction!(uleb128p1_to_int, m)?)?;
    m.add_function(wrap_pyfunction!(sleb128_to_int, m)?)?;
    m.add_function(wrap_pyfunction!(list_defined_classes, m)?)?;
    m.add_function(wrap_pyfunction!(is_dex, m)?)?;
    m.add_function(wrap_pyfunction!(is_zip, m)?)?;
    m.add_function(wrap_pyfunction!(replace_dex, m)?)?;
    Ok(())
}
