use crate::collections::VecSet;

use crate::protocol::eval::ValueGroup;

use super::ConnectorId;

use super::branch::BranchId;

use super::inbox::{

    Message, DataHeader, SyncHeader, ChannelAnnotation, BranchMarker,

    DataMessage,

    SyncCompMessage, SyncCompContent,

    SyncPortMessage, SyncPortContent,

    SyncControlMessage, SyncControlContent

};

use super::scheduler::{ComponentCtx, ComponentPortChange};

struct BranchAnnotation {

    channel_mapping: Vec<ChannelAnnotation>,

    cur_marker: BranchMarker,

}

#[derive(Debug)]

pub(crate) struct LocalSolution {

    component: ConnectorId,

    final_branch_id: BranchId,

    port_mapping: Vec<(ChannelId, BranchMarker)>,

}

#[derive(Debug, Clone)]

pub(crate) struct GlobalSolution {

    component_branches: Vec<(ConnectorId, BranchId)>,

    Success(BranchId),

}

// -----------------------------------------------------------------------------

// Consensus

// -----------------------------------------------------------------------------

struct Peer {

    id: ConnectorId,

    encountered_this_round: bool,

    expected_sync_round: u32,

}

/// The consensus algorithm. Currently only implemented to find the component

/// with the highest ID within the sync region and letting it handle all the

/// local solutions.

///

/// The type itself serves as an experiment to see how code should be organized.

// TODO: Flatten all datastructures

// TODO: Have a "branch+port position hint" in case multiple operations are

//  performed on the same port to prevent repeated lookups

// TODO: A lot of stuff should be batched. Like checking all the sync headers

//  and sending "I have a higher ID" messages. Should reduce locking by quite a

//  bit.

pub(crate) struct Consensus {

    // --- State that is cleared after each round

    // Local component's state

    highest_connector_id: ConnectorId,

    branch_annotations: Vec<BranchAnnotation>, // index is branch ID

    branch_markers: Vec<BranchId>, // index is branch marker, maps to branch

    // Gathered state from communication

    encountered_ports: VecSet<PortIdLocal>, // to determine if we should send "port remains silent" messages.

    solution_combiner: SolutionCombiner,

    handled_wave: bool, // encountered notification wave in this round

    conclusion: Option<RoundConclusion>,

    ack_remaining: u32,

    // --- Persistent state

    peers: Vec<Peer>,

    sync_round: u32,

    // --- Workspaces

    workspace_ports: Vec<PortIdLocal>,

}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]

pub(crate) enum Consistency {

    Valid,

