use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::atomic::Ordering;
use crate::runtime2::inbox2::ControlContent;

use super::{ScheduledConnector, RuntimeInner, ConnectorId, ConnectorKey, ConnectorVariant};
use super::port::{Port, PortState, PortIdLocal};
use super::native::Connector;
use super::branch::{BranchId};
use super::connector2::{ConnectorPDL, ConnectorScheduling};
use super::inbox2::{MessageFancy, DataMessageFancy, SyncMessageFancy, ControlMessageFancy};

// Because it contains pointers we're going to do a copy by value on this one
#[derive(Clone, Copy)]
pub(crate) struct SchedulerCtx<'a> {
    pub(crate) runtime: &'a RuntimeInner
}

pub(crate) struct Scheduler {
    runtime: Arc<RuntimeInner>,
    scheduler_id: u32,
}

impl Scheduler {
    pub fn new(runtime: Arc<RuntimeInner>, scheduler_id: u32) -> Self {
        return Self{ runtime, scheduler_id };
    }

    pub fn run(&mut self) {
        // Setup global storage and workspaces that are reused for every
        // connector that we run
        'thread_loop: loop {
            // Retrieve a unit of work
            self.debug("Waiting for work");
            let connector_key = self.runtime.wait_for_work();
            if connector_key.is_none() {
                // We should exit
                self.debug(" ... No more work, quitting");
                break 'thread_loop;
            }

            // We have something to do
            let connector_key = connector_key.unwrap();
            let connector_id = connector_key.downcast();
            self.debug_conn(connector_id, &format!(" ... Got work, running {}", connector_key.index));

            let scheduled = self.runtime.get_component_private(&connector_key);

            // Keep running until we should no longer immediately schedule the
            // connector.
            let mut cur_schedule = ConnectorScheduling::Immediate;
            while cur_schedule == ConnectorScheduling::Immediate {
                self.handle_inbox_messages(scheduled);

                // Run the main behaviour of the connector, depending on its
                // current state.
                if scheduled.shutting_down {
                    // Nothing to do. But we're stil waiting for all our pending
                    // control messages to be answered.
                    self.debug_conn(connector_id, &format!("Shutting down, {} Acks remaining", scheduled.router.num_pending_acks()));
                    if scheduled.router.num_pending_acks() == 0 {
                        // We're actually done, we can safely destroy the
                        // currently running connector
                        self.runtime.destroy_component(connector_key);
                        continue 'thread_loop;
                    } else {
                        cur_schedule = ConnectorScheduling::NotNow;
                    }
                } else {
                    self.debug_conn(connector_id, "Running ...");
                    let scheduler_ctx = SchedulerCtx{ runtime: &*self.runtime };
                    let new_schedule = scheduled.connector.run(scheduler_ctx, &mut scheduled.ctx_fancy);
                    self.debug_conn(connector_id, "Finished running");

                    // Handle all of the output from the current run: messages to
                    // send and connectors to instantiate.
                    self.handle_changes_in_context(scheduled);

                    cur_schedule = new_schedule;
                }
            }

            // If here then the connector does not require immediate execution.
            // So enqueue it if requested, and otherwise put it in a sleeping
            // state.
            match cur_schedule {
                ConnectorScheduling::Immediate => unreachable!(),
                ConnectorScheduling::Later => {
                    // Simply queue it again later
                    self.runtime.push_work(connector_key);
                },
                ConnectorScheduling::NotNow => {
                    // Need to sleep, note that we are the only ones which are
                    // allows to set the sleeping state to `true`, and since
                    // we're running it must currently be `false`.
                    self.try_go_to_sleep(connector_key, scheduled);
                },
                ConnectorScheduling::Exit => {
                    // Prepare for exit. Set the shutdown flag and broadcast
                    // messages to notify peers of closing channels
                    scheduled.shutting_down = true;
                    for port in &scheduled.ctx_fancy.ports {
                        if port.state != PortState::Closed {
                            let message = scheduled.router.prepare_closing_channel(
                                port.self_id, port.peer_id,
                                connector_id
                            );
                            self.debug_conn(connector_id, &format!("Sending message [ exit ] \n --- {:?}", message));
                            self.runtime.send_message(port.peer_connector, MessageFancy::Control(message));
                        }
                    }

                    if scheduled.router.num_pending_acks() == 0 {
                        self.runtime.destroy_component(connector_key);
                        continue 'thread_loop;
                    }

                    self.try_go_to_sleep(connector_key, scheduled);
                }
            }
        }
    }

    /// Receiving messages from the public inbox and handling them or storing
    /// them in the component's private inbox
    fn handle_inbox_messages(&mut self, scheduled: &mut ScheduledConnector) {
        let connector_id = scheduled.ctx_fancy.id;

        while let Some(message) = scheduled.public.inbox.take_message() {
            // Check for rerouting
            self.debug_conn(connector_id, &format!("Handling message from conn({}) at port({})\n --- {:?}", message.sending_connector.0, message.receiving_port.index, message));
            if let Some(other_connector_id) = scheduled.router.should_reroute(message.sending_connector, message.receiving_port) {
                self.debug_conn(connector_id, &format!(" ... Rerouting to connector {}", other_connector_id.0));
                self.runtime.send_message(other_connector_id, message);
                continue;
            }

            // Handle special messages here, messages for the component
            // will be added to the inbox.
            self.debug_conn(connector_id, " ... Handling message myself");
            match message {
                MessageFancy::Control(message) => {
                    match message.content {
                        ControlContent::PortPeerChanged(port_id, new_target_connector_id) => {
                            // Need to change port target
                            let port = scheduled.ctx_fancy.get_port_mut_by_id(port_id).unwrap();
                            port.peer_connector = new_target_connector_id;

                            // Note: for simplicity we program the scheduler to always finish
                            // running a connector with an empty outbox. If this ever changes
                            // then accepting the "port peer changed" message implies we need
                            // to change the recipient of the message in the outbox.
                            debug_assert!(scheduled.ctx_fancy.outbox.is_empty());

                            // And respond with an Ack
                            let ack_message = MessageFancy::Control(ControlMessageFancy{
                                id: content.id,
                                sending_component_id: connector_id,
                                content: ControlContent::Ack,
                            });
                            self.debug_conn(connector_id, &format!("Sending message [pp ack]\n --- {:?}", ack_message));
                            self.runtime.send_message(message.sending_component_id, ack_message);
                        },
                        ControlContent::CloseChannel(port_id) => {
                            // Mark the port as being closed
                            let port = scheduled.ctx_fancy.get_port_mut_by_id(port_id).unwrap();
                            port.state = PortState::Closed;

                            // Send an Ack
                            let ack_message = MessageFancy::Control(ControlMessageFancy{
                                id: content.id,
                                sending_component_id: connector_id,
                                content: ControlContent::Ack,
                            });
                            self.debug_conn(connector_id, &format!("Sending message [cc ack] \n --- {:?}", ack_message));
                            self.runtime.send_message(message.sending_component_id, ack_message);
                        },
                        ControlContent::Ack => {
                            scheduled.router.handle_ack(content.id);
                        },
                        ControlContent::Ping => {},
                    }
                },
                _ => {
                    // All other cases have to be handled by the component
                    scheduled.ctx_fancy.inbox_messages.push(message);
                }
            }
        }
    }

    /// Handles changes to the context that were made by the component. This is
    /// the way (due to Rust's borrowing rules) that we bubble up changes in the
    /// component's state that the scheduler needs to know about (e.g. a message
    /// that the component wants to send, a port that has been added).
    fn handle_changes_in_context(&mut self, scheduled: &mut ScheduledConnector) {
        let connector_id = scheduled.ctx_fancy.id;

        // Handling any messages that were sent
        while let Some(mut message) = scheduled.ctx_fancy.outbox.pop_front() {
            self.debug_conn(connector_id, &format!("Sending message [outbox] \n --- {:?}", message));

            let target_component_id = match &message {
                MessageFancy::Data(content) => {
                    // Data messages are always sent to a particular port, and
                    // may end up being rerouted.
                    let port_desc = scheduled.ctx_fancy.get_port_by_id(content.data_header.sending_port).unwrap();
                    debug_assert_eq!(port_desc.peer_id, content.data_header.target_port);

                    if port_desc.state == PortState::Closed {
                        todo!("handle sending over a closed port")
                    }

                    port_desc.peer_connector
                },
                MessageFancy::Sync(content) => {
                    // Sync messages are always sent to a particular component,
                    // the sender must make sure it actually wants to send to
                    // the specified component (and is not using an inconsistent
                    // component ID associated with a port).
                    content.sync_header.highest_component_id
                },
                MessageFancy::Control(_) => {
                    unreachable!("component sending control messages directly");
                }
            };

            self.runtime.send_message(target_component_id, message);
        }

        while let Some(state_change) = scheduled.ctx_fancy.state_changes.pop_front() {
            match state_change {
                ComponentStateChange::CreatedComponent(component, initial_ports) => {
                    // Add the new connector to the global registry
                    let new_key = self.runtime.create_pdl_component(component, false);
                    let new_connector = self.runtime.get_component_private(&new_key);

                    // Transfer ports
                    for port_id in initial_ports {
                        // Transfer messages associated with the transferred port
                        let mut message_idx = 0;
                        while message_idx < scheduled.ctx_fancy.inbox_messages.len() {
                            let message = &scheduled.ctx_fancy.inbox_messages[message_idx];
                            if message.receiving_port == *port_id {
                                // Need to transfer this message
                                let taken_message = scheduled.ctx_fancy.inbox_messages.remove(message_idx);
                                new_connector.ctx_fancy.inbox_messages.push(taken_message);
                            } else {
                                message_idx += 1;
                            }
                        }

                        // Transfer the port itself
                        let port_index = scheduled.ctx_fancy.ports.iter()
                            .position(|v| v.self_id == *port_id)
                            .unwrap();
                        let port = scheduled.ctx_fancy.ports.remove(port_index);
                        new_connector.ctx_fancy.ports.push(port.clone());

                        // Notify the peer that the port has changed
                        let reroute_message = scheduled.router.prepare_reroute(
                            port.self_id, port.peer_id, scheduled.ctx_fancy.id,
                            port.peer_connector, new_connector.ctx_fancy.id
                        );

                        self.debug_conn(connector_id, &format!("Sending message [newcon]\n --- {:?}", reroute_message));
                        self.runtime.send_message(port.peer_connector, MessageFancy::Control(reroute_message));
                    }

                    // Schedule new connector to run
                    self.runtime.push_work(new_key);
                },
                ComponentStateChange::CreatedPort(port) => {
                    scheduled.ctx_fancy.ports.push(port);
                },
                ComponentStateChange::ChangedPort(port_change) => {
                    if port_change.is_acquired {
                        scheduled.ctx_fancy.ports.push(port_change.port);
                    } else {
                        let index = scheduled.ctx_fancy.ports
                            .iter()
                            .position(|v| v.self_id == port_change.port.self_id)
                            .unwrap();
                        scheduled.ctx_fancy.ports.remove(index);
                    }
                }
            }
        }

        // Finally, check if we just entered or just left a sync region
        if scheduled.ctx_fancy.changed_in_sync {
            if scheduled.ctx_fancy.is_in_sync {
                // Just entered sync region
            } else {
                // Just left sync region. So clear inbox
                scheduled.ctx_fancy.inbox_messages.clear();
                scheduled.ctx_fancy.inbox_len_read = 0;
            }

            scheduled.ctx_fancy.changed_in_sync = false; // reset flag
        }
    }

    fn try_go_to_sleep(&self, connector_key: ConnectorKey, connector: &mut ScheduledConnector) {
        debug_assert_eq!(connector_key.index, connector.ctx_fancy.id.0);
        debug_assert_eq!(connector.public.sleeping.load(Ordering::Acquire), false);

        // This is the running connector, and only the running connector may
        // decide it wants to sleep again.
        connector.public.sleeping.store(true, Ordering::Release);

        // But due to reordering we might have received messages from peers who
        // did not consider us sleeping. If so, then we wake ourselves again.
        if !connector.public.inbox.is_empty() {
            // Try to wake ourselves up (needed because someone might be trying
            // the exact same atomic compare-and-swap at this point in time)
            let should_wake_up_again = connector.public.sleeping
                .compare_exchange(true, false, Ordering::SeqCst, Ordering::Acquire)
                .is_ok();

            if should_wake_up_again {
                self.runtime.push_work(connector_key)
            }
        }
    }

    // TODO: Remove, this is debugging stuff
    fn debug(&self, message: &str) {
        println!("DEBUG [thrd:{:02} conn:  ]: {}", self.scheduler_id, message);
    }

    fn debug_conn(&self, conn: ConnectorId, message: &str) {
        println!("DEBUG [thrd:{:02} conn:{:02}]: {}", self.scheduler_id, conn.0, message);
    }
}

