use crate::analysis::cpa::lattice::JoinSemiLattice;
use crate::analysis::cpa::lattice::pcode::PcodeAddressLattice;
use crate::analysis::pcode_store::{PcodeOpRef, PcodeStore};
use crate::modeling::machine::cpu::concrete::ConcretePcodeAddress;
use jingle_sleigh::PcodeOperation;
use std::borrow::Borrow;
use std::cmp::Ordering;
use std::fmt::{Debug, Display, Formatter, Result as FmtResult};

#[derive(Debug, Clone, Eq, PartialEq)]
pub enum MergeOutcome {
    NoOp,
    /// Merge occurred. Contains a clone of the state before the merge.
    Merged,
}

impl MergeOutcome {
    pub fn is_merged(&self) -> bool {
        matches!(self, MergeOutcome::Merged)
    }

    pub fn is_noop(&self) -> bool {
        matches!(self, MergeOutcome::NoOp)
    }
}

/// Object-safe trait for cloneable iterators stored inside `DynIterBox`.
///
/// Any iterator type that is also `Clone` will implement this via the
/// blanket impl below; the `clone_box` method allows cloning through the
/// trait object so the boxed iterator can be cloned cheaply.
pub trait CloneableIter<'a, T>: Iterator<Item = T> {
    fn clone_box(&self) -> Box<dyn CloneableIter<'a, T> + 'a>;
}

impl<'a, T, I> CloneableIter<'a, T> for I
where
    I: Iterator<Item = T> + Clone + 'a,
{
    fn clone_box(&self) -> Box<dyn CloneableIter<'a, T> + 'a> {
        Box::new(self.clone())
    }
}

/// Newtype wrapper around a boxed cloneable iterator.
///
/// This provides `Clone` and `Iterator` implementations for the boxed trait
/// object, so upstream types (like `Successor::IntoIter`) can be `Clone`,
/// which is required by combinators such as `iproduct`.
pub struct DynIterBox<'a, T>(Box<dyn CloneableIter<'a, T> + 'a>);

impl<'a, T> DynIterBox<'a, T> {
    pub fn new<I>(iter: I) -> Self
    where
        I: Iterator<Item = T> + Clone + 'a,
    {
        DynIterBox(Box::new(iter))
    }
}

impl<'a, T> Clone for DynIterBox<'a, T> {
    fn clone(&self) -> Self {
        DynIterBox(self.0.clone_box())
    }
}

impl<'a, T> Iterator for DynIterBox<'a, T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        self.0.next()
    }

    // Forward size_hint if available
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.0.size_hint()
    }
}

impl<'a, T> Debug for DynIterBox<'a, T> {
    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
        // We can't show the inner iterator state meaningfully, so provide a simple marker.
        f.write_str("DynIterBox { .. }")
    }
}

/// Iterator wrapper returned by `transfer` methods.
///
/// This stores a boxed, cloneable iterator so the `Successor` itself can be
/// `Clone` without forcing collection of items into a `Vec`.
///
/// This structure, [DynIterBox], and [ClonableIter] are my quick best-effort at
/// balancing the following goals:
///
/// * [AbstractState::transfer] returns an iterator.
/// * [AbstractState::transfer] impls can return heterogenous iterator types easily
///   (e.g. one path returns [std::iter::empty] and another returns [std::iter::once])
/// * [AbstractState] Implementations don't have to juggle an associated type or generic
///   arguments for the transfer function.
/// * The iterators produced by [AbstractState::transfer] can be cloned for usage with `iproduct`
///
/// There's definitely some trade-offs with this approach and this might change in the future.
pub struct Successor<'a, T>(DynIterBox<'a, T>);

impl<'a, T> Successor<'a, T> {
    /// Create an empty successor iterator.
    pub fn empty() -> Self
    where
        T: 'a,
    {
        // std::iter::Empty implements Clone, so it satisfies CloneableIter via the blanket impl.
        Self(DynIterBox::new(std::iter::empty::<T>()))
    }
}

impl<'a, T> Clone for Successor<'a, T> {
    fn clone(&self) -> Self {
        Successor(self.0.clone())
    }
}

impl<'a, T: 'a> IntoIterator for Successor<'a, T> {
    type Item = T;
    type IntoIter = DynIterBox<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        self.0
    }
}

/// Construct a `Successor` from any iterator that implements `Clone`.
///
/// We require the iterator to be `Clone` so that the boxed trait object can be
/// cloned cheaply. This avoids collecting into a `Vec`.
impl<'a, T: 'a, I> From<I> for Successor<'a, T>
where
    I: Iterator<Item = T> + Clone + 'a,
{
    fn from(value: I) -> Self {
        Self(DynIterBox::new(value))
    }
}

/// Core trait for abstract states used by the CPA.
pub trait AbstractState: JoinSemiLattice + Clone + Debug + Display {
    /// Merge `other` into `self`. Mutate `self` and return whether a merge occurred.
    /// The mutated `self` MUST be >= than it was before (monotonicity guarantee).
    /// If a merge occurred, returns `Merged`
    fn merge(&mut self, other: &Self) -> MergeOutcome;

    /// Default cartesian merge using `join`.
    fn merge_join(&mut self, new_state: &Self) -> MergeOutcome {
        if self == new_state {
            MergeOutcome::NoOp
        } else {
            self.join(new_state);
            MergeOutcome::Merged
        }
    }

    /// Default separate merge (no-op).
    fn merge_sep(&mut self, _: &Self) -> MergeOutcome {
        MergeOutcome::NoOp
    }

    /// Stop predicate: is `self` covered by any of `states`?
    fn stop<'a, T: Iterator<Item = &'a Self>>(&'a self, states: T) -> bool;

    /// Default stop predicate using piecewise ordering.
    fn stop_sep<'a, T: Iterator<Item = &'a Self>>(&'a self, mut states: T) -> bool {
        states.any(|s| {
            matches!(
                PartialOrd::partial_cmp(self, s),
                Some(Ordering::Less) | Some(Ordering::Equal)
            )
        })
    }

    /// Transfer function: return successor abstract states for a pcode operation.
    ///
    /// The returned `Successor` must be constructed from an iterator that
    /// implements `Clone`. This lets CPA clients clone the successor sequence
    /// cheaply when needed.
    fn transfer<'a, B: Borrow<PcodeOperation>>(&'a self, opcode: B) -> Successor<'a, Self>;
}

/// States that have a program location.
pub trait PcodeLocation {
    fn location(&self) -> PcodeAddressLattice;

    /// Returns the concrete p-code address for this state, if one exists.
    fn concrete_location(&self) -> Option<ConcretePcodeAddress> {
        self.location().value().cloned()
    }
}

/// States that know their program location.
pub trait LocationState: AbstractState + PcodeLocation {
    /// Returns all `(op, next_state)` transitions from this state.
    ///
    /// For forward analysis: `op` is the single op at the current location;
    /// each `next_state` is an abstract successor (one per CFG branch target).
    ///
    /// For backward analysis: each pair is `(op_at_predecessor, predecessor_state)`,
    /// where the abstract transfer has already been applied. The op belongs to
    /// the predecessor (which always has at least one outgoing edge to the current
    /// location, so its op is guaranteed to be present in the store).
    fn get_transitions<'op, T: PcodeStore<'op> + ?Sized>(
        &self,
        store: &'op T,
    ) -> Vec<(PcodeOpRef<'op>, Self)>;
}