    Inconsistent,

}

        debug_assert!(self.branch_annotations.len() == new_branch_id.index as usize);

        let parent_branch_annotations = &self.branch_annotations[parent_branch_id.index as usize];

        let new_marker = BranchMarker::new(self.branch_markers.len() as u32);

        let new_branch_annotations = BranchAnnotation{

            channel_mapping: parent_branch_annotations.channel_mapping.clone(),

            cur_marker: new_marker,

        };

        self.branch_annotations.push(new_branch_annotations);

        self.branch_markers.push(new_branch_id);

    }

    /// Notifies the consensus algorithm that a particular branch has

    /// encountered an unrecoverable error.

    pub fn notify_of_fatal_branch(&mut self, failed_branch_id: BranchId, ctx: &mut ComponentCtx) -> Option<RoundConclusion> {

        debug_assert!(self.is_in_sync());

        // Check for trivial case, where branch has not yet communicated within

        // the consensus algorithm

        let branch = &self.branch_annotations[failed_branch_id.index as usize];

        if branch.channel_mapping.iter().all(|v| v.registered_id.is_none()) {

            println!("DEBUG: Failure everything silent");

            return Some(RoundConclusion::Failure);

        }

    }

    /// Notifies the consensus algorithm that a branch has reached the end of

    /// the sync block. A final check for consistency will be performed that the

    /// caller has to handle. Note that

    pub fn notify_of_finished_branch(&self, branch_id: BranchId) -> Consistency {

        debug_assert!(self.is_in_sync());

        let branch = &self.branch_annotations[branch_id.index as usize];

        for mapping in &branch.channel_mapping {

            match mapping.expected_firing {

                Some(expected) => {

                    if expected != mapping.registered_id.is_some() {

                        // Inconsistent speculative state and actual state

                        debug_assert!(mapping.registered_id.is_none()); // because if we did fire on a silent port, we should've caught that earlier

                        return Consistency::Inconsistent;

                    }

                },

                None => {},

            }

        }

        return Consistency::Valid;

    }

    /// Prepares a message for sending. Caller should have made sure that

    /// sending the message is consistent with the speculative state.

    pub fn handle_message_to_send(&mut self, branch_id: BranchId, source_port_id: PortIdLocal, content: &ValueGroup, ctx: &mut ComponentCtx) -> (SyncHeader, DataHeader) {

        debug_assert!(self.is_in_sync());

        let branch = &mut self.branch_annotations[branch_id.index as usize];

        let port_info = ctx.get_port_by_id(source_port_id).unwrap();

        if cfg!(debug_assertions) {

            // Check for consistent mapping

            let port = branch.channel_mapping.iter()

                .find(|v| v.channel_id == port_info.channel_id)

                .unwrap();

            debug_assert!(port.expected_firing == None || port.expected_firing == Some(true));

        }

        // Check for ports that are being sent

        debug_assert!(self.workspace_ports.is_empty());

        find_ports_in_value_group(content, &mut self.workspace_ports);

        if !self.workspace_ports.is_empty() {

            todo!("handle sending ports");

            self.workspace_ports.clear();

        }

        // Construct data header

        let data_header = DataHeader{

            expected_mapping: branch.channel_mapping.iter()

                .filter(|v| v.registered_id.is_some() || v.channel_id == port_info.channel_id)

                .copied()

                .collect(),

            sending_port: port_info.self_id,

            target_port: port_info.peer_id,

            new_mapping: branch.cur_marker,

        };

        // Update port mapping

        for mapping in &mut branch.channel_mapping {

            if mapping.channel_id == port_info.channel_id {

                mapping.expected_firing = Some(true);

                mapping.registered_id = Some(branch.cur_marker);

            }

        }

        // Update branch marker

        let new_marker = BranchMarker::new(self.branch_markers.len() as u32);

        branch.cur_marker = new_marker;

        self.branch_markers.push(branch_id);

        self.encountered_ports.push(source_port_id);

                    let port_info = ctx.get_port_by_channel_id(mapping.channel_id).unwrap();

                    channels.push(LocalChannelPresence{

                        channel_id: mapping.channel_id,

                        is_closed: port_info.state == PortState::Closed,

                    });

                }

                let maybe_conlusion = self.send_to_leader_or_handle_as_leader(SyncCompContent::Presence(ComponentPresence{

                    component_id: ctx.id,

                    channels,

                }), ctx);

                return maybe_conlusion;

            }

        }

    }

    pub fn handle_new_sync_control_message(&mut self, message: SyncControlMessage, ctx: &mut ComponentCtx) -> Option<RoundConclusion> {

        if message.in_response_to_sync_round < self.sync_round {

            // Old message

            return None

        }

        match message.content {

            SyncControlContent::ChannelIsClosed(_) => {

            }

        }

    }

    pub fn notify_of_received_message(&mut self, branch_id: BranchId, message: &DataMessage, ctx: &ComponentCtx) {

        debug_assert!(self.branch_can_receive(branch_id, message));

        let target_port = ctx.get_port_by_id(message.data_header.target_port).unwrap();

        let branch = &mut self.branch_annotations[branch_id.index as usize];

        for mapping in &mut branch.channel_mapping {

            if mapping.channel_id == target_port.channel_id {

                // Found the port in which the message should be inserted

                mapping.registered_id = Some(message.data_header.new_mapping);

                // Check for sent ports

                debug_assert!(self.workspace_ports.is_empty());

                find_ports_in_value_group(&message.content, &mut self.workspace_ports);

                if !self.workspace_ports.is_empty() {

                    todo!("handle received ports");

                    self.workspace_ports.clear();

                }

                return;

            }

                if owner_b != ctx.id {

                    if encountered.push(owner_b) {

                        let message = SyncCompMessage{

                            sync_header: self.create_sync_header(ctx),

                            target_component_id: owner_b,

                            content: SyncCompContent::GlobalFailure,

                        };

                        ctx.submit_message(Message::SyncComp(message)).unwrap();

                    }

                }

            }

        }

        println!("DEBUERINO: Leader entering error state, we need to wait on {:?}", encountered.iter().map(|v| v.0).collect::<Vec<_>>());

        self.conclusion = Some(RoundConclusion::Failure);

        if encountered.is_empty() {

            // We don't have to wait on Acks

            return Some(RoundConclusion::Failure);

        } else {

            self.ack_remaining = encountered.len() as u32;

            return None;

        }

    }

    #[inline]

    fn create_sync_header(&self, ctx: &ComponentCtx) -> SyncHeader {

        return SyncHeader{

            sending_component_id: ctx.id,

            highest_component_id: self.highest_connector_id,

            sync_round: self.sync_round,

        }

    }

    fn forward_local_data_to_new_leader(&mut self, ctx: &mut ComponentCtx) {

        debug_assert_ne!(self.highest_connector_id, ctx.id);

        if let Some(partial_solution) = self.solution_combiner.drain() {

            let message = SyncCompMessage {

                sync_header: self.create_sync_header(ctx),

                target_component_id: self.highest_connector_id,

                content: SyncCompContent::PartialSolution(partial_solution),

            };

            ctx.submit_message(Message::SyncComp(message)).unwrap(); // unwrap: sending to component instead of through channel

        }

    }

}

