use super::*;
use androscalpel_serializer::Instruction as InsFormat;
use androscalpel_serializer::*;
use std::fs::File;
use std::io;
use std::sync::OnceLock;

fn get_dex(filename: &str) -> Vec<u8> {
    let hello_world_dex = format!("{}/src/tests/{}", env!("CARGO_MANIFEST_DIR"), filename);
    let mut file = File::open(&hello_world_dex).expect(&format!("{} not found", filename));
    let mut data = vec![];
    io::copy(&mut file, &mut data).unwrap();
    data
}

fn get_hello_world_dex() -> &'static [u8] {
    static HELLO_WORLD_DEX: OnceLock<Vec<u8>> = OnceLock::new();
    HELLO_WORLD_DEX.get_or_init(|| get_dex("classes_hello_world.dex"))
}

fn get_hello_world_apk() -> &'static Apk {
    static HELLO_WORLD_APK: OnceLock<Apk> = OnceLock::new();
    HELLO_WORLD_APK.get_or_init(|| {
        let mut apk = Apk::new();
        apk.add_dex_file(get_hello_world_dex()).unwrap();
        apk
    })
}

fn get_hello_world_recompilled() -> &'static [u8] {
    static HELLO_WORLD_RECOMP: OnceLock<Vec<u8>> = OnceLock::new();
    HELLO_WORLD_RECOMP.get_or_init(|| get_hello_world_apk().gen_raw_dex().unwrap().pop().unwrap())
}

fn get_class_dex<'a, 'b>(name: &'a str, dex: &'b DexFileReader) -> Option<&'b ClassDefItem> {
    let name = name.into();
    for class in dex.get_class_defs() {
        let class_name = dex
            .get_string(
                dex.get_type_id(class.class_idx as usize)
                    .unwrap()
                    .descriptor_idx,
            )
            .unwrap();
        if class_name == name {
            return Some(class);
        }
    }
    None
}

fn get_method_code_dex(name: &str, dex: &DexFileReader) -> Option<CodeItem> {
    let method = IdMethod::from_smali(name).unwrap();
    let class_name: String = (&method.class_.0).into();
    let class = get_class_dex(&class_name, dex);
    let class = if let Some(class) = class {
        class
    } else {
        return None;
    };
    if class.class_data_off == 0 {
        return None;
    }
    let data = dex
        .get_struct_at_offset::<ClassDataItem>(class.class_data_off)
        .unwrap();
    let mut idx = 0;
    for meth in data.direct_methods {
        idx += meth.method_idx_diff.0;
        let meth_id = Apk::get_id_method_from_idx(idx as usize, dex).unwrap();
        println!("{}", meth_id.__str__());
        if meth_id == method && meth.code_off.0 != 0 {
            return Some(
                dex.get_struct_at_offset::<CodeItem>(meth.code_off.0)
                    .unwrap(),
            );
        }
    }
    let mut idx = 0;
    for meth in data.virtual_methods {
        idx += meth.method_idx_diff.0;
        let meth_id = Apk::get_id_method_from_idx(idx as usize, dex).unwrap();
        println!("{}", meth_id.__str__());
        if meth_id == method && meth.code_off.0 != 0 {
            return Some(
                dex.get_struct_at_offset::<CodeItem>(meth.code_off.0)
                    .unwrap(),
            );
        }
    }
    None
}

#[test]
fn test_generated_data_size() {
    let dex_bin = get_hello_world_recompilled();
    let dex = DexFileReader::new(dex_bin).unwrap();

    assert_eq!(
        dex.get_header().data_off + dex.get_header().data_size,
        dex_bin.len() as u32
    )
    // TODO: check for all pool concerned if the pool span outside the data section?
    //for item in dex.get_map_list().list() {}
}

#[test]
fn test_generated_apk_equivalence() {
    let new_dex = get_hello_world_recompilled();
    let mut new_apk = Apk::new();
    new_apk.add_dex_file(&new_dex).unwrap();
    assert_eq!(get_hello_world_apk(), &new_apk);
}

#[test]
fn test_string_order_in_dex() {
    use std::collections::HashSet;
    let dex = DexFileReader::new(get_hello_world_dex()).unwrap();
    let new_dex = DexFileReader::new(get_hello_world_recompilled()).unwrap();
    let mut string_set = HashSet::new();
    for i in 0..new_dex.get_string_ids().len() {
        string_set.insert(DexString(new_dex.get_string(i as u32).unwrap()));
    }
    for i in 0..dex.get_string_ids().len() {
        let string = DexString(dex.get_string(i as u32).unwrap());
        if !string_set.contains(&string) {
            println!("{}", string.__str__());
        }
    }
    assert_eq!(
        dex.get_string_ids().len(),
        new_dex.get_string_ids().len(),
        "Not the same number of strings"
    );
    /* this is hell to debug
    let strings: Vec<_> = (0..dex.get_string_ids().len())
        .map(|idx| DexString(dex.get_string(idx as u32).unwrap()))
        .collect();
    let new_strings: Vec<_> = (0..new_dex.get_string_ids().len())
        .map(|idx| DexString(new_dex.get_string(idx as u32).unwrap()))
        .collect();
    assert_eq!(strings, new_strings);
    */
    for idx in 0..dex.get_string_ids().len() {
        let string = DexString(dex.get_string(idx as u32).unwrap());
        let new_string = DexString(new_dex.get_string(idx as u32).unwrap());
        assert_eq!(
            string, new_string,
            "Strings {idx} do not match, left original, right recompiled"
        );
    }
}

