//! Representation of a method.

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

#[cfg(feature = "python")]
use pyo3::prelude::*;

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

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

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

/// The code run by a method.
#[cfg_attr(feature = "python", pyclass(eq))]
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Code {
    // TODO: remove and compute this value from code?
    /// The number of registers used by the code
    #[cfg_attr(feature = "python", pyo3(get))]
    pub registers_size: u16,
    // TODO: what does it means? is it computable?
    /// The number of words of incoming arguments to the method
    #[cfg_attr(feature = "python", pyo3(get))]
    pub ins_size: u16,
    // TODO: what does it means? is it computable?
    /// The number of words of outgoing argument space
    #[cfg_attr(feature = "python", pyo3(get))]
    pub outs_size: u16,
    /// The names of the parameters if given
    #[cfg_attr(feature = "python", pyo3(get))]
    pub parameter_names: Option<Vec<Option<DexString>>>,
    /// The instructions.
    #[cfg_attr(feature = "python", pyo3(get))]
    pub insns: Vec<Instruction>,
}

// TODO reimplement PartialEq: label should become address independant
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.insns == comparable_other.insns)
    }
}

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

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

    #[cfg_attr(feature = "python", new)]
    #[cfg_attr(feature = "python", pyo3(signature = (registers_size, ins_size, outs_size, insns, parameter_names=None)))]
    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,
        }
    }

    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.iter().flatten() {
                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
    }

    /// 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 { label } => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfEq { label, .. } => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfNe { label, .. } => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfLt { label, .. } => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfGe { label, .. } => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfGt { label, .. } => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfLe { label, .. } => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfEqZ { label, .. } => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfNeZ { label, .. } => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfLtZ { label, .. } => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfGeZ { label, .. } => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfGtZ { label, .. } => {
                    used_labels.insert(label.clone());
                }
                Instruction::IfLeZ { label, .. } => {
                    used_labels.insert(label.clone());
                }
                Instruction::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 { 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 { name } => {
                    if !used_labels.contains(name) {
                        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.contains_key(label) {
                debug!("{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 { label } => {
                    last_ins_was_a_label = false;
                    let label = new_labels
                        .get(&label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            label
                        ))?
                        .clone();
                    new_insns.push(Instruction::Goto { label });
                }
                Instruction::IfEq { label, a, b } => {
                    last_ins_was_a_label = false;
                    let label = new_labels
                        .get(&label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfEq { label, a, b });
                }
                Instruction::IfNe { label, a, b } => {
                    last_ins_was_a_label = false;
                    let label = new_labels
                        .get(&label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfNe { label, a, b });
                }
                Instruction::IfLt { label, a, b } => {
                    last_ins_was_a_label = false;
                    let label = new_labels
                        .get(&label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfLt { label, a, b });
                }
                Instruction::IfGe { label, a, b } => {
                    last_ins_was_a_label = false;
                    let label = new_labels
                        .get(&label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfGe { label, a, b });
                }
                Instruction::IfGt { label, a, b } => {
                    last_ins_was_a_label = false;
                    let label = new_labels
                        .get(&label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfGt { label, a, b });
                }
                Instruction::IfLe { label, a, b } => {
                    last_ins_was_a_label = false;
                    let label = new_labels
                        .get(&label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfLe { label, a, b });
                }
                Instruction::IfEqZ { label, a } => {
                    last_ins_was_a_label = false;
                    let label = new_labels
                        .get(&label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfEqZ { label, a });
                }
                Instruction::IfNeZ { label, a } => {
                    last_ins_was_a_label = false;
                    let label = new_labels
                        .get(&label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfNeZ { label, a });
                }
                Instruction::IfLtZ { label, a } => {
                    last_ins_was_a_label = false;
                    let label = new_labels
                        .get(&label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfLtZ { label, a });
                }
                Instruction::IfGeZ { label, a } => {
                    last_ins_was_a_label = false;
                    let label = new_labels
                        .get(&label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfGeZ { label, a });
                }
                Instruction::IfGtZ { label, a } => {
                    last_ins_was_a_label = false;
                    let label = new_labels
                        .get(&label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfGtZ { label, a });
                }
                Instruction::IfLeZ { label, a } => {
                    last_ins_was_a_label = false;
                    let label = new_labels
                        .get(&label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            label
                        ))?
                        .clone();
                    new_insns.push(Instruction::IfLeZ { label, a });
                }
                Instruction::Try {
                    end_label,
                    mut handlers,
                    default_handler,
                } => {
                    last_ins_was_a_label = false;
                    let end_label = new_labels
                        .get(&end_label)
                        .ok_or(anyhow!(
                            "Internal error: {} not found in renamed label",
                            end_label
                        ))?
                        .clone();
                    for handler in &mut handlers {
                        handler.1 = new_labels
                            .get(&handler.1)
                            .ok_or(anyhow!(
                                "Internal error: {} not found in renamed label",
                                handler.1
                            ))?
                            .clone();
                    }
                    let default_handler = default_handler
                        .map(|label| {
                            new_labels
                                .get(&label)
                                .ok_or(anyhow!(
                                    "Internal error: {} not found in renamed label",
                                    label
                                ))
                                .cloned()
                        })
                        .transpose()?;
                    new_insns.push(Instruction::Try {
                        end_label,
                        handlers,
                        default_handler,
                    });
                }
                Instruction::Switch { mut branches, reg } => {
                    last_ins_was_a_label = false;
                    for label in 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 { branches, reg });
                }
                Instruction::Label { name } => {
                    if !used_labels.contains(&name) {
                        //println!("{name} not used");
                        continue;
                    }
                    if !last_ins_was_a_label {
                        new_insns.push(Instruction::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()
        })
    }
}

impl<V: Visitor> Visitable<V> for Code {
    fn default_visit(&self, v: &mut V) -> Result<()> {
        for ins in &self.insns {
            v.visit_instruction(ins)?;
        }
        Ok(())
    }
}
impl<V: VisitorMut> VisitableMut<V> for Code {
    fn default_visit_mut(self, v: &mut V) -> Result<Self> {
        let parameter_names = if let Some(parameter_names) = self.parameter_names {
            let mut new_param = vec![];
            for param in parameter_names {
                if let Some(param) = param {
                    new_param.push(Some(v.visit_string(param)?));
                } else {
                    new_param.push(None);
                }
            }
            Some(new_param)
        } else {
            None
        };
        Ok(Self {
            parameter_names,
            insns: self
                .insns
                .into_iter()
                .map(|ins| v.visit_instruction(ins))
                .collect::<Result<_>>()?,
            ..self
        })
    }
}
