use std::collections::HashMap;
use crate::bus::{Response, BusConnection, RawReceiver, RequestEnvelope, ResponseEnvelope, RawBusClient};
use crate::utils::asyn::{AsyncQueue, Deferred};
use futures::Stream;
use async_std::sync::{Arc, Mutex};
use async_macros::utils::task::{Context, Poll};
use async_macros::utils::pin::Pin;

#[derive(Clone, Debug)]
pub struct ProcessReceiver {
    state: Arc<Mutex<ProcessReceiverState>>,
    request_queue: AsyncQueue<RequestEnvelope>,
}

impl ProcessReceiver {
    pub fn new() -> Self {
        ProcessReceiver {
            state: Arc::new(Mutex::new(ProcessReceiverState {
                id: 0,
                callbacks: Default::default(),
                registered_client: false,
            })),
            request_queue: Default::default(),
        }
    }

    async fn run_callback(&mut self, id: u64, response: Option<Response>) {
        let mut state = self.state.lock().await;
        if let Some(mut deferred) = state.callbacks.remove(&id) {
            deferred.complete(ResponseEnvelope {
                id,
                response,
            })
        }
    }
}

#[async_trait]
impl RawReceiver for ProcessReceiver {
    fn name<'s>(&self) -> &'s str {
        "in-process"
    }

    async fn disconnect(&mut self) {}

    async fn get(&mut self) -> Option<Box<dyn BusConnection>> {
        let mut state = self.state.lock().await;
        if state.registered_client {
            None
        } else {
            state.registered_client = true;
            let copy = Box::new(self.clone()) as Box<dyn BusConnection>;
            Some(copy)
        }
    }
}

#[derive(Debug)]
struct ProcessReceiverState {
    id: u64,
    callbacks: HashMap<u64, Deferred<ResponseEnvelope>>,
    registered_client: bool,
}

#[async_trait]
impl BusConnection for ProcessReceiver {
    async fn close(&mut self) {}

    async fn ack(&mut self, id: u64) {
        self.run_callback(id, None).await;
    }

    async fn send(&mut self, id: u64, message: Response) {
        self.run_callback(id, Some(message)).await;
    }
}

impl Stream for ProcessReceiver {
    type Item = RequestEnvelope;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        Stream::poll_next(unsafe { Pin::new_unchecked(&mut self.request_queue) }, cx)
    }
}

#[async_trait]
impl RawBusClient for ProcessReceiver {
    async fn send(&mut self, request: RequestEnvelope) -> ResponseEnvelope {
        let deferred = Deferred::new();
        let def_clone = deferred.clone();
        {
            let mut state = self.state.lock().await;
            state.callbacks.insert(request.id, def_clone);
        }

        self.request_queue.push(request);
        deferred.await
    }
}