#[test]
fn test_switch_table_alignement() {
    let new_dex = DexFileReader::new(get_hello_world_recompilled()).unwrap();
    let code = get_method_code_dex(
        "Landroidx/appcompat/app/AppCompatViewInflater;\
        ->createView(\
          Landroid/view/View;Ljava/lang/String;\
          Landroid/content/Context;Landroid/util/AttributeSet;\
          ZZZZ\
        )Landroid/view/View;",
        &new_dex,
    )
    .unwrap();
    let mut addr = 0;
    for ins in code.insns {
        match ins {
            InsFormat::FormatPackedSwitchPayload { .. } => {
                assert_eq!(addr % 4, 0, "Switch table not aligned")
            }
            InsFormat::FormatSparseSwitchPayload { .. } => {
                assert_eq!(addr % 4, 0, "Switch table not aligned")
            }
            _ => (),
        }
        addr += ins.size();
    }
}

#[test]
fn test_array_table_alignement() {
    let new_dex = DexFileReader::new(get_hello_world_recompilled()).unwrap();
    let code = get_method_code_dex(
        "Landroidx/constraintlayout/core/widgets/ConstraintWidget;-><init>()V",
        &new_dex,
    )
    .unwrap();
    let mut addr = 0;
    for ins in code.insns {
        match ins {
            InsFormat::FormatFillArrayDataPayload { .. } => {
                assert_eq!(addr % 4, 0, "Array Data table not aligned")
            }
            _ => (),
        }
        addr += ins.size();
    }
}

#[test]
fn test_parse_type_list() {
    let types = IdType::get_list_from_str("IILjavalangObject;LjavalangObject;Z").unwrap();
    assert_eq!(
        types,
        vec![
            IdType::int(),
            IdType::int(),
            IdType::class("javalangObject"),
            IdType::class("javalangObject"),
            IdType::boolean(),
        ]
    )
}

#[test]
fn test_parse_proto_smali() {
    let proto = IdMethodType::from_smali(
        "(LandroidxcoreutilPredicate;LandroidxcoreutilPredicate;LjavalangObject;)Z",
    )
    .unwrap();
    assert_eq!(
        proto,
        IdMethodType::new(
            IdType::boolean(),
            vec![
                IdType::class("androidxcoreutilPredicate"),
                IdType::class("androidxcoreutilPredicate"),
                IdType::class("javalangObject")
            ]
        )
    );
}

#[test]
fn test_parse_method_smali() {
    let id_method = IdMethod::from_smali(
         "Landroidx/core/util/Predicate;->lambda$and$0(Landroidx/core/util/Predicate;Landroidx/core/util/Predicate;Ljava/lang/Object;)Z"
       ).unwrap();

    assert_eq!(
        id_method,
        IdMethod::new(
            "lambda$and$0".into(),
            IdMethodType::new(
                IdType::boolean(),
                vec![
                    IdType::class("androidx/core/util/Predicate"),
                    IdType::class("androidx/core/util/Predicate"),
                    IdType::class("java/lang/Object")
                ]
            ),
            IdType::class("androidx/core/util/Predicate")
        )
    );
}

#[test]
fn test_parse_field_smali() {
    let proto = IdField::from_smali(
        "Ljava/lang/annotation/ElementType;->METHOD:Ljava/lang/annotation/ElementType;",
    )
    .unwrap();
    assert_eq!(
        proto,
        IdField::new(
            "METHOD".into(),
            IdType::class("java/lang/annotation/ElementType"),
            IdType::class("java/lang/annotation/ElementType"),
        )
    );
}

#[test]
fn test_sort_fields() {
    let f1 = IdField::from_smali(
        "Landroidx/lifecycle/WithLifecycleStateKt$suspendWithStateAtLeastUnchecked$2$observer$1;\
        ->$co:Lkotlinx/coroutines/CancellableContinuation;",
    )
    .unwrap();
    let f2 = IdField::from_smali(
        "Lcom/google/android/material/search/SearchBarAnimationHelper$$ExternalSyntheticLambda4;\
        ->f$0:Landroid/animation/Animator;",
    )
    .unwrap();
    let mut list1 = vec![f2.clone(), f1.clone()];
    let mut list2 = vec![f1.clone(), f2.clone()];
    list1.sort();
    list2.sort();
    assert_eq!(list1, list2);
    assert_eq!(list1, vec![f1.clone(), f2.clone()]);
}

#[test]
fn test_sort_types() {
    let t1 = IdType::from_smali(
        "Landroidx/lifecycle/WithLifecycleStateKt$suspendWithStateAtLeastUnchecked$2$observer$1;",
    )
    .unwrap();
    let t2 = IdType::from_smali(
        "Lcom/google/android/material/search/SearchBarAnimationHelper$$ExternalSyntheticLambda4;",
    )
    .unwrap();
    let mut list1 = vec![t2.clone(), t1.clone()];
    let mut list2 = vec![t1.clone(), t2.clone()];
    list1.sort();
    list2.sort();
    assert_eq!(list1, list2);
    assert_eq!(list1, vec![t1.clone(), t2.clone()]);
}

#[test]
fn test_sort_strings() {
    let s1: DexString =
        "Landroidx/lifecycle/WithLifecycleStateKt$suspendWithStateAtLeastUnchecked$2$observer$1;"
            .into();
    let s2: DexString =
        "Lcom/google/android/material/search/SearchBarAnimationHelper$$ExternalSyntheticLambda4;"
            .into();

    let mut list1 = vec![s2.clone(), s1.clone()];
    let mut list2 = vec![s1.clone(), s2.clone()];
    list1.sort();
    list2.sort();
    assert_eq!(list1, list2);
    assert_eq!(list1, vec![s1.clone(), s2.clone()]);
}