// -----------------------------------------------------------------------------

    fn combine(&mut self, combiner: SolutionCombiner) -> Option<LeaderConclusion> {

        self.failure_reported = self.failure_reported || combiner.failure_reported;

        // Handle local solutions

        if self.local.is_empty() {

            // Trivial case

            self.local = combiner.local;

        } else {

            for local in combiner.local {

                for matched in local.solutions {

                    let local_solution = LocalSolution{

                        component: local.component,

                        final_branch_id: matched.final_branch_id,

                        port_mapping: matched.channel_mapping,

                    };

                    let maybe_solution = self.add_solution_and_check_for_global_solution(local_solution);

                    if let Some(global_solution) = maybe_solution {

                        return Some(LeaderConclusion::Solution(global_solution));

                    }

                }

            }

        }

        // Handle channel presence

        if self.presence.is_empty() {

            // Trivial case

        } else {

            for presence in combiner.presence {

                match self.presence.iter_mut().find(|v| v.id == presence.id) {

                    Some(entry) => {

                        // Combine entries. Take first that has Closed, then

                        // check first that has both, then check if they are

                        // combinable

                        if entry.state == PresenceState::Closed {

                            // Do nothing

                        } else if presence.state == PresenceState::Closed {

                            entry.owner_a = presence.owner_a;

                            entry.owner_b = presence.owner_b;

                            entry.state = PresenceState::Closed;

                        } else if entry.state == PresenceState::BothPresent {

                            // Again: do nothing

                        } else if presence.state == PresenceState::BothPresent {

                            entry.owner_a = presence.owner_a;

                            entry.owner_b = presence.owner_b;

                            entry.state = PresenceState::BothPresent;

                        } else {

                            // Both have one presence, combine into both present

                            debug_assert!(entry.state == PresenceState::OnePresent && presence.state == PresenceState::OnePresent);

                            entry.owner_b = Some(presence.owner_a);

                        }

                    },

                    None => {

                        self.presence.push(presence);

                    }

                }

            }

            // After adding everything we might have immediately found a solution

            if self.check_for_global_failure() {

                return Some(LeaderConclusion::Failure);

            }

        }

        return None;

    }

    fn clear(&mut self) {

        self.local.clear();

        self.presence.clear();

        self.failure_reported = false;

    }

}

// -----------------------------------------------------------------------------

// Generic Helpers

// -----------------------------------------------------------------------------

/// Recursively goes through the value group, attempting to find ports.

/// Duplicates will only be added once.

