//! Annotations (for class, fields, methods and parameters alike).
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};

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

use crate::hashmap_vectorize;
use crate::{
    dex_id::IdType, value::DexValue, DexString, IdField, IdMethod, IdMethodType, MethodHandle,
    Result, Visitable, VisitableMut, Visitor, VisitorMut,
};

/// Annotation with a visibility
#[cfg_attr(feature = "python", pyclass(eq))]
#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
pub struct DexAnnotationItem {
    // TODO: the get/set will probably be wonky on the python side when edditing
    /// The actual annotation
    #[cfg_attr(feature = "python", pyo3(get))]
    pub annotation: DexAnnotation,
    // TODO: enforce exclusivity
    /// If the annotation visibility is set to build
    #[cfg_attr(feature = "python", pyo3(get))]
    pub visibility_build: bool,
    /// If the annotation visibility is set to runtime
    #[cfg_attr(feature = "python", pyo3(get))]
    pub visibility_runtime: bool,
    /// If the annotation visibility is set to system
    #[cfg_attr(feature = "python", pyo3(get))]
    pub visibility_system: bool,
}

#[cfg_attr(feature = "python", pymethods)]
impl DexAnnotationItem {
    #[cfg_attr(feature = "python", new)]
    pub fn new(annotation: DexAnnotation) -> Self {
        Self {
            annotation,
            visibility_build: false,
            visibility_runtime: false,
            visibility_system: false,
        }
    }

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

    pub fn __repr__(&self) -> String {
        let visibility =
            if !(self.visibility_build || self.visibility_runtime || self.visibility_system) {
                "none".into()
            } else {
                let mut visibility = String::new();
                if self.visibility_build {
                    visibility += "build";
                }
                if self.visibility_runtime {
                    visibility += "runtime";
                }
                if self.visibility_system {
                    visibility += "system";
                }
                visibility
            };
        let annotation = self.annotation.__repr__();
        format!("AnnotationItem(visibility: {visibility}, {annotation})")
    }

    /// Return all strings referenced in the annotation.
    pub fn get_all_strings(&self) -> HashSet<DexString> {
        self.annotation.get_all_strings()
    }

    /// Return all types referenced in the annotation.
    pub fn get_all_types(&self) -> HashSet<IdType> {
        self.annotation.get_all_types()
    }

    /// Return all prototypes referenced in the annotation.
    pub fn get_all_protos(&self) -> HashSet<IdMethodType> {
        self.annotation.get_all_protos()
    }

    /// Return all field ids referenced in the annotation.
    pub fn get_all_field_ids(&self) -> HashSet<IdField> {
        self.annotation.get_all_field_ids()
    }

    /// Return all method ids referenced in the annotation.
    pub fn get_all_method_ids(&self) -> HashSet<IdMethod> {
        self.annotation.get_all_method_ids()
    }

    /// Return all method handles referenced in the annotation.
    pub fn get_all_method_handles(&self) -> HashSet<MethodHandle> {
        self.annotation.get_all_method_handles()
    }

    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)?)
    }
}

impl<V: Visitor> Visitable<V> for DexAnnotationItem {
    fn default_visit(&self, v: &mut V) -> Result<()> {
        v.visit_annotation(&self.annotation)
    }
}
impl<V: VisitorMut> VisitableMut<V> for DexAnnotationItem {
    fn default_visit_mut(self, v: &mut V) -> Result<Self> {
        Ok(Self {
            annotation: v.visit_annotation(self.annotation)?,
            ..self
        })
    }
}

/// An annotation.
#[cfg_attr(feature = "python", pyclass(eq))]
#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
pub struct DexAnnotation {
    // TODO: check the relation between type and encoded_value.
    /// The type of the annotation.
    #[cfg_attr(feature = "python", pyo3(get))]
    pub type_: IdType,
    // TODO: the get/set will probably be wonky on the python side when edditing
    /// The annotation elements.
    #[cfg_attr(feature = "python", pyo3(get))]
    #[serde(with = "hashmap_vectorize")]
    pub elements: HashMap<DexString, DexValue>, // TODO: check MemberName syntax?
}

#[cfg_attr(feature = "python", pymethods)]
impl DexAnnotation {
    #[cfg_attr(feature = "python", new)]
    pub fn new(type_: IdType, elements: HashMap<DexString, DexValue>) -> Self {
        Self { type_, elements }
    }

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

    pub fn __repr__(&self) -> String {
        let type_ = self.type_.__str__();
        let mut elts: String = "{".into();
        for (key, val) in self.elements.iter() {
            elts += &format!("{}: {}, ", key.__str__(), val.__str__());
        }
        elts += "}";
        format!("Annotation({type_}, {elts})")
    }

    /// Return all strings referenced in the annotation.
    pub fn get_all_strings(&self) -> HashSet<DexString> {
        let mut strings = HashSet::new();
        strings.extend(self.type_.get_all_strings());
        for (name, value) in &self.elements {
            strings.insert(name.clone());
            strings.extend(value.get_all_strings());
        }
        strings
    }

    /// Return all types referenced in the annotation.
    pub fn get_all_types(&self) -> HashSet<IdType> {
        let mut types = HashSet::new();
        types.insert(self.type_.clone());
        for value in self.elements.values() {
            types.extend(value.get_all_types());
        }
        types
    }

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

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

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

    /// Return all method handles referenced in the annotation.
    pub fn get_all_method_handles(&self) -> HashSet<MethodHandle> {
        let mut methods = HashSet::new();
        for value in self.elements.values() {
            methods.extend(value.get_all_method_handles());
        }
        methods
    }

    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)?)
    }
}

impl<V: Visitor> Visitable<V> for DexAnnotation {
    fn default_visit(&self, v: &mut V) -> Result<()> {
        v.visit_type(&self.type_)?;
        for (id, val) in &self.elements {
            v.visit_string(id)?;
            v.visit_value(val)?;
        }
        Ok(())
    }
}

impl<V: VisitorMut> VisitableMut<V> for DexAnnotation {
    fn default_visit_mut(self, v: &mut V) -> Result<Self> {
        let type_ = v.visit_type(self.type_)?;
        let mut elements = HashMap::new();
        for (id, val) in self.elements {
            let id = v.visit_string(id)?;
            let val = v.visit_value(val)?;
            elements.insert(id, val);
        }
        Ok(Self { type_, elements })
    }
}
