use super::{
    AbstractDomain, AbstractIdentifier, HasBitSize, HasTop, PointerDomain, RegisterDomain,
};
use crate::bil::*;
use crate::prelude::*;
use std::collections::{BTreeMap, BTreeSet};
use std::fmt::Display;

/// An abstract domain representing either a pointer or a non-pointer value.
/// Both non-pointer values and offsets of pointers are represented by the same abstract domain `T`.
#[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Clone)]
pub enum DataDomain<T: RegisterDomain> {
    Top(BitSize),
    Pointer(PointerDomain<T>),
    Value(T),
}

impl<T: RegisterDomain> DataDomain<T> {
    /// For pointer values replace an abstract identifier with another one and add the offset_adjustment to the pointer offset.
    /// This is needed to adjust stack pointer on call and return instructions.
    pub fn replace_abstract_id(
        &mut self,
        old_id: &AbstractIdentifier,
        new_id: &AbstractIdentifier,
        offset_adjustment: &T,
    ) {
        if let Self::Pointer(pointer) = self {
            pointer.replace_abstract_id(old_id, new_id, offset_adjustment);
        }
    }

    /// Return a set of all referenced abstract IDs. The set is empty if `self` is not a pointer.
    pub fn referenced_ids(&self) -> BTreeSet<AbstractIdentifier> {
        if let Self::Pointer(pointer) = self {
            pointer.ids().cloned().collect()
        } else {
            BTreeSet::new()
        }
    }

    /// If *self* is a pointer, remove all provided IDs from the target list of it.
    /// If this would leave the pointer without any targets, replace it with *Top*.
    pub fn remove_ids(&mut self, ids_to_remove: &BTreeSet<AbstractIdentifier>) {
        if let Self::Pointer(pointer) = self {
            let remaining_targets: BTreeMap<AbstractIdentifier, T> = pointer
                .targets()
                .iter()
                .filter_map(|(id, offset)| {
                    if ids_to_remove.get(id).is_none() {
                        Some((id.clone(), offset.clone()))
                    } else {
                        None
                    }
                })
                .collect();
            if remaining_targets.is_empty() {
                *self = Self::new_top(self.bitsize());
            } else {
                *self = Self::Pointer(PointerDomain::with_targets(remaining_targets));
            }
        }
    }
}

impl<T: RegisterDomain> HasBitSize for DataDomain<T> {
    // Return the bitsize of `self`.
    fn bitsize(&self) -> BitSize {
        use DataDomain::*;
        match self {
            Top(size) => *size,
            Pointer(pointer) => pointer.bitsize(),
            Value(bitvec) => bitvec.bitsize(),
        }
    }
}

impl<T: RegisterDomain> HasTop for DataDomain<T> {
    // Generate a new *Top* element with the same bitsize as `self`.
    fn top(&self) -> Self {
        DataDomain::new_top(self.bitsize())
    }
}

impl<T: RegisterDomain> RegisterDomain for DataDomain<T> {
    // Return a new *Top* element with the given bitsize
    fn new_top(bitsize: BitSize) -> Self {
        Self::Top(bitsize)
    }

