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

use crate::runtime2::scheduler::*;

use crate::runtime2::runtime::*;

use crate::runtime2::communication::*;

use super::component_context::*;

pub struct PortAnnotation {

    self_comp_id: CompId,

    self_port_id: PortId,

    peer_comp_id: CompId, // only valid for getter ports

    peer_port_id: PortId, // only valid for getter ports

    peer_discovered: bool, // only valid for getter ports

    mapping: Option<u32>,

    kind: PortKind,

}

impl PortAnnotation {

    fn new(comp_id: CompId, port_id: PortId, kind: PortKind) -> Self {

        return Self{

            self_comp_id: comp_id,

            self_port_id: port_id,

            peer_comp_id: CompId::new_invalid(),

            peer_port_id: PortId::new_invalid(),

            peer_discovered: false,

            mapping: None,

            kind,

        }

    }

}

#[derive(Debug, Eq, PartialEq)]

enum Mode {

    NonSync,

    SyncBusy,

    SyncAwaitingSolution,

    SelectBusy,

    SelectWait,

}

struct SolutionCombiner {

    solution: SyncPartialSolution,

    matched_channels: usize,

}

impl SolutionCombiner {

    fn new() -> Self {

        return Self {

            solution: SyncPartialSolution::default(),

            matched_channels: 0,

        }

    }

    #[inline]

    fn has_contributions(&self) -> bool {

        return !self.solution.channel_mapping.is_empty();

    }

    /// Returns a decision for the current round. If there is no decision (yet)

    /// then `RoundDecision::None` is returned.

    fn get_decision(&self) -> SyncRoundDecision {

        if self.matched_channels == self.solution.channel_mapping.len() {

            debug_assert_ne!(self.solution.decision, SyncRoundDecision::None);

            return self.solution.decision;

        }

        return SyncRoundDecision::None; // even in case of failure: wait for everyone.

    }

    fn combine_with_partial_solution(&mut self, partial: SyncPartialSolution) {

        debug_assert_ne!(self.solution.decision, SyncRoundDecision::Solution);

        debug_assert_ne!(partial.decision, SyncRoundDecision::Solution);

        if partial.decision == SyncRoundDecision::Failure {

            self.solution.decision = SyncRoundDecision::Failure;

        }

        for entry in partial.channel_mapping {

            let channel_index = if entry.getter.is_some() && entry.putter.is_some() {

                let channel_index = self.solution.channel_mapping.len();

                self.solution.channel_mapping.push(entry);

                channel_index

            } else if let Some(putter) = entry.putter {

                self.combine_with_putter_port(putter)

            } else if let Some(getter) = entry.getter {

                self.combine_with_getter_port(getter)

            } else {

                unreachable!(); // both putter and getter are None

            };

            let channel = &self.solution.channel_mapping[channel_index];

            if let Some(consistent) = Self::channel_is_consistent(channel) {

                if !consistent {

                    self.solution.decision = SyncRoundDecision::Failure;

                }

                self.matched_channels += 1;

            }

        }

        self.update_solution();

    }

    /// Combines the currently stored global solution (if any) with the newly

    /// provided local solution. Make sure to check the `has_decision` return

    /// value afterwards.

    fn combine_with_local_solution(&mut self, _comp_id: CompId, solution: SyncLocalSolution) {

        debug_assert_ne!(self.solution.decision, SyncRoundDecision::Solution);

        // Combine partial solution with the local solution entries

        for entry in solution {

            // Match the current entry up with its peer endpoint, or add a new

            // entry.

            let channel_index = match entry {

                SyncLocalSolutionEntry::Putter(putter) => {

                    self.combine_with_putter_port(putter)

                },

                SyncLocalSolutionEntry::Getter(getter) => {

                    self.combine_with_getter_port(getter)

                }

            };

            // Check if channel is now consistent

            let channel = &self.solution.channel_mapping[channel_index];

            if let Some(consistent) = Self::channel_is_consistent(channel) {

                if !consistent {

                    self.solution.decision = SyncRoundDecision::Failure;

                }

                self.matched_channels += 1;

            }

        }

        self.update_solution();

    }

    /// Takes whatever partial solution is present in the solution combiner and

    /// returns it. The solution combiner's solution will end up being empty.

    /// This is used when a new leader is found and we need to pass along our

    /// partial results.

    fn take_partial_solution(&mut self) -> SyncPartialSolution {

        let mut partial_solution = SyncPartialSolution::default();

        std::mem::swap(&mut partial_solution, &mut self.solution);

        self.clear();

        return partial_solution;

    }

    fn clear(&mut self) {

        self.solution.channel_mapping.clear();

        self.solution.decision = SyncRoundDecision::None;

        self.matched_channels = 0;

    }