// -----------------------------------------------------------------------------
// ComponentCtx
// -----------------------------------------------------------------------------

enum ComponentStateChange {
    CreatedComponent(ConnectorPDL, Vec<PortIdLocal>),
    CreatedPort(Port),
    ChangedPort(ComponentPortChange),
}

#[derive(Clone)]
pub(crate) struct ComponentPortChange {
    pub is_acquired: bool, // otherwise: released
    pub port: Port,
}

/// The component context (better name may be invented). This was created
/// because part of the component's state is managed by the scheduler, and part
/// of it by the component itself. When the component starts a sync block or
/// exits a sync block the partially managed state by both component and
/// scheduler need to be exchanged.
pub(crate) struct ComponentCtxFancy {
    // Mostly managed by the scheduler
    pub(crate) id: ConnectorId,
    ports: Vec<Port>,
    inbox_messages: Vec<MessageFancy>, // never control or ping messages
    inbox_len_read: usize,
    // Submitted by the component
    is_in_sync: bool,
    changed_in_sync: bool,
    outbox: VecDeque<MessageFancy>,
    state_changes: VecDeque<ComponentStateChange>,
    // Workspaces that may be used by components to (generally) prevent
    // allocations. Be a good scout and leave it empty after you've used it.
    // TODO: Move to scheduler ctx, this is the wrong place
    pub workspace_ports: Vec<PortIdLocal>,
    pub workspace_branches: Vec<BranchId>,
}