    /// Compute the (abstract) result of a binary operation
    fn bin_op(&self, op: BinOpType, rhs: &Self) -> Self {
        use BinOpType::*;
        use DataDomain::*;
        match (self, op, rhs) {
            (Value(left), _, Value(right)) => Value(left.bin_op(op, right)),
            (Pointer(pointer), PLUS, Value(value)) | (Value(value), PLUS, Pointer(pointer)) => {
                Pointer(pointer.add_to_offset(value))
            }
            (Pointer(pointer), MINUS, Value(value)) => Pointer(pointer.sub_from_offset(value)),
            (Pointer(pointer_lhs), MINUS, Pointer(pointer_rhs)) => {
                if pointer_lhs.ids().len() == 1 && pointer_rhs.ids().len() == 1 {
                    let (id_lhs, offset_lhs) = pointer_lhs.targets().iter().next().unwrap();
                    let (id_rhs, offset_rhs) = pointer_rhs.targets().iter().next().unwrap();
                    if id_lhs == id_rhs {
                        Self::Value(offset_lhs.bin_op(MINUS, offset_rhs))
                    } else {
                        Self::Top(self.bitsize())
                    }
                } else {
                    // We cannot be sure that both pointers point to the same target
                    Self::Top(self.bitsize())
                }
            }
            (_, EQ, _) | (_, NEQ, _) | (_, LT, _) | (_, LE, _) | (_, SLT, _) | (_, SLE, _) => {
                T::new_top(1).into()
            }
            (_, PLUS, _)
            | (_, MINUS, _)
            | (_, TIMES, _)
            | (_, DIVIDE, _)
            | (_, SDIVIDE, _)
            | (_, MOD, _)
            | (_, SMOD, _)
            | (_, LSHIFT, _)
            | (_, RSHIFT, _)
            | (_, ARSHIFT, _)
            | (_, AND, _)
            | (_, OR, _)
            | (_, XOR, _) => Self::new_top(self.bitsize()),
        }
    }

    /// Compute the (abstract) result of a unary operation
    fn un_op(&self, op: UnOpType) -> Self {
        if let Self::Value(value) = self {
            Self::Value(value.un_op(op))
        } else {
            Self::new_top(self.bitsize())
        }
    }

    /// extract a sub-bitvector
    fn extract(&self, low_bit: BitSize, high_bit: BitSize) -> Self {
        if let Self::Value(value) = self {
            Self::Value(value.extract(low_bit, high_bit))
        } else {
            Self::new_top(high_bit - low_bit + 1)
        }
    }

    /// Extend or shrink a bitvector using the given cast type
    fn cast(&self, kind: CastType, width: BitSize) -> Self {
        if self.bitsize() == width {
            // The cast is a no-op.
            return self.clone();
        }
        if let Self::Value(value) = self {
            Self::Value(value.cast(kind, width))
        } else {
            // The result of extending or shrinking pointers is undefined.
            Self::new_top(width)
        }
    }

    /// Concatenate two bitvectors.
    fn concat(&self, other: &Self) -> Self {
        if let (Self::Value(upper_bits), Self::Value(lower_bits)) = (self, other) {
            Self::Value(upper_bits.concat(lower_bits))
        } else {
            Self::new_top(self.bitsize() + other.bitsize())
        }
    }
}

impl<T: RegisterDomain> AbstractDomain for DataDomain<T> {
    // Merge `self` with `other`.
    fn merge(&self, other: &Self) -> Self {
        use DataDomain::*;
        match (self, other) {
            (Top(bitsize), _) | (_, Top(bitsize)) => Top(*bitsize),
            (Pointer(pointer1), Pointer(pointer2)) => Pointer(pointer1.merge(pointer2)),
            (Value(val1), Value(val2)) => Value(val1.merge(val2)),
            (Pointer(_), Value(_)) | (Value(_), Pointer(_)) => Top(self.bitsize()),
        }
    }

    /// Return whether the element represents a top element or not.
    fn is_top(&self) -> bool {
        matches!(self, Self::Top(_))
    }
}

impl<T: RegisterDomain> From<PointerDomain<T>> for DataDomain<T> {
    fn from(val: PointerDomain<T>) -> Self {
        Self::Pointer(val)
    }
}

impl<T: RegisterDomain> From<T> for DataDomain<T> {
    fn from(value: T) -> Self {
        Self::Value(value)
    }
}

impl<T: RegisterDomain + From<Bitvector>> From<Bitvector> for DataDomain<T> {
    fn from(bitvector: Bitvector) -> Self {
        Self::Value(bitvector.into())
    }
}