    // --- Small utilities for combining solutions

    fn combine_with_putter_port(&mut self, putter: SyncSolutionPutterPort) -> usize {

        let channel_index = self.get_channel_index_for_putter(putter.self_comp_id, putter.self_port_id);

        if let Some(channel_index) = channel_index {

            let channel = &mut self.solution.channel_mapping[channel_index];

            debug_assert!(channel.putter.is_none());

            channel.putter = Some(putter);

            return channel_index;

        } else {

            let channel_index = self.solution.channel_mapping.len();

            self.solution.channel_mapping.push(SyncSolutionChannel{

                putter: Some(putter),

                getter: None,

            });

            return channel_index;

        }

    }

    fn combine_with_getter_port(&mut self, getter: SyncSolutionGetterPort) -> usize {

        let channel_index = self.get_channel_index_for_getter(getter.peer_comp_id, getter.peer_port_id);

        if let Some(channel_index) = channel_index {

                        self_port_id: port_info.self_id,

                        mapping

                    }),

                    PortKind::Getter => SyncLocalSolutionEntry::Getter(SyncSolutionGetterPort{

                        self_comp_id: comp_ctx.id,

                        self_port_id: port_info.self_id,

                        peer_comp_id: port.peer_comp_id,

                        peer_port_id: port.peer_port_id,

                        mapping

                    })

                };

                local_solution.push(new_entry);

            }

        }

        let decision = self.handle_local_solution(sched_ctx, comp_ctx, comp_ctx.id, local_solution);

        return decision;

    }

    /// Notifies that a decision has been reached. Note that the caller should

    /// still take the appropriate actions based on the decision it is supplying

    /// to the consensus layer.

    pub(crate) fn notify_sync_decision(&mut self, _decision: SyncRoundDecision) {

        // Reset everything for the next round

        debug_assert_eq!(self.mode, Mode::SyncAwaitingSolution);

        self.mode = Mode::NonSync;

        self.round_index = self.round_index.wrapping_add(1);

        for port in self.ports.iter_mut() {

            port.mapping = None;

        }

        self.solution.clear();

    }

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

    // Handling inbound and outbound messages

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

    /// Prepares a set of values to be sent of a channel.

    pub(crate) fn annotate_data_message(&mut self, comp_ctx: &CompCtx, port_info: &Port, content: ValueGroup) -> DataMessage {

        debug_assert_eq!(self.mode, Mode::SyncBusy); // can only send between sync start and sync end

        debug_assert!(self.ports.iter().any(|v| v.self_port_id == port_info.self_id));

        let data_header = self.create_data_header_and_update_mapping(port_info);

        let sync_header = self.create_sync_header(comp_ctx);

        return DataMessage{ data_header, sync_header, content };

    }

    /// Handles the arrival of a new data message (needs to be called for every

    /// new data message, even though it might not end up being received). This

    /// is used to determine peers of `get`ter ports.

    pub(crate) fn handle_new_data_message(&mut self, comp_ctx: &CompCtx, message: &DataMessage) {

        let target_handle = comp_ctx.get_port_handle(message.data_header.target_port);

        let target_index = comp_ctx.get_port_index(target_handle);

        let annotation = &mut self.ports[target_index];

        debug_assert!(

            !annotation.peer_discovered || (

                annotation.peer_comp_id == message.sync_header.sending_id &&

                annotation.peer_port_id == message.data_header.source_port

            )

        );

        annotation.peer_comp_id = message.sync_header.sending_id;

        annotation.peer_port_id = message.data_header.source_port;

        annotation.peer_discovered = true;

    }

    /// Checks if the data message can be received (due to port annotations), if

    /// it can then `true` is returned and the caller is responsible for handing

    /// the message of to the PDL code. Otherwise the message cannot be

    /// received.

    pub(crate) fn try_receive_data_message(&mut self, sched_ctx: &SchedulerCtx, comp_ctx: &mut CompCtx, message: &DataMessage) -> bool {

        debug_assert_eq!(self.mode, Mode::SyncBusy);

        debug_assert!(self.ports.iter().any(|v| v.self_port_id == message.data_header.target_port));

        // Make sure the expected mapping matches the currently stored mapping

        for (peer_port_kind, expected_annotation) in &message.data_header.expected_mapping {

            // Determine our annotation, in order to do so we need to find the

            // port matching the peer ports

            let mut self_annotation = None;

            let mut self_annotation_found = false;

            match peer_port_kind {

                PortAnnotationKind::Putter(peer_port) => {

                    for self_port in &self.ports {

                        if self_port.kind == PortKind::Getter &&

                            self_port.peer_discovered &&

                            self_port.peer_comp_id == peer_port.self_comp_id &&

                            self_port.peer_port_id == peer_port.self_port_id

                        {

                            self_annotation = self_port.mapping;

                            self_annotation_found = true;

                            break;

                        }

                    }

                },

                PortAnnotationKind::Getter(peer_port) => {

                        // Peer indicates that we talked to it

                        let self_port_handle = comp_ctx.get_port_handle(peer_port.peer_port_id);

                        let self_port_index = comp_ctx.get_port_index(self_port_handle);

                        self_annotation = self.ports[self_port_index].mapping;

                        self_annotation_found = true;

                    }

                }

            }

            if !self_annotation_found {

                continue

            }

            if self_annotation != *expected_annotation {

                return false;

            }

        }

        // Expected mapping matches current mapping, so we will receive the message

        self.set_annotation(message.sync_header.sending_id, &message.data_header);

        // Handle the sync header embedded within the data message

        self.handle_sync_header(sched_ctx, comp_ctx, &message.sync_header);

        return true;

    }

    /// Receives the sync message and updates the consensus state appropriately.

    pub(crate) fn receive_sync_message(&mut self, sched_ctx: &SchedulerCtx, comp_ctx: &mut CompCtx, message: SyncMessage) -> SyncRoundDecision {

        // Whatever happens: handle the sync header (possibly changing the

        // currently registered leader)

        self.handle_sync_header(sched_ctx, comp_ctx, &message.sync_header);

        match message.content {

            SyncMessageContent::NotificationOfLeader => {

                return SyncRoundDecision::None;

            },

            SyncMessageContent::LocalSolution(solution_generator_id, local_solution) => {

                return self.handle_local_solution(sched_ctx, comp_ctx, solution_generator_id, local_solution);

            },

            SyncMessageContent::PartialSolution(partial_solution) => {

                return self.handle_partial_solution(sched_ctx, comp_ctx, partial_solution);

            },

            SyncMessageContent::GlobalSolution => {

                debug_assert_eq!(self.mode, Mode::SyncAwaitingSolution); // leader can only find global- if we submitted local solution

                return SyncRoundDecision::Solution;

            },

            SyncMessageContent::GlobalFailure => {

                debug_assert_eq!(self.mode, Mode::SyncAwaitingSolution);

                return SyncRoundDecision::Failure;

            }

        }

    }

    fn handle_sync_header(&mut self, sched_ctx: &SchedulerCtx, comp_ctx: &mut CompCtx, header: &MessageSyncHeader) {

        if header.highest_id.0 > self.highest_id.0 {

            // Sender knows of someone with a higher ID. So store highest ID,

            // notify all peers, and forward local solutions

            self.highest_id = header.highest_id;

            for peer in comp_ctx.iter_peers() {

                if peer.id == header.sending_id {

                    continue; // do not send to sender: it has the higher ID

                }

                // also: only send if we received a message in this round

                let mut performed_communication = false; // TODO: Revise, temporary fix

                for port in self.ports.iter() {

                    if port.peer_comp_id == peer.id && port.mapping.is_some() {

                        performed_communication = true;

                        break;

                    }

                }

                if !performed_communication {

                    continue;

                }

                let message = SyncMessage{

                    sync_header: self.create_sync_header(comp_ctx),

                    content: SyncMessageContent::NotificationOfLeader,

                };

                peer.handle.send_message(sched_ctx, Message::Sync(message), true);

            }

            self.forward_partial_solution(sched_ctx, comp_ctx);

        } else if header.highest_id.0 < self.highest_id.0 {

            // Sender has a lower ID, so notify it of our higher one

            let message = SyncMessage{

                sync_header: self.create_sync_header(comp_ctx),

                content: SyncMessageContent::NotificationOfLeader,

            };

            let peer_handle = comp_ctx.get_peer_handle(header.sending_id);

            let peer_info = comp_ctx.get_peer(peer_handle);

            peer_info.handle.send_message(sched_ctx, Message::Sync(message), true);

        } // else: exactly equal

    }

        s32 a = 5;

        auto b = 5 + a;

    }

    ").expect("compilation");

    let rt = Runtime::new(1, true, pd);

    for _i in 0..20 {

        create_component(&rt, "", "nothing_at_all", no_args());

    }

}

