// SPDX-FileCopyrightText: 2023 Thomas Kramer <code@tkramer.ch>
//
// SPDX-License-Identifier: AGPL-3.0-or-later

//! Wrap delay/constraint models and add the capability for keeping track of clock sources
//! which drive a signal.

use std::sync::Arc;

use crate::traits::{ConstraintBase, DelayBase, LoadBase, Signal, TimingBase};

/// Mark the driving clock of a signal.
#[derive(Clone, Debug, PartialEq, Default)]
pub enum ClockId {
    /// No clock ID.
    #[default]
    None,
    /// A single clock ID.
    Single(u16),
    /// Multiple clock IDs.
    Multiple(ArcSet<u16>),
}

impl ClockId {
    /// Keep track of the clock IDs which possibly drive a signal when joining two
    /// signals through a gate.
    fn join(&self, other: &Self) -> Self {
        use ClockId::*;
        match (self, other) {
            (None, x) | (x, None) => x.clone(),
            (Multiple(ids), Single(x)) | (Single(x), Multiple(ids)) => Multiple(ids.add(*x)),
            (Single(id1), Single(id2)) if id1 == id2 => Single(*id1),
            (Single(id1), Single(id2)) => {
                let mut ids = VecSet {
                    elements: vec![*id1, *id2],
                };
                Multiple(ArcSet {
                    elements: ids.into(),
                })
            }
            (Multiple(ids1), Multiple(ids2)) => Multiple(ids1.union(ids2)),
        }
    }
}

/// Signal representation wich can keep track of the clock(s) which launch the signal transition.
#[derive(Debug, Clone)]
pub struct SignalClocked<S> {
    /// Underlying signal representation.
    inner: S,
    /// Clock which drives this signal.
    clock_id: ClockId,
}

impl<S> SignalClocked<S> {
    /// Set the ID of the clock which triggers this signal.
    pub fn with_clock_id(mut self, clock_id: u16) -> Self {
        self.clock_id = ClockId::Single(clock_id);
        self
    }

    /// Get the ID of the clock which launches this signal.
    pub fn clock_id(&self) -> &ClockId {
        &self.clock_id
    }
}

impl<S> Signal for SignalClocked<S>
where
    S: Signal,
{
    type LogicValue = S::LogicValue;

    fn logic_value(&self) -> Self::LogicValue {
        self.inner.logic_value()
    }
}

/// Wrap a delay and/or constraint model and add the capability of tracking
/// clock sources.
pub struct Model<M> {
    /// Underlying delay/constraint model.
    inner: M,
}

impl<M> LoadBase for Model<M>
where
    M: LoadBase,
{
    type Load = M::Load;

    fn sum_loads(&self, load1: &Self::Load, load2: &Self::Load) -> Self::Load {
        self.inner.sum_loads(load1, load2)
    }
}

impl<M> TimingBase for Model<M>
where
    M: TimingBase,
{
    type Signal = SignalClocked<M::Signal>;

    type LogicValue = M::LogicValue;
}

impl<M> DelayBase for Model<M>
where
    M: DelayBase,
{
    type Delay = M::Delay;

    fn summarize_delays(&self, signal1: &Self::Signal, signal2: &Self::Signal) -> Self::Signal {
        let s = self.inner.summarize_delays(&signal1.inner, &signal2.inner);

        SignalClocked {
            inner: s,
            clock_id: signal1.clock_id.join(&signal2.clock_id),
        }
    }

    fn get_delay(&self, from: &Self::Signal, to: &Self::Signal) -> Self::Delay {
        self.inner.get_delay(&from.inner, &to.inner)
    }
}

impl<M> ConstraintBase for Model<M>
where
    M: ConstraintBase,
{
    type Constraint = M::Constraint;

    type RequiredSignal = M::RequiredSignal;

    type Slack = M::Slack;

    fn summarize_constraints(
        &self,
        constraint1: &Self::RequiredSignal,
        constraint2: &Self::RequiredSignal,
    ) -> Self::RequiredSignal {
        self.inner.summarize_constraints(constraint1, constraint2)
    }

    fn solve_delay_constraint(
        &self,
        actual_delay: &Self::Delay,
        required_output: &Self::RequiredSignal,
        actual_signal: &Self::Signal,
    ) -> Self::RequiredSignal {
        self.inner
            .solve_delay_constraint(actual_delay, required_output, &actual_signal.inner)
    }

    fn get_slack(
        &self,
        actual_signal: &Self::Signal,
        required_signal: &Self::RequiredSignal,
    ) -> Self::Slack {
        self.inner.get_slack(&actual_signal.inner, required_signal)
    }
}

