//! Representation of a method.

use anyhow::anyhow;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};

use pyo3::prelude::*;

use crate::{
    ins, ins::Instruction, DexString, IdField, IdMethod, IdMethodType, IdType, MethodHandle, Result,
};

// TODO: make this easy to edit/manipulate, maybe move to Method

// type TmpHandlerType = (Vec<(IdType, u32)>, Option<u32>);

/// The code run by a method.
#[pyclass]
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Code {
    // TODO: remove and compute this value from code?
    /// The number of registers used by the code
    #[pyo3(get)]
    pub registers_size: u16,
    // TODO: what does it means? is it computable?
    /// The number of words of incoming arguments to the method
    #[pyo3(get)]
    pub ins_size: u16,
    // TODO: what does it means? is it computable?
    /// The number of words of outgoing argument space
    #[pyo3(get)]
    pub outs_size: u16,
    // TODO: implement
    /// The debug info
    #[pyo3(get)]
    pub debug_info: (u32, Vec<u8>), // Should be stripped, copying like this just don't work
    /// The names of the parameters if given
    #[pyo3(get)]
    pub parameter_names: Option<Vec<Option<DexString>>>,
    /// The instructions.
    #[pyo3(get)]
    pub insns: Vec<Instruction>,
}

impl PartialEq for Code {
    fn eq(&self, other: &Self) -> bool {
        let comparable_self = self.semantic_comparable().unwrap();
        let comparable_other = other.semantic_comparable().unwrap();
        (comparable_self.registers_size == comparable_other.registers_size)
            && (comparable_self.ins_size == comparable_other.ins_size)
            && (comparable_self.outs_size == comparable_other.outs_size)
            && (comparable_self.debug_info == comparable_other.debug_info)
            && (comparable_self.insns == comparable_other.insns)
    }
}

// TODO reimplement PartialEq: label should become address independant

#[pymethods]
impl Code {
    pub fn to_json(&self) -> Result<String> {
        Ok(serde_json::to_string(self)?)
    }

    #[staticmethod]
    pub fn from_json(json: &str) -> Result<Self> {
        Ok(serde_json::from_str(json)?)
    }

    #[new]
    pub fn new(
        registers_size: u16,
        ins_size: u16,
        outs_size: u16,
        insns: Vec<Instruction>,
        parameter_names: Option<Vec<Option<DexString>>>,
    ) -> Self {
        Self {
            registers_size,
            ins_size,
            outs_size,
            insns,
            parameter_names,
            debug_info: (0, vec![]),
        }
    }

    pub fn __str__(&self) -> String {
        self.__repr__()
    }

    pub fn __repr__(&self) -> String {
        "Code()".into()
    }

    /// Return all strings referenced in the codes.
    pub fn get_all_strings(&self) -> HashSet<DexString> {
        let mut strings = HashSet::new();
        // TODO debug info contains types and strings
        for ins in &self.insns {
            strings.extend(ins.get_all_strings());
        }
        if let Some(names) = &self.parameter_names {
            for name in names {
                if let Some(name) = name {
                    strings.insert(name.clone());
                }
            }
        }
        strings
    }

    /// Return all types referenced in the codes.
    pub fn get_all_types(&self) -> HashSet<IdType> {
        let mut type_ids = HashSet::new();
        // TODO debug info contains types and strings
        for ins in &self.insns {
            type_ids.extend(ins.get_all_types());
        }
        type_ids
    }

    /// Return all prototypes referenced in the codes.
    pub fn get_all_protos(&self) -> HashSet<IdMethodType> {
        let mut protos = HashSet::new();
        for ins in &self.insns {
            protos.extend(ins.get_all_protos());
        }
        protos
    }

    /// Return all field ids referenced in the codes.
    pub fn get_all_field_ids(&self) -> HashSet<IdField> {
        let mut fields = HashSet::new();
        for ins in &self.insns {
            fields.extend(ins.get_all_field_ids());
        }
        fields
    }