impl ComponentCtxFancy {
    pub(crate) fn new_empty() -> Self {
        return Self{
            id: ConnectorId::new_invalid(),
            ports: Vec::new(),
            inbox_messages: Vec::new(),
            inbox_len_read: 0,
            is_in_sync: false,
            changed_in_sync: false,
            outbox: VecDeque::new(),
            state_changes: VecDeque::new(),
            workspace_ports: Vec::new(),
            workspace_branches: Vec::new(),
        };
    }

    /// Notify the runtime that the component has created a new component. May
    /// only be called outside of a sync block.
    pub(crate) fn push_component(&mut self, component: ConnectorPDL, initial_ports: Vec<PortIdLocal>) {
        debug_assert!(!self.is_in_sync);
        self.state_changes.push_back(ComponentStateChange::CreatedComponent(component, initial_ports));
    }

    /// Notify the runtime that the component has created a new port. May only
    /// be called outside of a sync block (for ports received during a sync
    /// block, pass them when calling `notify_sync_end`).
    pub(crate) fn push_port(&mut self, port: Port) {
        debug_assert!(!self.is_in_sync);
        self.state_changes.push_back(ComponentStateChange::CreatedPort(port))
    }

    pub(crate) fn get_port_by_id(&self, id: PortIdLocal) -> Option<&Port> {
        return self.ports.iter().find(|v| v.self_id == id);
    }