pub(crate) fn find_ports_in_value_group(value_group: &ValueGroup, ports: &mut Vec<PortIdLocal>) {

    // Helper to check a value for a port and recurse if needed.

    use crate::protocol::eval::Value;

                        // Transfer messages associated with the transferred port

                        let mut message_idx = 0;

                        while message_idx < scheduled.ctx.inbox_messages.len() {

                            let message = &scheduled.ctx.inbox_messages[message_idx];

                            if Self::get_message_target_port(message) == Some(port_id) {

                                // Need to transfer this message

                                // TODO: Revise messages, this is becoming messy and error-prone

                                let message = scheduled.ctx.inbox_messages.remove(message_idx);

                                if message_idx < scheduled.ctx.inbox_len_read {

                                    scheduled.ctx.inbox_len_read -= 1;

                                }

                                new_connector.ctx.inbox_messages.push(message);

                            } else {

                                message_idx += 1;

                            }

                        }

                        // Transfer the port itself

                        let port_index = scheduled.ctx.ports.iter()

                            .position(|v| v.self_id == port_id)

                            .unwrap();

                        let port = scheduled.ctx.ports.remove(port_index);

                        new_connector.ctx.ports.push(port.clone());

                    }

                    // Schedule new connector to run

                    self.runtime.push_work(new_key);

                },

                ComponentStateChange::CreatedPort(port) => {

                    scheduled.ctx.ports.push(port);

                },

                ComponentStateChange::ChangedPort(port_change) => {

                    if port_change.is_acquired {

                        scheduled.ctx.ports.push(port_change.port);

                    } else {

                        let index = scheduled.ctx.ports

                            .iter()

                            .position(|v| v.self_id == port_change.port.self_id)

                            .unwrap();

                        scheduled.ctx.ports.remove(index);

                    }

                }

            }

        }

        // Finally, check if we just entered or just left a sync region

        if scheduled.ctx.changed_in_sync {

mod network_shapes;

mod api_component;

mod speculation;

mod data_transmission;

mod sync_failure;

use super::*;

use crate::{PortId, ProtocolDescription};

use crate::common::Id;

use crate::protocol::eval::*;

use crate::runtime2::native::{ApplicationSyncAction};

// Generic testing constants, use when appropriate to simplify stress-testing

// pub(crate) const NUM_THREADS: u32 = 3;     // number of threads in runtime

// pub(crate) const NUM_INSTANCES: u32 = 7;   // number of test instances constructed

// pub(crate) const NUM_LOOPS: u32 = 8;       // number of loops within a single test (not used by all tests)

fn create_runtime(pdl: &str) -> Runtime {

    let protocol = ProtocolDescription::parse(pdl.as_bytes()).expect("parse pdl");

    let runtime = Runtime::new(NUM_THREADS, protocol);

    return runtime;

}

fn run_test_in_runtime<F: Fn(&mut ApplicationInterface)>(pdl: &str, constructor: F) {

    let protocol = ProtocolDescription::parse(pdl.as_bytes())

        .expect("parse PDL");

    let runtime = Runtime::new(NUM_THREADS, protocol);

    let mut api = runtime.create_interface();

    for _ in 0..NUM_INSTANCES {

        constructor(&mut api);

    }

}

pub(crate) struct TestTimer {

    name: &'static str,

    started: std::time::Instant

}

    primitive immediate_failure_inside_sync() {

        u32[] only_allows_index_0 = { 1 };

        while (true) sync { // note the infinite loop

            auto value = only_allows_index_0[1];

        }

    }

    primitive immediate_failure_outside_sync() {

        u32[] only_allows_index_0 = { 1 };

        auto never_gonna_get = only_allows_index_0[1];

        while (true) sync {}

    }

    ";

    // let thing = TestTimer::new("local_sync_failure");

    run_test_in_runtime(CODE, |api| {

        api.create_connector("", "immediate_failure_outside_sync", ValueGroup::new_stack(Vec::new()))

            .expect("create component");

        api.create_connector("", "immediate_failure_inside_sync", ValueGroup::new_stack(Vec::new()))

            .expect("create component");

    })

}

        }

    }

        for variant in 0..4 { // all `Location` enum variants, except `Never`.

            // Create the channels

                Value::Enum(variant)

            ])).expect("create connector");

        }

    })

}