use rand::thread_rng;
use rand::seq::IteratorRandom;
use std::collections::HashMap;
use crate::encoder;
use capstone::prelude::*;
use keystone::*;


#[derive(Debug, Clone)]
pub struct Instruction {
    pub address: u64,
    pub mnemonic: String,
    pub operand: String
}

impl Instruction {

    pub fn new(addr: u64, mnemonic: String, operand: String) -> Self {
        Instruction {address: addr, mnemonic: mnemonic, operand: operand}
    }
}

#[derive(Debug, Clone)]
pub struct AsmProgram {
    //pub instructions: HashMap<u64, Instruction>
    pub instructions: Vec<Instruction>
}

impl AsmProgram {
    pub fn new(bytecode: &Vec<u8>) -> Self {
        let cs = Capstone::new()
            .x86()
            .mode(arch::x86::ArchMode::Mode64)
            .syntax(arch::x86::ArchSyntax::Intel)
            .detail(true)
            .build()
            .expect("Failed to create Capstone object");


        let disas = cs.disasm_all(&bytecode, 0x0).unwrap();
        let mut instructions: Vec<Instruction> = Vec::new();
        for ins in disas.iter() {

            let tmpins = Instruction::new(ins.address(),
                                ins.mnemonic().unwrap().to_string(),
                                ins.op_str().unwrap().to_string());
            instructions.push(tmpins);
        }

        AsmProgram { instructions: instructions }
    }

    /// Output the ASM Instructions into a newline
    /// separated string
    pub fn to_string(&self) -> String {
        // Map the instructions to a vector of strings
        let s: Vec<String> = self.instructions.iter()
            .map(|c| format!("{} {}", c.mnemonic, c.operand))
            .collect();
        // return the joined() version of the instructions
        s.join("\n")
    }

    pub fn reorder(mut self) -> Vec<u8> {

        println!("[+] Patching JMP/Loop/Call Instructions");
        self.patch_jumps();
        println!("[*] Reordering Blocks");
        loop {
            match self.reorder_blocks() {
                Ok(v) => return v,
                Err(v) => {
                    match v.bits() {
                        161 => {
                            // KS_ERR_ASM_FIXUP_INVALID
                            continue;
                        }
                        _ => {
                            println!("[-] Failed to reorder blocks");
                            std::process::exit(1);
                        }
                    }
                }
            }
        }
    }
    fn patch_jumps(&mut self) {

        let mut label_ctr = 0;

        // As we need to modify several items inside
        let mut local_instructions = self.instructions.clone();
        // to track the indizes into the local instructions

        for instr in self.instructions.iter() {
            // Check if the Instruction is a JMP instruction
            if instr.mnemonic.chars().next().unwrap() == 'j'  || instr.mnemonic == "call" || instr.mnemonic == "loop" {

                // If an JMP instruction is found,
                // obtain the destination addrs and
                // insert a label before the corresponding
                // destination.

                // First, check if the operand begins with `0x`, if not,
                // the destination is a register
                if !instr.operand.contains("0x") {
                    // If the operand is a register, no action is needed as the value
                    // should be set correctly anyway
                    continue;
                }
                // else, the operand is parsed as u64
                let wo_prefix = instr.operand.trim_start_matches("0x");
                let jmp_destination = u64::from_str_radix(wo_prefix, 16).unwrap();

                let new_label = format!("patched_jmp{}:", label_ctr);

                // Get the position of the JMP destination in the instruction array
                let local_jmp_pos = local_instructions.clone().iter().position(|c| c.address == jmp_destination).unwrap();

                // Change the operand of the JMP instruction (JMP Label_{})
                let pos_local = local_instructions.clone().iter().position(|c| c.address == instr.address).unwrap();
                local_instructions[pos_local].operand = format!("patched_jmp{}", label_ctr);

                // Create the patch
                let new_instruction = Instruction::new(0, new_label, String::new());
                local_instructions.insert(local_jmp_pos, new_instruction);

                label_ctr += 1;
            }
        }

        // replace the instance vector with the local vector
        self.instructions = local_instructions;

    }

    /// Reorder the instructions
    pub fn reorder_blocks(&self) -> Result<Vec<u8>, keystone::Error> {

        let asm = self.to_string();

        const INS_PER_BLOCK: usize = 6;


        let instructions: Vec<String> = asm.split("\n").map(|c| c.to_string()).collect();


        // Split in Chunks, (Vector of String Vectors)
        let chunks: Vec<Vec<String>> = instructions.chunks(INS_PER_BLOCK).map(|c| c.to_vec()).collect();

        let in_between_chunks = &chunks[1..chunks.len()-1];
        let in_between_chunks = in_between_chunks.to_vec();


        let mut hm: HashMap<String, Vec<String>> = HashMap::new();

        let mut label_counter = 0;

        for mut chunk in in_between_chunks {
            chunk.insert(0, format!("label_{}:", label_counter));
            chunk.push(format!("jmp label_{}", label_counter+1));
            hm.insert(format!("label_{}", label_counter), chunk);
            label_counter += 1;
        }

        let mut shuffled_payload: Vec<String> = Vec::new();
        shuffled_payload.push(chunks[0].join("\n"));
        shuffled_payload.push(format!("jmp label_0"));
        //hm.shuffle(&mut thread_rng());


        // Loop though the len of bloocks and choose a random block
        for _ in 0..hm.len() {

            // Get the labels of the blocks
            let labels =  hm.keys().cloned();

            // Choose a random label
            let key = labels.choose(&mut thread_rng()).unwrap();

            // Get the Instructions for that choosen label, generate string and
            // Attach to the final, shuffled payload
            let instructions = hm.get(&key).unwrap();
            let instructions = instructions.join("\n");
            shuffled_payload.push(instructions);

            // Remove that key, as its already processe
            hm.remove(&key);
        }

        let chunk_len = chunks.len();
        shuffled_payload.push(format!("label_{}:", chunk_len-2));
        shuffled_payload.push(chunks[chunk_len-1].join("\n"));

        let asm_out = shuffled_payload.join("\n");

        // return the assembled bytes
        encoder::assemble_asm(asm_out, OptionValue::SYNTAX_INTEL)

    }


}
