// Copyright Marcus Del Favero 2025
// Licensed under the GNU AGPLv3 with an exception, see `README.md` for details
use std::iter;

use tokio::sync::mpsc::{self, Sender as MpscSender, Receiver as MpscReceiver, error::TryRecvError};
use serde::de::DeserializeOwned;

use super::{CHANNEL_SIZE, error::ErrorContainer};

use crate::protocol::message::Message;
use crate::net::serp::{Error, Receiver};

pub struct ReceiverTaskBuilder {
	sender: MpscSender<Message>,
	receiver: MpscReceiver<Message>,
	error: ErrorContainer,
}

pub trait GetMessages {
	fn get(self) -> impl Iterator<Item = Message> + Send;
}

impl<M> GetMessages for M where M: Into<Message> + Send {
	fn get(self) -> impl Iterator<Item = Message> + Send {
		iter::once(self.into())
	}
}

impl<M> GetMessages for Vec<M> where M: Into<Message> + Send {
	fn get(self) -> impl Iterator<Item = Message> + Send {
		self.into_iter().map(Into::into)
	}
}

impl ReceiverTaskBuilder {
	pub fn new(error: ErrorContainer) -> ReceiverTaskBuilder {
		let (sender, receiver) = mpsc::channel(CHANNEL_SIZE);
		ReceiverTaskBuilder { sender, receiver, error }
	}

	pub fn add_stream<M>(self, mut stream: Receiver, size_limit: usize) -> Self where M: GetMessages + DeserializeOwned + Send {
		let sender = self.sender.clone();
		let error = self.error.clone();
		tokio::spawn(async move {
			error.insert('outer: loop {
				match stream.recv_with_limit::<M>(size_limit).await {
					Ok(messages) => {
						for msg in messages.get() {
							if sender.send(msg).await.is_err() {
								break 'outer Error::Disconnected(None);
							}
						}
					},
					Err(err) => break err,
				}
			});
		});

		self
	}

	/**
	 * Receives a message right now from the currently added streams.
	 */
	#[cfg(feature = "client")]
	pub async fn receive(&mut self, error: &ErrorContainer) -> Result<Message, Error> {
		self.receiver.recv().await.ok_or_else(|| error.get())
	}

	pub fn get_sender(&self) -> MpscSender<Message> {
		self.sender.clone()
	}

	pub fn build(self) -> ReceiverTask {
		ReceiverTask(self.receiver)
	}
}

pub struct ReceiverTask(MpscReceiver<Message>);

impl ReceiverTask {
	pub fn receive(&mut self, error: &ErrorContainer) -> Result<Option<Message>, Error> {
		match self.0.try_recv() {
			Ok(msg) => Ok(Some(msg)),
			Err(TryRecvError::Empty) => Ok(None),
			Err(TryRecvError::Disconnected) => Err(error.get()),
		}
	}
}