    fn get_port_mut_by_id(&mut self, id: PortIdLocal) -> Option<&mut Port> {
        return self.ports.iter_mut().find(|v| v.self_id == id);
    }

    /// Notify that component will enter a sync block. Note that after calling
    /// this function you must allow the scheduler to pick up the changes in
    /// the context by exiting your `Component::run` function with an
    /// appropriate scheduling value.
    pub(crate) fn notify_sync_start(&mut self) -> &[Port] {
        debug_assert!(!self.is_in_sync);

        self.is_in_sync = true;
        self.changed_in_sync = true;
        return &self.ports
    }

    #[inline]
    pub(crate) fn is_in_sync(&self) -> bool {
        return self.is_in_sync;
    }

    /// Submit a message for the scheduler to send to the appropriate receiver.
    /// May only be called inside of a sync block.
    pub(crate) fn submit_message(&mut self, contents: MessageFancy) {
        debug_assert!(self.is_in_sync);
        self.outbox.push_back(contents);
    }

    /// Notify that component just finished a sync block. Like
    /// `notify_sync_start`: drop out of the `Component::Run` function.
    pub(crate) fn notify_sync_end(&mut self, changed_ports: &[ComponentPortChange]) {
        debug_assert!(self.is_in_sync);

        self.is_in_sync = false;
        self.changed_in_sync = true;

        self.state_changes.reserve(changed_ports.len());
        for changed_port in changed_ports {
            self.state_changes.push_back(ComponentStateChange::ChangedPort(changed_port.clone()));
        }
    }