    /// Return all method ids referenced in the codes.
    pub fn get_all_method_ids(&self) -> HashSet<IdMethod> {
        let mut methods = HashSet::new();
        for ins in &self.insns {
            methods.extend(ins.get_all_method_ids())
        }
        methods
    }

    /// Return all method handles referenced in the codes.
    pub fn get_all_method_handles(&self) -> HashSet<MethodHandle> {
        let mut handles = HashSet::new();
        for ins in &self.insns {
            handles.extend(ins.get_all_method_handles());
        }
        handles
    }

    pub fn __eq__(&self, other: &Self) -> bool {
        self == other
    }

    /// Return all the labels used by instrutions in the code.
    pub fn get_referenced_label(&self) -> HashSet<String> {
        let mut used_labels = HashSet::new();
        for ins in &self.insns {
            match ins {
                Instruction::Goto(ins::Goto { label }) => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfEq(ins::IfEq { label, .. }) => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfNe(ins::IfNe { label, .. }) => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfLt(ins::IfLt { label, .. }) => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfGe(ins::IfGe { label, .. }) => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfGt(ins::IfGt { label, .. }) => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfLe(ins::IfLe { label, .. }) => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfEqZ(ins::IfEqZ { label, .. }) => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfNeZ(ins::IfNeZ { label, .. }) => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfLtZ(ins::IfLtZ { label, .. }) => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfGeZ(ins::IfGeZ { label, .. }) => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfGtZ(ins::IfGtZ { label, .. }) => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfLeZ(ins::IfLeZ { label, .. }) => {
                    used_labels.insert(label.clone());
                }
                Instruction::Try(ins::Try {
                    end_label,
                    handlers,
                    default_handler,
                }) => {
                    used_labels.insert(end_label.clone());
                    for (_, label) in handlers {
                        used_labels.insert(label.clone());
                    }
                    if let Some(label) = default_handler {
                        used_labels.insert(label.clone());
                    }
                }
                Instruction::Switch(ins::Switch { branches, .. }) => {
                    for label in branches.values() {
                        used_labels.insert(label.clone());
                    }
                }
                _ => (),
            };
        }
        used_labels
    }

