use std::ops::{Index, IndexMut};

use crate::protocol::ComponentState;

/// Generic branch ID. A component will always have one branch: the
/// non-speculative branch. This branch has ID 0. Hence in a speculative context
/// we use this fact to let branch ID 0 denote the ID being invalid.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct BranchId {
    pub index: u32
}

impl BranchId {
    #[inline]
    fn new_invalid() -> Self {
        return Self{ index: 0 };
    }

    #[inline]
    fn new(index: u32) -> Self {
        debug_assert!(index != 0);
        return Self{ index };
    }

    #[inline]
    pub(crate) fn is_valid(&self) -> bool {
        return self.index != 0;
    }
}

#[derive(Debug, PartialEq, Eq)]
pub(crate) enum SpeculativeState {
    // Non-synchronous variants
    RunningNonSync,         // regular execution of code
    Error,                  // encountered a runtime error
    Finished,               // finished executing connector's code
    // Synchronous variants
    RunningInSync,          // running within a sync block
    HaltedAtBranchPoint,    // at a branching point (at a `get` call)
    ReachedSyncEnd,         // reached end of sync block, branch represents a local solution
    Inconsistent,           // branch can never represent a local solution, so halted
}

/// The execution state of a branch. This envelops the PDL code and the
/// execution state. And derived from that: if we're ready to keep running the
/// code, or if we're halted for some reason (e.g. waiting for a message).
pub(crate) struct Branch {
    pub id: BranchId,
    pub parent_id: BranchId,
    // Execution state
    pub code_state: ComponentState,
    pub sync_state: SpeculativeState,
    pub next_in_queue: BranchId, // used by `ExecTree`/`BranchQueue`
}

impl Branch {
    /// Creates a new non-speculative branch
    pub(crate) fn new_non_sync(component_state: ComponentState) -> Self {
        Branch {
            id: BranchId::new_invalid(),
            parent_id: BranchId::new_invalid(),
            code_state: component_state,
            sync_state: SpeculativeState::RunningNonSync,
            next_in_queue: BranchId::new_invalid(),
        }
    }

    /// Constructs a sync branch. The provided branch is assumed to be the
    /// parent of the new branch within the execution tree.
    fn new_sync(new_index: u32, parent_branch: &Branch) -> Self {
        debug_assert!(
            (parent_branch.sync_state == SpeculativeState::RunningNonSync && !parent_branch.parent_index.is_valid()) ||
                (parent_branch.sync_state == SpeculativeState::HaltedAtBranchPoint)
        );
        debug_assert!(parent_branch.prepared_channel.is_none());

        Branch {
            id: BranchId::new(new_index),
            parent_id: parent_branch.index,
            code_state: parent_branch.code_state.clone(),
            sync_state: SpeculativeState::RunningInSync,
            next_in_queue: BranchId::new_invalid(),
        }
    }
}

/// Queue of branches. Just a little helper.
#[derive(Copy, Clone)]
struct BranchQueue {
    first: BranchId,
    last: BranchId,
}

impl BranchQueue {
    #[inline]
    fn new() -> Self {
        Self{
            first: BranchId::new_invalid(),
            last: BranchId::new_invalid()
        }
    }

    #[inline]
    fn is_empty(&self) -> bool {
        debug_assert!(self.first.is_valid() == self.last.is_valid());
        return !self.first.is_valid();
    }
}

const NUM_QUEUES: usize = 3;

pub(crate) enum QueueKind {
    Runnable,
    AwaitingMessage,
    FinishedSync,
}

impl QueueKind {
    fn as_index(&self) -> usize {
        return match self {
            QueueKind::Runnable => 0,
            QueueKind::AwaitingMessage => 1,
            QueueKind::FinishedSync => 2,
        }
    }
}

/// Execution tree of branches. Tries to keep the extra information stored
/// herein to a minimum. So the execution tree is aware of the branches, their
/// execution state and the way they're dependent on each other, but the
/// execution tree should not be aware of e.g. sync algorithms.
///
/// Note that the tree keeps track of multiple lists of branches. Each list
/// contains branches that ended up in a particular execution state. The lists
/// are described by the various `BranchQueue` instances and the `next_in_queue`
/// field in each branch.
pub(crate) struct ExecTree {
    // All branches. the `parent_id` field in each branch implies the shape of
    // the tree. Branches are index stable throughout a sync round.
    pub branches: Vec<Branch>,
    pub queues: [BranchQueue; NUM_QUEUES]
}