impl<T: RegisterDomain + Display> DataDomain<T> {
    /// Get a more compact json-representation of the data domain.
    /// Intended for pretty printing, not useable for serialization/deserialization.
    pub fn to_json_compact(&self) -> serde_json::Value {
        match self {
            Self::Top(bitsize) => serde_json::Value::String(format!("Top:{}", bitsize)),
            Self::Pointer(pointer) => {
                let target_iter = pointer.targets().iter().map(|(id, offset)| {
                    (
                        format!("{}", id),
                        serde_json::Value::String(format!("{}", offset)),
                    )
                });
                let targets = serde_json::Value::Object(target_iter.collect());
                let mut obj_map = serde_json::Map::new();
                obj_map.insert("Pointer".to_string(), targets);
                serde_json::Value::Object(obj_map)
            }
            Self::Value(bitvector) => serde_json::Value::String(format!("Value: {}", bitvector)),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::super::*;
    use super::*;

    type Data = DataDomain<BitvectorDomain>;

    fn bv(value: i64) -> BitvectorDomain {
        BitvectorDomain::Value(Bitvector::from_i64(value))
    }

    fn new_id(name: &str) -> AbstractIdentifier {
        AbstractIdentifier::new(
            Tid::new("time0"),
            AbstractLocation::Register(name.into(), 64),
        )
    }

    fn new_pointer_domain(location: &str, offset: i64) -> PointerDomain<BitvectorDomain> {
        let id = new_id(location);
        PointerDomain::new(id, bv(offset))
    }

    fn new_pointer(location: &str, offset: i64) -> Data {
        Data::Pointer(new_pointer_domain(location, offset))
    }

    fn new_value(value: i64) -> Data {
        Data::Value(bv(value))
    }

    #[test]
    fn data_abstract_domain() {
        let pointer = new_pointer("Rax".into(), 0);
        let data = new_value(42);
        assert_eq!(pointer.merge(&pointer), pointer);
        assert_eq!(pointer.merge(&data), Data::new_top(64));
        assert_eq!(
            data.merge(&new_value(41)),
            Data::Value(BitvectorDomain::new_top(64))
        );

        let other_pointer = new_pointer("Rbx".into(), 0);
        match pointer.merge(&other_pointer) {
            Data::Pointer(_) => (),
            _ => panic!(),
        }
    }

    #[test]
    fn data_register_domain() {
        use crate::bil::BinOpType::*;
        let data = new_value(42);
        assert_eq!(data.bitsize(), 64);

        let three = new_value(3);
        let pointer = new_pointer("Rax".into(), 0);
        assert_eq!(data.bin_op(PLUS, &three), new_value(45));
        assert_eq!(pointer.bin_op(PLUS, &three), new_pointer("Rax".into(), 3));
        assert_eq!(three.un_op(crate::bil::UnOpType::NEG), new_value(-3));

        assert_eq!(
            three.extract(0, 31),
            Data::Value(BitvectorDomain::Value(Bitvector::from_i32(3)))
        );

        assert_eq!(data.cast(crate::bil::CastType::SIGNED, 128).bitsize(), 128);

        let one = Data::Value(BitvectorDomain::Value(Bitvector::from_i32(1)));
        let two = Data::Value(BitvectorDomain::Value(Bitvector::from_i32(2)));
        let concat = new_value((1 << 32) + 2);
        assert_eq!(one.concat(&two), concat);
    }

    #[test]
    fn remove_ids() {
        let mut targets = BTreeMap::new();
        targets.insert(new_id("Rax"), bv(1));
        targets.insert(new_id("Rbx"), bv(2));
        let mut data: Data = PointerDomain::with_targets(targets).into();

        let mut ids_to_remove = BTreeSet::new();
        ids_to_remove.insert(new_id("Rbx"));
        ids_to_remove.insert(new_id("Rcx"));

        data.remove_ids(&ids_to_remove);
        assert_eq!(
            data.referenced_ids(),
            vec![new_id("Rax")].into_iter().collect()
        );

        data = bv(42).into();
        data.remove_ids(&ids_to_remove);
        assert_eq!(data, bv(42).into());
    }
}