    /// Generate a new code with normalized labels and no nops.
    /// This allows to compare codes with same semantic but different labels
    /// and padding.
    /// (ig when the same code was reserialized)
    pub fn semantic_comparable(&self) -> Result<Self> {
        // TODO: keep the nops that are not used for padding
        let used_labels = self.get_referenced_label();
        let mut new_labels = HashMap::new();
        let mut label_id = 0;
        let mut last_ins_was_a_label = false;

        for ins in &self.insns {
            match ins {
                Instruction::Label(ins::Label { name }) => {
                    if used_labels.get(name).is_none() {
                        continue;
                    }
                    let new_label_id = if last_ins_was_a_label {
                        label_id
                    } else {
                        label_id += 1;
                        label_id
                    };
                    new_labels.insert(name.clone(), format!("label_{new_label_id}"));

                    last_ins_was_a_label = true;
                }
                _ => last_ins_was_a_label = false,
            }
        }

        for label in &used_labels {
            if new_labels.get(label).is_none() {
                println!("{label} use but not in new_labels");
            }
        }

        let mut new_insns = vec![];
        last_ins_was_a_label = false;
        for ins in self.insns.iter().cloned() {
            match ins {
                Instruction::Goto(mut instr) => {
                    last_ins_was_a_label = false;
                    instr.label = new_labels
                        .get(&instr.label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            instr.label
                        ))?
                        .clone();
                    new_insns.push(Instruction::Goto(instr));
                }
                Instruction::IfEq(mut instr) => {
                    last_ins_was_a_label = false;
                    instr.label = new_labels
                        .get(&instr.label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            instr.label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfEq(instr));
                }
                Instruction::IfNe(mut instr) => {
                    last_ins_was_a_label = false;
                    instr.label = new_labels
                        .get(&instr.label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            instr.label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfNe(instr));
                }
                Instruction::IfLt(mut instr) => {
                    last_ins_was_a_label = false;
                    instr.label = new_labels
                        .get(&instr.label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            instr.label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfLt(instr));
                }
                Instruction::IfGe(mut instr) => {
                    last_ins_was_a_label = false;
                    instr.label = new_labels
                        .get(&instr.label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            instr.label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfGe(instr));
                }
                Instruction::IfGt(mut instr) => {
                    last_ins_was_a_label = false;
                    instr.label = new_labels
                        .get(&instr.label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            instr.label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfGt(instr));
                }
                Instruction::IfLe(mut instr) => {
                    last_ins_was_a_label = false;
                    instr.label = new_labels
                        .get(&instr.label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            instr.label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfLe(instr));
                }
                Instruction::IfEqZ(mut instr) => {
                    last_ins_was_a_label = false;
                    instr.label = new_labels
                        .get(&instr.label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            instr.label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfEqZ(instr));
                }
                Instruction::IfNeZ(mut instr) => {
                    last_ins_was_a_label = false;
                    instr.label = new_labels
                        .get(&instr.label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            instr.label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfNeZ(instr));
                }
                Instruction::IfLtZ(mut instr) => {
                    last_ins_was_a_label = false;
                    instr.label = new_labels
                        .get(&instr.label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            instr.label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfLtZ(instr));
                }
                Instruction::IfGeZ(mut instr) => {
                    last_ins_was_a_label = false;
                    instr.label = new_labels
                        .get(&instr.label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            instr.label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfGeZ(instr));
                }
                Instruction::IfGtZ(mut instr) => {
                    last_ins_was_a_label = false;
                    instr.label = new_labels
                        .get(&instr.label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            instr.label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfGtZ(instr));
                }
                Instruction::IfLeZ(mut instr) => {
                    last_ins_was_a_label = false;
                    instr.label = new_labels
                        .get(&instr.label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            instr.label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfLeZ(instr));
                }
                Instruction::Try(mut instr) => {
                    last_ins_was_a_label = false;
                    instr.end_label = new_labels
                        .get(&instr.end_label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            instr.end_label
                        ))?
                        .clone();
                    for i in 0..instr.handlers.len() {
                        instr.handlers[i].1 = new_labels
                            .get(&instr.handlers[i].1)
                            .ok_or(anyhow!(
                                "Internal error: {} not found in renamed label",
                                instr.handlers[i].1
                            ))?
                            .clone();
                    }
                    if let Some(label) = instr.default_handler {
                        instr.default_handler = Some(
                            new_labels
                                .get(&label)
                                .ok_or(anyhow!(
                                    "Internal error: {} not found in renamed label",
                                    label
                                ))?
                                .clone(),
                        );
                    }
                    new_insns.push(Instruction::Try(instr));
                }
                Instruction::Switch(mut instr) => {
                    last_ins_was_a_label = false;
                    for label in instr.branches.values_mut() {
                        *label = new_labels
                            .get(label)
                            .ok_or(anyhow!(
                                "Internal error: {} not found in renamed label",
                                label
                            ))?
                            .clone();
                    }
                    new_insns.push(Instruction::Switch(instr));
                }
                Instruction::Label(ins::Label { name }) => {
                    if used_labels.get(&name).is_none() {
                        //println!("{name} not used");
                        continue;
                    }
                    if !last_ins_was_a_label {
                        new_insns.push(Instruction::Label(ins::Label {
                            name: new_labels
                                .get(&name)
                                .ok_or(anyhow!(
                                    "Internal error: {} not found in renamed label",
                                    name
                                ))?
                                .clone(),
                        }));
                    }
                    last_ins_was_a_label = true;
                }
                Instruction::Nop(_) => (),
                instr => {
                    last_ins_was_a_label = false;
                    new_insns.push(instr);
                }
            };
        }
        Ok(Self {
            insns: new_insns,
            ..self.clone()
        })
    }
}
