use std::pin::pin;
use std::time::Duration;

use futures::FutureExt;
use mm1_address::address::Address;
use mm1_common::errors::error_of::ErrorOf;
use mm1_common::impl_error_kind;
use mm1_proc_macros::dispatch;
use mm1_proto::message;
use mm1_proto_system::Down;
use tracing::warn;

use super::{ForkErrorKind, Messaging, RecvErrorKind};
use crate::context::{Fork, Watching};

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[message(base_path = ::mm1_proto)]
pub enum ShutdownErrorKind {
    InternalError,
    Fork(ForkErrorKind),
    Recv(RecvErrorKind),
}

pub trait Stop {
    fn shutdown(
        &mut self,
        peer: Address,
        stop_timeout: Duration,
    ) -> impl Future<Output = Result<(), ErrorOf<ShutdownErrorKind>>> + Send
    where
        Self: Watching + Fork + Messaging,
    {
        async move {
            let mut fork = self
                .fork()
                .await
                .map_err(|e| e.map_kind(ShutdownErrorKind::Fork))?;

            let watch_ref = fork.watch(peer).await;

            let mut shutdown_sequence = pin!(
                async {
                    self.exit(peer).await;
                    tokio::time::sleep(stop_timeout).await;
                    self.kill(peer).await;
                }
                .fuse()
            );

            let mut recv_result = pin!(
                async {
                    loop {
                        dispatch!(match fork
                            .recv()
                            .await
                            .map_err(|e| e.map_kind(ShutdownErrorKind::Recv))?
                        {
                            down @ Down { .. }
                                if down.watch_ref == watch_ref && down.peer == peer =>
                            {
                                break Ok(())
                            },

                            unexpected @ _ => {
                                warn!(msg = ?unexpected, "unexpected message");
                            },
                        })
                    }
                }
                .fuse()
            );

            loop {
                tokio::select! {
                    recv_result = recv_result.as_mut() => { break recv_result },
                    () = shutdown_sequence.as_mut() => {}
                }
            }
        }
    }

    fn exit(&mut self, peer: Address) -> impl Future<Output = bool> + Send;
    fn kill(&mut self, peer: Address) -> impl Future<Output = bool> + Send;
}

impl_error_kind!(ShutdownErrorKind);