/// A simple set implementation based on sorted and deduplicated `Vec`.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
struct VecSet<T> {
    elements: Vec<T>,
}

impl<T> VecSet<T> {
    pub fn new() -> Self {
        Self {
            elements: Default::default(),
        }
    }
}

impl<T> VecSet<T>
where
    T: Ord + Eq + Clone,
{
    pub fn insert(&mut self, item: T) {
        let pos = self.elements.binary_search(&item);
        match pos {
            Ok(_) => {
                // Contains element, nothing to do.
            }
            Err(pos) => self.elements.insert(pos, item),
        }
        //debug_assert!(self.elements.is_sorted());
    }

    pub fn contains(&self, item: &T) -> bool {
        self.elements.binary_search(&item).is_ok()
    }
    pub fn len(&self) -> usize {
        self.elements.len()
    }
    pub fn is_empty(&self) -> bool {
        self.elements.is_empty()
    }
    pub fn contains_set(&self, other: &Self) -> bool {
        // use iter_set::intersection(..).count()
        if other.len() > self.len() {
            // Cannot contain a larger set.
            false
        } else {
            // TODO: more efficient algorithm using the fact that items are sorted.
            // Set is expected to be quite small though. Therefore the naive implementation
            // should not be too bad.
            other.elements.iter().all(|e| self.contains(e))
        }
    }
}

/// A clone-on-write reference counted set.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ArcSet<T> {
    /// Sorted list of unique elements in the set.
    elements: Arc<VecSet<T>>,
}

impl<T> ArcSet<T> {
    pub fn new() -> Self {
        Self {
            elements: VecSet::new().into(),
        }
    }
}

impl<T> ArcSet<T>
where
    T: Ord + Eq + Clone,
{
    /// Create a new set by adding the item.
    /// If the item already exists, creates only a reference and not a full
    /// clone of the set.
    pub fn add(&self, item: T) -> Self {
        let mut new_set = Arc::clone(&self.elements);
        if self.elements.contains(&item) {
            Self { elements: new_set }
        } else {
            let s = Arc::make_mut(&mut new_set);
            s.insert(item);

            Self { elements: new_set }
        }
    }

    /// Create the union of the current set and the other set.
    /// Creates a clone only when necessary.
    pub fn union(&self, other: &Self) -> Self {
        if Arc::ptr_eq(&self.elements, &other.elements) || self == other {
            Self {
                elements: Arc::clone(&self.elements),
            }
        } else {
            // Merge the IDs.
            if self.elements.contains_set(&other.elements) {
                Self {
                    elements: Arc::clone(&self.elements),
                }
            } else if other.elements.contains_set(&self.elements) {
                Self {
                    elements: Arc::clone(&other.elements),
                }
            } else {
                let mut new_set = Arc::clone(&self.elements);
                let s = Arc::make_mut(&mut new_set);
                // Naive
                // TODO: use iter_set crate
                other
                    .elements
                    .elements
                    .iter()
                    .for_each(|e| s.insert(e.clone()));

                Self { elements: new_set }
            }
        }
    }
}

#[test]
fn test_arc_set() {
    let a = ArcSet::new();
    let b = ArcSet::new();

    assert_eq!(a, b);

    let a = a.add(1);
    let b = b.add(2);
    assert_ne!(a, b);

    let a = a.add(2);
    let b = b.add(1);
    assert_eq!(a, b);

    assert!(Arc::ptr_eq(&a.elements, &a.add(1).elements));
}

#[test]
fn test_arc_set_union() {
    let a = ArcSet::new().add(1).add(3);
    let b = ArcSet::new().add(2);

    assert_eq!(a.union(&b), ArcSet::new().add(1).add(2).add(3));

    let a = ArcSet::new().add(1);
    let b = ArcSet::new().add(1);

    assert!(Arc::ptr_eq(&a.elements, &a.union(&b).elements));
}