use crate::common::*;

use crate::runtime::{actors::*, endpoint::*, errors::*, *};

impl Controller {

    fn end_round_with_decision(&mut self, decision: Decision) -> Result<(), SyncErr> {

        log!(&mut self.inner.logger, "ENDING ROUND WITH DECISION! {:?}", &decision);

        let ret = match &decision {

            Decision::Success(predicate) => {

                // overwrite MonoN/P

                self.inner.mono_n =

                    self.ephemeral.poly_n.take().unwrap().choose_mono(predicate).unwrap();

                self.inner.mono_ps.clear();

                self.inner.mono_ps.extend(

                    self.ephemeral

                        .poly_ps

                        .drain(..)

                        .map(|poly_p| poly_p.choose_mono(predicate).unwrap()),

                );

                Ok(())

            }

            Decision::Failure => Err(SyncErr::Timeout),

        };

        let announcement = CommMsgContents::Announce { decision }.into_msg(self.inner.round_index);

        for &child_ekey in self.inner.family.children_ekeys.iter() {

            log!(

                &mut self.inner.logger,

                "Forwarding {:?} to child with ekey {:?}",

                &announcement,

                child_ekey

            );

            self.inner

                .endpoint_exts

                .get_mut(child_ekey)

                .expect("eefef")

                .endpoint

                .send(announcement.clone())?;

        }

        self.inner.round_index += 1;

        self.ephemeral.clear();

        ret

    }

    // Drain self.ephemeral.solution_storage and handle the new locals. Return decision if one is found

    fn handle_locals_maybe_decide(&mut self) -> Result<bool, SyncErr> {

        if let Some(parent_ekey) = self.inner.family.parent_ekey {

            // I have a parent -> I'm not the leader

            let parent_endpoint =

                &mut self.inner.endpoint_exts.get_mut(parent_ekey).expect("huu").endpoint;

            for partial_oracle in self.ephemeral.solution_storage.iter_new_local_make_old() {

                let msg =

                    CommMsgContents::Elaborate { partial_oracle }.into_msg(self.inner.round_index);

                log!(&mut self.inner.logger, "Sending {:?} to parent {:?}", &msg, parent_ekey);

                parent_endpoint.send(msg)?;

            }

            Ok(false)

        } else {

            // I have no parent -> I'm the leader

            assert!(self.inner.family.parent_ekey.is_none());

            let maybe_predicate = self.ephemeral.solution_storage.iter_new_local_make_old().next();

            Ok(if let Some(predicate) = maybe_predicate {

                let decision = Decision::Success(predicate);

                log!(&mut self.inner.logger, "DECIDE ON {:?} AS LEADER!", &decision);

                self.end_round_with_decision(decision)?;

                true

            } else {

                false

            })

        }

    }

    fn kick_off_native(

        &mut self,

        sync_batches: impl Iterator<Item = SyncBatch>,

    ) -> Result<PolyN, EndpointErr> {

        let MonoN { ekeys, .. } = self.inner.mono_n.clone();

        let Self { inner: ControllerInner { endpoint_exts, round_index, .. }, .. } = self;

        let mut branches = HashMap::<_, _>::default();

        for (sync_batch_index, SyncBatch { puts, gets }) in sync_batches.enumerate() {

            let ekey_to_channel_id = |ekey| endpoint_exts.get(ekey).unwrap().info.channel_id;

            let all_ekeys = ekeys.iter().copied();

            let all_channel_ids = all_ekeys.map(ekey_to_channel_id);

            let mut predicate = Predicate::new_trivial();

            // assign TRUE for puts and gets

            let true_ekeys = puts.keys().chain(gets.iter()).copied();

            let true_channel_ids = true_ekeys.clone().map(ekey_to_channel_id);

            predicate.batch_assign_nones(true_channel_ids, true);

            // assign FALSE for all in interface not assigned true

            predicate.batch_assign_nones(all_channel_ids.clone(), false);

            if branches.contains_key(&predicate) {

                // TODO what do I do with redundant predicates?

                unimplemented!(

                    "Having multiple batches with the same

                    predicate requires the support of oracle boolean variables"

                )

            }

            let branch = BranchN { to_get: gets, gotten: Default::default(), sync_batch_index };

            for (ekey, payload) in puts {

                log!(

                    &mut self.inner.logger,

                    "... ... Initial native put msg {:?} pred {:?} batch {:?}",

                    &payload,

                    &predicate,

                    sync_batch_index,

                );

                let msg =

                    CommMsgContents::SendPayload { payload_predicate: predicate.clone(), payload }

                        .into_msg(*round_index);

                endpoint_exts.get_mut(ekey).unwrap().endpoint.send(msg)?;

            }

            log!(

                &mut self.inner.logger,

                "... Initial native branch batch index={} with pred {:?}",

                sync_batch_index,

                &predicate

            );

            if branch.to_get.is_empty() {

                self.ephemeral.solution_storage.submit_and_digest_subtree_solution(

                    &mut self.inner.logger,

                    SubtreeId::PolyN,

                    predicate.clone(),

                );

            }

            branches.insert(predicate, branch);

        }

        Ok(PolyN { ekeys, branches })

    }

