//! Representation of a class.

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

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

use crate::{
    DexAnnotationItem, DexString, Field, IdField, IdMethod, IdMethodType, IdType, Method,
    MethodHandle, Result, Visitable, VisitableMut, Visitor, VisitorMut,
};
use androscalpel_serializer::consts::*;

/// Represent an apk
#[cfg_attr(feature = "python", pyclass(eq))]
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
pub struct Class {
    /// Type, format described at
    /// <https://source.android.com/docs/core/runtime/dex-format#typedescriptor>
    #[cfg_attr(feature = "python", pyo3(get))]
    pub descriptor: IdType,
    /// If the class is visible everywhere
    #[cfg_attr(feature = "python", pyo3(get))]
    pub is_public: bool,
    /// If the class is subclassable
    #[cfg_attr(feature = "python", pyo3(get))]
    pub is_final: bool,
    /// If the class is a 'multipy-implementable abstract class' AKA an interface
    #[cfg_attr(feature = "python", pyo3(get))]
    pub is_interface: bool,
    /// If the class is instanciable
    #[cfg_attr(feature = "python", pyo3(get))]
    pub is_abstract: bool,
    /// If the class is not directly defined in the source code
    #[cfg_attr(feature = "python", pyo3(get))]
    pub is_synthetic: bool,
    /// If the class is an annotation
    #[cfg_attr(feature = "python", pyo3(get))]
    pub is_annotation: bool,
    /// If the class is an enum
    #[cfg_attr(feature = "python", pyo3(get))]
    pub is_enum: bool,
    /// Name of the superclass, format described at
    /// <https://source.android.com/docs/core/runtime/dex-format#typedescriptor>
    #[cfg_attr(feature = "python", pyo3(get))]
    pub superclass: Option<IdType>,
    /// List of the interfaces that class implement, format of the interfaces
    /// name is discribed at
    /// <https://source.android.com/docs/core/runtime/dex-format#typedescriptor>
    #[cfg_attr(feature = "python", pyo3(get))]
    pub interfaces: Vec<IdType>,
    /// Name of the source file where this class is defined.
    #[cfg_attr(feature = "python", pyo3(get))]
    pub source_file: Option<DexString>,

    /// The static fields
    #[cfg_attr(feature = "python", pyo3(get))]
    pub static_fields: HashMap<IdField, Field>,
    /// The instance fields
    #[cfg_attr(feature = "python", pyo3(get))]
    pub instance_fields: HashMap<IdField, Field>,
    /// The direct (static, private or constructor) methods of the class
    #[cfg_attr(feature = "python", pyo3(get))]
    pub direct_methods: HashMap<IdMethod, Method>,
    /// The virtual (ie non direct) methods of the class
    #[cfg_attr(feature = "python", pyo3(get))]
    pub virtual_methods: HashMap<IdMethod, Method>,
    // Do we need to distinguish direct and virtual (all the other) methods?
    // Maybe overlapping descriptor (same name, class and proto?)
    /// The annotation related to this class (note: this does not include the
    /// methods/field/parameters annotations, they are stored in the methods and fields
    /// structutres)
    #[cfg_attr(feature = "python", pyo3(get))]
    pub annotations: Vec<DexAnnotationItem>,
}