    /// Retrieves messages matching a particular port and branch id. But only
    /// those messages that have been previously received with
    /// `read_next_message`.
    pub(crate) fn get_read_data_messages(&self, match_port_id: PortIdLocal) -> MessagesIter {
        return MessagesIter {
            messages: &self.inbox_messages,
            next_index: 0,
            max_index: self.inbox_len_read,
            match_port_id
        };
    }

    /// Retrieves the next unread message from the inbox `None` if there are no
    /// (new) messages to read.
    // TODO: Fix the clone of the data message, entirely unnecessary
    pub(crate) fn read_next_message(&mut self) -> Option<MessageFancy> {
        if !self.is_in_sync { return None; }
        if self.inbox_len_read == self.inbox_messages.len() { return None; }

        // We want to keep data messages in the inbox, because we need to check
        // them in the future. We don't want to keep sync messages around, we
        // should only handle them once. Control messages should never be in
        // here.
        let message = &self.inbox_messages[self.inbox_len_read];
        match message {
            MessageFancy::Data(content) => {
                self.inbox_len_read += 1;
                return Some(MessageFancy::Data(content.clone()));
            },
            MessageFancy::Sync(_) => {
                let message = self.inbox_messages.remove(self.inbox_len_read);
                return Some(message);
            },
            MessageFancy::Control(_) => unreachable!("control message ended up in component inbox"),
        }
    }
}

pub(crate) struct MessagesIter<'a> {
    messages: &'a [MessageFancy],
    next_index: usize,
    max_index: usize,
    match_port_id: PortIdLocal,
}

impl<'a> Iterator for MessagesIter<'a> {
    type Item = &'a DataMessageFancy;

    fn next(&mut self) -> Option<Self::Item> {
        // Loop until match is found or at end of messages
        while self.next_index < self.max_index {
            let message = &self.messages[self.next_index];
            if let MessageFancy::Data(message) = &message {
                if message.data_header.target_port == self.match_port_id {
                    // Found a match
                    self.next_index += 1;
                    return Some(data_message);
                }
            } else {
                // Unreachable because:
                //  1. We only iterate over messages that were previously retrieved by `read_next_message`.
                //  2. Inbox does not contain control/ping messages.
                //  3. If `read_next_message` encounters anything else than a data message, it is removed from the inbox.
                unreachable!();
            }

            self.next_index += 1;
        }

        // No more messages
        return None;
    }
}