#[test]

fn test_component_communication() {

    let pd = ProtocolDescription::parse(b"

    primitive sender(out<u32> o, u32 outside_loops, u32 inside_loops) {

        u32 outside_index = 0;

        while (outside_index < outside_loops) {

            u32 inside_index = 0;

            sync while (inside_index < inside_loops) {

                put(o, inside_index);

                inside_index += 1;

            }

            outside_index += 1;

        }

    }

    primitive receiver(in<u32> i, u32 outside_loops, u32 inside_loops) {

        u32 outside_index = 0;

        while (outside_index < outside_loops) {

            u32 inside_index = 0;

            sync while (inside_index < inside_loops) {

                auto val = get(i);

                while (val != inside_index) {} // infinite loop if incorrect value is received

                inside_index += 1;

            }

            outside_index += 1;

        }

    }

    composite constructor() {

        channel o_orom -> i_orom;

        channel o_mrom -> i_mrom;

        channel o_ormm -> i_ormm;

        channel o_mrmm -> i_mrmm;

        // one round, one message per round

        new sender(o_orom, 1, 1);

        new receiver(i_orom, 1, 1);

        // multiple rounds, one message per round

        new sender(o_mrom, 5, 1);

        new receiver(i_mrom, 5, 1);

        // one round, multiple messages per round

        new sender(o_ormm, 1, 5);

        new receiver(i_ormm, 1, 5);

        // multiple rounds, multiple messages per round

        new sender(o_mrmm, 5, 5);

        new receiver(i_mrmm, 5, 5);

    }").expect("compilation");

    let rt = Runtime::new(3, true, pd);

    create_component(&rt, "", "constructor", no_args());

}

#[test]

fn test_intermediate_messenger() {

    let pd = ProtocolDescription::parse(b"

    primitive receiver<T>(in<T> rx, u32 num) {

        auto index = 0;

        while (index < num) {

            sync { auto v = get(rx); }

            index += 1;

        }

    }

    primitive middleman<T>(in<T> rx, out<T> tx, u32 num) {

        auto index = 0;

        while (index < num) {

            sync { put(tx, get(rx)); }

            index += 1;

        }

    }

    primitive sender<T>(out<T> tx, u32 num) {

        auto index = 0;

        while (index < num) {

            sync put(tx, 1337);

            index += 1;

        }

    }

    composite constructor_template<T>() {

        auto num = 0;

        channel<T> tx_a -> rx_a;

        channel tx_b -> rx_b;

        new receiver(rx_b, 3);

    }

    composite constructor() {

        new constructor_template<u16>();

    }

    ").expect("compilation");

    create_component(&rt, "", "constructor", no_args());

}

#[test]

fn test_simple_select() {

    let pd = ProtocolDescription::parse(b"

    func infinite_assert<T>(T val, T expected) -> () {

        while (val != expected) { print(\"nope!\"); }

        return ();

    }

    primitive receiver(in<u32> in_a, in<u32> in_b, u32 num_sends) {

        auto num_from_a = 0;

        auto num_from_b = 0;

        while (num_from_a + num_from_b < 2 * num_sends) {

            sync select {

                auto v = get(in_a) -> {

                    print(\"got something from A\");

                    auto _ = infinite_assert(v, num_from_a);

                    num_from_a += 1;

                }

                auto v = get(in_b) -> {

                    print(\"got something from B\");

                    auto _ = infinite_assert(v, num_from_b);

                    num_from_b += 1;

                }

            }

        }

    }

    primitive sender(out<u32> tx, u32 num_sends) {

        auto index = 0;

        while (index < num_sends) {

            sync {

                put(tx, index);

                index += 1;

            }

        }

    }

    composite constructor() {

        auto num_sends = 15;

        channel tx_a -> rx_a;

        channel tx_b -> rx_b;

        new sender(tx_a, num_sends);

        new receiver(rx_a, rx_b, num_sends);

        new sender(tx_b, num_sends);

    }

    ").expect("compilation");

    let rt = Runtime::new(3, false, pd);

    create_component(&rt, "", "constructor", no_args());

}

#[test]

fn test_unguarded_select() {

    let pd = ProtocolDescription::parse(b"

    primitive constructor_outside_select() {

        u32 index = 0;

        while (index < 5) {

            sync select { auto v = () -> print(\"hello\"); }

            index += 1;

        }

    }

    primitive constructor_inside_select() {

        u32 index = 0;

        while (index < 5) {

            sync select { auto v = () -> index += 1; }

        }

    }

    ").expect("compilation");

    let rt = Runtime::new(3, false, pd);

    create_component(&rt, "", "constructor_outside_select", no_args());

    create_component(&rt, "", "constructor_inside_select", no_args());

}

#[test]

fn test_empty_select() {

    let pd = ProtocolDescription::parse(b"

    primitive constructor() {

        u32 index = 0;

        while (index < 5) {

            sync select {}

            index += 1;

        }

    }

    ").expect("compilation");

    let rt = Runtime::new(3, false, pd);

    create_component(&rt, "", "constructor", no_args());

}