#[cfg_attr(feature = "python", pymethods)]
impl Class {
    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)]
    pub fn new(name: DexString) -> Result<Self> {
        Ok(Self {
            descriptor: IdType::new(name)?,
            superclass: Some(IdType::new("Ljava/lang/Object;".into())?),
            interfaces: vec![],
            source_file: None,
            is_public: true,
            is_final: false,
            is_interface: false,
            is_abstract: false,
            is_synthetic: false,
            is_annotation: false,
            is_enum: false,
            static_fields: HashMap::new(),
            instance_fields: HashMap::new(),
            direct_methods: HashMap::new(),
            virtual_methods: HashMap::new(),
            annotations: vec![],
        })
    }

    pub fn __str__(&self) -> String {
        let name: String = self.descriptor.get_name().__str__();
        let file = if let Some(file) = &self.source_file {
            let file: String = file.__str__();
            format!("  defined in {file}\n")
        } else {
            "".into()
        };
        let superclass = if let Some(spcl) = &self.superclass {
            let spcl: String = spcl.get_name().__str__();
            format!("    extends: {spcl}\n")
        } else {
            "".into()
        };
        let interfaces = if self.interfaces.is_empty() {
            "".into()
        } else {
            let mut interfaces: String = "  implements:\n".into();
            for it in &self.interfaces {
                let it: String = it.get_name().__str__();
                interfaces += &format!("    {it}\n");
            }
            interfaces
        };

        format!("{name}\n{file}{superclass}{interfaces}")
    }

    pub fn __repr__(&self) -> String {
        let name: String = self.descriptor.get_name().__str__();
        format!("Class({name})")
    }

    /// Return all strings referenced in the class.
    pub fn get_all_strings(&self) -> HashSet<DexString> {
        let mut strings = HashSet::new();
        strings.extend(self.descriptor.get_all_strings());
        if let Some(ty) = &self.superclass {
            strings.extend(ty.get_all_strings());
        }
        for ty in &self.interfaces {
            strings.extend(ty.get_all_strings());
        }
        if let Some(string) = &self.source_file {
            strings.insert(string.clone());
        }
        for (id, field) in &self.static_fields {
            strings.extend(id.get_all_strings());
            strings.extend(field.get_all_strings());
        }
        for (id, field) in &self.instance_fields {
            strings.extend(id.get_all_strings());
            strings.extend(field.get_all_strings());
        }
        for (id, method) in &self.direct_methods {
            strings.extend(id.get_all_strings());
            strings.extend(method.get_all_strings());
        }
        for (id, method) in &self.virtual_methods {
            strings.extend(id.get_all_strings());
            strings.extend(method.get_all_strings());
        }
        for annot in &self.annotations {
            strings.extend(annot.get_all_strings());
        }
        strings
    }

    /// Return all types referenced in the class.
    pub fn get_all_types(&self) -> HashSet<IdType> {
        let mut types = HashSet::new();
        types.insert(self.descriptor.clone());
        if let Some(superclass) = &self.superclass {
            types.insert(superclass.clone());
        }
        types.extend(self.interfaces.clone());
        for (id, field) in &self.static_fields {
            types.extend(id.get_all_types());
            types.extend(field.get_all_types());
        }
        for (id, field) in &self.instance_fields {
            types.extend(id.get_all_types());
            types.extend(field.get_all_types());
        }
        for (id, method) in &self.direct_methods {
            types.extend(id.get_all_types());
            types.extend(method.get_all_types());
        }
        for (id, method) in &self.virtual_methods {
            types.extend(id.get_all_types());
            types.extend(method.get_all_types());
        }
        for annot in &self.annotations {
            types.extend(annot.get_all_types());
        }
        types
    }

    /// Return all protoypes referenced in the class.
    pub fn get_all_protos(&self) -> HashSet<IdMethodType> {
        let mut protos = HashSet::new();
        for field in self.static_fields.values() {
            protos.extend(field.get_all_protos());
        }
        for field in self.instance_fields.values() {
            protos.extend(field.get_all_protos());
        }
        for (id, method) in &self.direct_methods {
            protos.extend(id.get_all_protos());
            protos.extend(method.get_all_protos());
        }
        for (id, method) in &self.virtual_methods {
            protos.extend(id.get_all_protos());
            protos.extend(method.get_all_protos());
        }
        for annot in &self.annotations {
            protos.extend(annot.get_all_protos());
        }
        protos
    }

    /// Return all fields id referenced in the class.
    /// This **not** the concatenation of the static and instances fields attributes:
    /// this also contains reference to fields in other classes used by methods/values
    /// in this class.
    pub fn get_all_field_ids(&self) -> HashSet<IdField> {
        let mut fields = HashSet::new();
        for (id, field) in &self.static_fields {
            fields.insert(id.clone());
            fields.extend(field.get_all_field_ids());
        }
        for (id, field) in &self.instance_fields {
            fields.insert(id.clone());
            fields.extend(field.get_all_field_ids());
        }
        for method in self.direct_methods.values() {
            fields.extend(method.get_all_field_ids());
        }
        for method in self.virtual_methods.values() {
            fields.extend(method.get_all_field_ids());
        }
        for annot in &self.annotations {
            fields.extend(annot.get_all_field_ids());
        }
        fields
    }

    /// Return all method ids referenced in the class.
    /// This **not** the concatenation of the direct and virtual method attributs:
    /// this also contains reference to method in other classes used by methods/values
    /// in this class.
    pub fn get_all_method_ids(&self) -> HashSet<IdMethod> {
        let mut methods = HashSet::new();
        for field in self.static_fields.values() {
            methods.extend(field.get_all_method_ids());
        }
        for field in self.instance_fields.values() {
            methods.extend(field.get_all_method_ids());
        }
        for (id, method) in &self.direct_methods {
            methods.insert(id.clone());
            methods.extend(method.get_all_method_ids());
        }
        for (id, method) in &self.virtual_methods {
            methods.insert(id.clone());
            methods.extend(method.get_all_method_ids());
        }
        for annot in &self.annotations {
            methods.extend(annot.get_all_method_ids());
        }
        methods
    }

    /// Return all method handles referenced in the class.
    pub fn get_all_method_handles(&self) -> HashSet<MethodHandle> {
        let mut methods = HashSet::new();
        for field in self.static_fields.values() {
            methods.extend(field.get_all_method_handles());
        }
        for field in self.instance_fields.values() {
            methods.extend(field.get_all_method_handles());
        }
        for method in self.direct_methods.values() {
            methods.extend(method.get_all_method_handles());
        }
        for method in self.virtual_methods.values() {
            methods.extend(method.get_all_method_handles());
        }
        for annot in &self.annotations {
            methods.extend(annot.get_all_method_handles());
        }
        methods
    }

    /// If a data item is needed for this class.
    pub fn has_data_item(&self) -> bool {
        !self.static_fields.is_empty()
            || !self.instance_fields.is_empty()
            || !self.direct_methods.is_empty()
            || !self.virtual_methods.is_empty()
    }

    /// If a static_values array is needed for this class.
    pub fn has_static_values_array(&self) -> bool {
        self.static_fields
            .values()
            .any(|field| field.value.is_some())
    }

    /// Check if the class or its fields/methods have annotations
    pub fn has_annotations(&self) -> bool {
        !self.annotations.is_empty()
            || self
                .static_fields
                .values()
                .any(|field| field.has_annotations())
            || self
                .instance_fields
                .values()
                .any(|field| field.has_annotations())
            || self
                .direct_methods
                .values()
                .any(|method| method.has_annotations())
            || self
                .virtual_methods
                .values()
                .any(|method| method.has_annotations())
    }

    /// Check if the class have a field of method with hiddenapi information
    pub fn has_hiddenapi(&self) -> bool {
        self.static_fields
            .values()
            .any(|field| field.hiddenapi.is_some())
            || self
                .instance_fields
                .values()
                .any(|field| field.hiddenapi.is_some())
            || self
                .direct_methods
                .values()
                .any(|method| method.hiddenapi.is_some())
            || self
                .virtual_methods
                .values()
                .any(|method| method.hiddenapi.is_some())
    }

    /// Return the binary representation of access flags.
    pub fn get_raw_access_flags(&self) -> u32 {
        let mut flags = 0u32;
        if self.is_public {
            flags |= ACC_PUBLIC;
        }
        if self.is_final {
            flags |= ACC_FINAL;
        }
        if self.is_interface {
            flags |= ACC_INTERFACE;
        }
        if self.is_abstract {
            flags |= ACC_ABSTRACT;
        }
        if self.is_synthetic {
            flags |= ACC_SYNTHETIC;
        }
        if self.is_annotation {
            flags |= ACC_ANNOTATION;
        }
        if self.is_enum {
            flags |= ACC_ENUM;
        }
        flags
    }
}