// -----------------------------------------------------------------------------
// Control messages
// -----------------------------------------------------------------------------

struct ControlEntry {
    id: u32,
    variant: ControlVariant,
}

enum ControlVariant {
    ChangedPort(ControlChangedPort),
    ClosedChannel(ControlClosedChannel),
}

struct ControlChangedPort {
    target_port: PortIdLocal,       // if send to this port, then reroute
    source_connector: ConnectorId,  // connector we expect messages from
    target_connector: ConnectorId,  // connector we need to reroute to
}

struct ControlClosedChannel {
    source_port: PortIdLocal,
    target_port: PortIdLocal,
}

pub(crate) struct ControlMessageHandler {
    id_counter: u32,
    active: Vec<ControlEntry>,
}

impl ControlMessageHandler {
    pub fn new() -> Self {
        ControlMessageHandler {
            id_counter: 0,
            active: Vec::new(),
        }
    }

    /// Prepares a message indicating that a channel has closed, we keep a local
    /// entry to match against the (hopefully) returned `Ack` message.
    pub fn prepare_closing_channel(
        &mut self, self_port_id: PortIdLocal, peer_port_id: PortIdLocal,
        self_connector_id: ConnectorId
    ) -> ControlMessageFancy {
        let id = self.take_id();

        self.active.push(ControlEntry{
            id,
            variant: ControlVariant::ClosedChannel(ControlClosedChannel{
                source_port: self_port_id,
                target_port: peer_port_id,
            }),
        });

        return ControlMessageFancy{
            id,
            sending_component_id: self_connector_id,
            content: ControlContent::CloseChannel(peer_port_id),
        };
    }

    /// Prepares rerouting messages due to changed ownership of a port. The
    /// control message returned by this function must be sent to the
    /// transferred port's peer connector.
    pub fn prepare_reroute(
        &mut self,
        port_id: PortIdLocal, peer_port_id: PortIdLocal,
        self_connector_id: ConnectorId, peer_connector_id: ConnectorId,
        new_owner_connector_id: ConnectorId
    ) -> ControlMessageFancy {
        let id = self.take_id();

        self.active.push(ControlEntry{
            id,
            variant: ControlVariant::ChangedPort(ControlChangedPort{
                target_port: port_id,
                source_connector: peer_connector_id,
                target_connector: new_owner_connector_id,
            }),
        });

        return ControlMessageFancy{
            id,
            sending_component_id: self_connector_id,
            content: ControlContent::PortPeerChanged(peer_port_id, new_owner_connector_id),
        };
    }

    /// Returns true if the supplied message should be rerouted. If so then this
    /// function returns the connector that should retrieve this message.
    pub fn should_reroute(&self, sending_connector: ConnectorId, target_port: PortIdLocal) -> Option<ConnectorId> {
        for entry in &self.active {
            if let ControlVariant::ChangedPort(entry) = &entry.variant {
                if entry.target_port == target_port {
                    // Need to reroute this message
                    return Some(entry.target_connector);
                }
            }
        }

        return None;
    }

    /// Handles an Ack as an answer to a previously sent control message
    pub fn handle_ack(&mut self, id: u32) {
        let index = self.active.iter()
            .position(|v| v.id == id);

        match index {
            Some(index) => { self.active.remove(index); },
            None => { todo!("handling of nefarious ACKs"); },
        }
    }

    /// Retrieves the number of responses we still expect to receive from our
    /// peers
    #[inline]
    pub fn num_pending_acks(&self) -> usize {
        return self.active.len();
    }

    fn take_id(&mut self) -> u32 {
        let generated_id = self.id_counter;
        let (new_id, _) = self.id_counter.overflowing_add(1);
        self.id_counter = new_id;

        return generated_id;
    }
}