impl ExecTree {
    /// Constructs a new execution tree with a single non-sync branch.
    pub fn new(component: ComponentState) -> Self {
        return Self {
            branches: vec![Branch::new_non_sync(component)],
            queues: [BranchQueue::new(); 3]
        }
    }

    // --- Generic branch (queue) management

    /// Returns if tree is in speculative mode
    pub fn is_in_sync(&self) -> bool {
        return self.branches.len() != 1;
    }

    /// Pops a branch (ID) from a queue.
    pub fn pop_from_queue(&mut self, kind: QueueKind) -> Option<BranchId> {
        let queue = &mut self.queues[kind.as_index()];
        if queue.is_empty() {
            return None;
        } else {
            let first_branch = &mut self.branches[queue.first.index as usize];
            queue.first = first_branch.next_in_queue;
            first_branch.next_in_queue = BranchId::new_invalid();
            if !queue.first.is_valid() {
                queue.last = BranchId::new_invalid();
            }

            return Some(first_branch.id);
        }
    }

    /// Pushes a branch (ID) into a queue.
    pub fn push_into_queue(&mut self, kind: QueueKind, id: BranchId) {
        let queue = &mut self.queues[kind.as_index()];
        if queue.is_empty() {
            queue.first = id;
            queue.last = id;
        } else {
            let last_branch = &mut self.branches[queue.last.index as usize];
            last_branch.next_in_queue = id;
            queue.last = id;
        }
    }

    pub fn iter_queue(&self, kind: QueueKind) -> BranchIter {
        let queue = &self.queues[kind.as_index()];
        let index = queue.first as usize;
        return BranchIter{
            branches: self.branches.as_slice(),
            index,
        }
    }

    // --- Preparing and finishing a speculative round

    /// Starts a synchronous round by cloning the non-sync branch and marking it
    /// as the root of the speculative tree.
    pub fn start_sync(&mut self) {
        debug_assert!(!self.is_in_sync());
        let sync_branch = Branch::new_sync(1, &self.branches[0]);
        let sync_branch_id = sync_branch.id;
        self.branches.push(sync_branch);
        self.push_into_queue(QueueKind::Runnable, sync_branch_id);
    }

    /// Creates a new speculative branch based on the provided one. The index to
    /// retrieve this new branch will be returned.
    pub fn fork_branch(&mut self, parent_branch_id: BranchId, initial_queue: Option<QueueKind>) -> BranchId {
        debug_assert!(self.is_in_sync());
        let parent_branch = &self[parent_branch_id];
        let new_branch = Branch::new_sync(1, parent_branch);
        let new_branch_id = new_branch.id;

        if let Some(kind) = initial_queue {
            self.push_into_queue(kind, new_branch_id);
        }

        return new_branch_id;
    }

    /// Collapses the speculative execution tree back into a deterministic one,
    /// using the provided branch as the final sync result.
    pub fn end_sync(&mut self, branch_id: BranchId) {
        debug_assert!(self.is_in_sync());
        debug_assert!(self.iter_queue(QueueKind::FinishedSync).any(|v| v.id == branch_id));

        // Swap indicated branch into the first position
        self.branches.swap(0, branch_id.index as usize);
        self.branches.truncate(1);

        // Reset all values to non-sync defaults
        let branch = &mut self.branches[0];
        branch.id = BranchId::new_invalid();
        branch.parent_id = BranchId::new_invalid();
        branch.sync_state = SpeculativeState::RunningNonSync;
        branch.next_in_queue = BranchId::new_invalid();

        // Clear out all the queues
        for queue_idx in 0..NUM_QUEUES {
            self.queues[queue_idx] = BranchQueue::new();
        }
    }
}

impl Index<BranchId> for ExecTree {
    type Output = Branch;

    fn index(&self, index: BranchId) -> &Self::Output {
        debug_assert!(index.is_valid());
        return &self.branches[index.index as usize];
    }
}

impl IndexMut<BranchId> for ExecTree {
    fn index_mut(&mut self, index: BranchId) -> &mut Self::Output {
        debug_assert!(index.is_valid());
        return &mut self.branches[index.index as usize];
    }
}

pub struct BranchIter<'a> {
    branches: &'a [Branch],
    index: usize,
}

impl<'a> Iterator for BranchIter<'a> {
    type Item = &'a Branch;

    fn next(&mut self) -> Option<Self::Item> {
        if self.index == 0 {
            // i.e. the invalid branch index
            return None;
        }

        let branch = &self.branches[self.index];
        self.index = branch.next_in_queue.index as usize;
        return Some(branch);
    }
}