use crate::rpc;
use jam_types::ServiceId;
use jsonrpsee::{SubscriptionMessage, SubscriptionSink};
use parking_lot::Mutex;
use schnellru::{ByLength, LruMap};
use std::collections::VecDeque;

/// Log message of a particular service.
struct LogMessages {
	/// Circular buffer of messages.
	messages: VecDeque<rpc::LogMessage>,
	/// Total size in bytes of all messages.
	total_size: u32,
	/// RPC clients subscribed to the logs of this service.
	subscribers: Vec<SubscriptionSink>,
}

impl LogMessages {
	fn new() -> Self {
		Self { messages: VecDeque::new(), total_size: 0, subscribers: Vec::new() }
	}

	/// Push a new message to the buffer.
	///
	/// At least one message is retained in the buffer after this call even if it alone doesn't fit
	/// into `max_size` bytes.
	fn push(&mut self, message: rpc::LogMessage, max_size: u32) {
		let message_len = message.size_in_bytes() as u32;
		self.total_size += message_len;
		self.remove_old_messages(max_size);
		self.notify_subscribers(message.as_ref());
		self.messages.push_back(message);
	}

	#[inline]
	fn remove_old_messages(&mut self, max_size: u32) {
		while self.total_size > max_size {
			let Some(old_message) = self.messages.pop_front() else {
				break;
			};
			self.total_size -= old_message.size_in_bytes() as u32;
		}
	}

	#[inline]
	fn notify_subscribers(&mut self, message: rpc::LogMessageRef<'_>) {
		self.subscribers.retain_mut(|sink| {
			let mut retain = true;
			if let Ok(msg) =
				SubscriptionMessage::new(sink.method_name(), sink.subscription_id(), &message)
			{
				if let Err(jsonrpsee::TrySendError::Closed(..)) = sink.try_send(msg) {
					retain = false;
				}
			}
			retain
		});
	}
}

pub struct LogCache {
	cache: Mutex<LruMap<ServiceId, LogMessages>>,
	max_bytes_per_service: u32,
}

impl LogCache {
	pub fn new(max_services: u32, max_bytes_per_service: u32) -> Self {
		Self { cache: Mutex::new(LruMap::new(ByLength::new(max_services))), max_bytes_per_service }
	}

	pub(crate) fn get(&self, service_id: ServiceId) -> Option<Vec<rpc::LogMessage>> {
		self.cache
			.lock()
			.get(&service_id)
			.map(|messages| messages.messages.iter().cloned().collect())
	}

	#[must_use]
	pub(crate) fn subscribe(&self, service_id: ServiceId, mut sink: SubscriptionSink) -> bool {
		match self.cache.lock().get(&service_id) {
			Some(messages) => {
				// Send old log messages.
				for message in messages.messages.iter() {
					if let Ok(msg) = SubscriptionMessage::new(
						sink.method_name(),
						sink.subscription_id(),
						&message.as_ref(),
					) {
						if let Err(jsonrpsee::TrySendError::Closed(..)) = sink.try_send(msg) {
							return false;
						}
					}
				}
				messages.subscribers.push(sink);
				true
			},
			None => false,
		}
	}

	pub(crate) fn insert(
		&self,
		service_id: ServiceId,
		timestamp: u64,
		level: log::Level,
		target: String,
		message: String,
	) {
		let message = rpc::LogMessage { timestamp, level, target, message };
		let mut cache = self.cache.lock();
		match cache.get(&service_id) {
			Some(messages) => messages.push(message, self.max_bytes_per_service),
			None => {
				let mut messages = LogMessages::new();
				messages.push(message, self.max_bytes_per_service);
				cache.insert(service_id, messages);
			},
		}
	}
}