impl<V: Visitor> Visitable<V> for Class {
    fn default_visit(&self, v: &mut V) -> Result<()> {
        v.visit_type(&self.descriptor)?;
        if let Some(superclass) = &self.superclass {
            v.visit_type(superclass)?;
        }
        for interface in &self.interfaces {
            v.visit_type(interface)?;
        }
        if let Some(source_file) = &self.source_file {
            v.visit_string(source_file)?;
        }
        for (id, field) in &self.static_fields {
            v.visit_field_id(id)?;
            v.visit_field(field)?;
        }
        for (id, field) in &self.instance_fields {
            v.visit_field_id(id)?;
            v.visit_field(field)?;
        }
        for (id, meth) in &self.direct_methods {
            v.visit_method_id(id)?;
            v.visit_method(meth)?;
        }
        for (id, meth) in &self.virtual_methods {
            v.visit_method_id(id)?;
            v.visit_method(meth)?;
        }
        for annot in &self.annotations {
            v.visit_annotation_item(annot)?;
        }
        Ok(())
    }
}
impl<V: VisitorMut> VisitableMut<V> for Class {
    fn default_visit_mut(self, v: &mut V) -> Result<Self> {
        Ok(Self {
            descriptor: v.visit_type(self.descriptor)?,
            superclass: self
                .superclass
                .map(|superclass| v.visit_type(superclass))
                .transpose()?,
            interfaces: self
                .interfaces
                .into_iter()
                .map(|interface| v.visit_type(interface))
                .collect::<Result<_>>()?,
            source_file: self
                .source_file
                .map(|source_file| v.visit_string(source_file))
                .transpose()?,
            static_fields: {
                let mut static_fields = HashMap::new();
                for (id, field) in self.static_fields.into_iter() {
                    static_fields.insert(v.visit_field_id(id)?, v.visit_field(field)?);
                }
                static_fields
            },
            instance_fields: {
                let mut instance_fields = HashMap::new();
                for (id, field) in self.instance_fields.into_iter() {
                    instance_fields.insert(v.visit_field_id(id)?, v.visit_field(field)?);
                }
                instance_fields
            },
            direct_methods: {
                let mut direct_methods = HashMap::new();
                for (id, meth) in self.direct_methods.into_iter() {
                    direct_methods.insert(v.visit_method_id(id)?, v.visit_method(meth)?);
                }
                direct_methods
            },
            virtual_methods: {
                let mut virtual_methods = HashMap::new();
                for (id, meth) in self.virtual_methods.into_iter() {
                    virtual_methods.insert(v.visit_method_id(id)?, v.visit_method(meth)?);
                }
                virtual_methods
            },
            annotations: self
                .annotations
                .into_iter()
                .map(|annotation| v.visit_annotation_item(annotation))
                .collect::<Result<_>>()?,
            ..self
        })
    }
}