    // Runs a synchronous round until all the actors are in decided state OR 1+ are inconsistent.

    // If a native requires setting up, arg `sync_batches` is Some, and those are used as the sync batches.

        &mut self,

        mut deadline: Option<Instant>,

        sync_batches: Option<impl Iterator<Item = SyncBatch>>,

    ) -> Result<(), SyncErr> {

        log!(

            &mut self.inner.logger,

            "~~~~~~~~ SYNC ROUND STARTS! ROUND={} ~~~~~~~~~",

            self.inner.round_index

        );

        assert!(self.ephemeral.is_clear());

        // 1. Run the Mono for each Mono actor (stored in `self.mono_ps`).

        //    Some actors are dropped. some new actors are created.

        //    Ultimately, we have 0 Mono actors and a list of unnamed sync_actors

        self.ephemeral.mono_ps.extend(self.inner.mono_ps.iter().cloned());

        log!(&mut self.inner.logger, "Got {} MonoP's to run!", self.ephemeral.mono_ps.len());

        while let Some(mut mono_p) = self.ephemeral.mono_ps.pop() {

            let mut m_ctx = MonoPContext {

                ekeys: &mut mono_p.ekeys,

                mono_ps: &mut self.ephemeral.mono_ps,

                inner: &mut self.inner,

            };

            // cross boundary into crate::protocol

            let blocker = mono_p.state.pre_sync_run(&mut m_ctx, &self.protocol_description);

            log!(&mut self.inner.logger, "... MonoP's pre_sync_run got blocker {:?}", &blocker);

            match blocker {

                MonoBlocker::Inconsistent => return Err(SyncErr::Inconsistent),

                MonoBlocker::ComponentExit => drop(mono_p),

                MonoBlocker::SyncBlockStart => self.ephemeral.poly_ps.push(mono_p.into()),

            }

        }

        log!(

            &mut self.inner.logger,

            "Finished running all MonoPs! Have {} PolyPs waiting",

            self.ephemeral.poly_ps.len()

        );

        // 3. define the mapping from ekey -> actor

        //    this is needed during the event loop to determine which actor

        //    should receive the incoming message.

        //    TODO: store and update this mapping rather than rebuilding it each round.

        let ekey_to_holder: HashMap<Key, PolyId> = {

            use PolyId::*;

            let n = self.inner.mono_n.ekeys.iter().map(move |&e| (e, N));

            let p = self

                .ephemeral

                .poly_ps

                .iter()

                .enumerate()

                .flat_map(|(index, m)| m.ekeys.iter().map(move |&e| (e, P { index })));

            n.chain(p).collect()

        };

        log!(

            &mut self.inner.logger,

            "SET OF PolyPs and MonoPs final! ekey lookup map is {:?}",

            &ekey_to_holder

        );

        // 4. Create the solution storage. it tracks the solutions of "subtrees"

        //    of the controller in the overlay tree.

        self.ephemeral.solution_storage.reset({

            let n = std::iter::once(SubtreeId::PolyN);

            let m = (0..self.ephemeral.poly_ps.len()).map(|index| SubtreeId::PolyP { index });

            let c = self

                .inner

                .family

                .children_ekeys

                .iter()

                .map(|&ekey| SubtreeId::ChildController { ekey });

            let subtree_id_iter = n.chain(m).chain(c);

            log!(

                &mut self.inner.logger,

                "Solution Storage has subtree Ids: {:?}",

                &subtree_id_iter.clone().collect::<Vec<_>>()

            );

            subtree_id_iter

        });

        // 5. kick off the synchronous round of the native actor if it exists

        log!(&mut self.inner.logger, "Kicking off native's synchronous round...");

        self.ephemeral.poly_n = if let Some(sync_batches) = sync_batches {

            // using if let because of nested ? operator

            // TODO check that there are 1+ branches or NO SOLUTION

            let poly_n = self.kick_off_native(sync_batches)?;

            log!(

                &mut self.inner.logger,

                "PolyN kicked off, and has branches with predicates... {:?}",

                poly_n.branches.keys().collect::<Vec<_>>()

            );

            Some(poly_n)

        } else {

            log!(&mut self.inner.logger, "NO NATIVE COMPONENT");

            None

        };

        // 6. Kick off the synchronous round of each protocol actor

        //    If just one actor becomes inconsistent now, there can be no solution!

        //    TODO distinguish between completed and not completed poly_p's?

        log!(&mut self.inner.logger, "Kicking off {} PolyP's.", self.ephemeral.poly_ps.len());

        for (index, poly_p) in self.ephemeral.poly_ps.iter_mut().enumerate() {

            let my_subtree_id = SubtreeId::PolyP { index };

            let m_ctx = PolyPContext {

                my_subtree_id,

                inner: &mut self.inner,

                solution_storage: &mut self.ephemeral.solution_storage,

            };

            use SyncRunResult as Srr;

            let blocker = poly_p.poly_run(m_ctx, &self.protocol_description)?;

            log!(&mut self.inner.logger, "... PolyP's poly_run got blocker {:?}", &blocker);

            match blocker {

                Srr::NoBranches => return Err(SyncErr::Inconsistent),

                Srr::AllBranchesComplete | Srr::BlockingForRecv => (),

            }

        }

        log!(&mut self.inner.logger, "All Poly machines have been kicked off!");

        // 7. `solution_storage` may have new solutions for this controller

        //    handle their discovery. LEADER => announce, otherwise => send to parent

        {

            let peeked = self.ephemeral.solution_storage.peek_new_locals().collect::<Vec<_>>();

            log!(

                &mut self.inner.logger,

                "Got {} controller-local solutions before a single RECV: {:?}",

                peeked.len(),

                peeked

            );

        }

        if self.handle_locals_maybe_decide()? {

            return Ok(());

        }

        // 4. Receive incoming messages until the DECISION is made OR some unrecoverable error

        log!(&mut self.inner.logger, "`No decision yet`. Time to recv messages");

        self.undelay_all();

        'recv_loop: loop {

            log!(&mut self.inner.logger, "`POLLING` with deadline {:?}...", deadline);

            let received = match deadline {

                None => {

                    // we have personally timed out. perform a "long" poll.

                    self.recv(Instant::now() + Duration::from_secs(10))?.expect("DRIED UP")

                }

                Some(d) => match self.recv(d)? {

                    // we have not yet timed out. performed a time-limited poll

                    Some(received) => received,

                    None => {

                        // timed out! send a FAILURE message to the sink,

                        // and henceforth don't time out on polling.

                        deadline = None;

                        match self.inner.family.parent_ekey {

                            None => {

                                // I am the sink! announce failure and return.

                                return self.end_round_with_decision(Decision::Failure);

                            }

                            Some(parent_ekey) => {

                                // I am not the sink! send a failure message.

                                let announcement = Msg::CommMsg(CommMsg {

                                    round_index: self.inner.round_index,

                                    contents: CommMsgContents::Failure,

                                });

                                log!(

                                    &mut self.inner.logger,

                                    "Forwarding {:?} to parent with ekey {:?}",

                                    &announcement,

                                    parent_ekey

                                );

                                self.inner

                                    .endpoint_exts

                                    .get_mut(parent_ekey)

                                    .expect("ss")

                                    .endpoint

                                    .send(announcement.clone())?;

                                continue; // poll some more

                            }

                        }

                    }

                },

            };

            log!(&mut self.inner.logger, "::: message {:?}...", &received);

            let current_content = match received.msg {

                Msg::SetupMsg(s) => {

                    // This occurs in the event the connector was malformed during connect()

                    println!("WASNT EXPECTING {:?}", s);

                    return Err(SyncErr::UnexpectedSetupMsg);

                }

                Msg::CommMsg(CommMsg { round_index, .. })

                    if round_index < self.inner.round_index =>

                {

                    // Old message! Can safely discard

                    log!(&mut self.inner.logger, "...and its OLD! :(");

                    drop(received);

                    continue 'recv_loop;

                }

                Msg::CommMsg(CommMsg { round_index, .. })

                    if round_index > self.inner.round_index =>

                {

                    // Message from a next round. Keep for later!

                    log!(&mut self.inner.logger, "... DELAY! :(");

                    self.delay(received);

                    continue 'recv_loop;

                }

                Msg::CommMsg(CommMsg { contents, round_index }) => {