use std::{collections::HashMap, time::Instant};

use conduwuit::{
	Event, PduEvent, debug, debug_info, debug_warn, trace,
	utils::{BoolExt, IterStream, stream::BroadbandExt},
};
use futures::StreamExt;
use ruma::{MilliSecondsSinceUnixEpoch, RoomId, ServerName};

use crate::rooms::event_handler::{build_local_dag, fetch_and_handle_outliers::DagBuilderTree};

impl super::Service {
	/// Fetches any missing prev_events for this event and persists them before
	/// returning. The caller is responsible for then handling the incoming PDU.
	pub(super) async fn fetch_prevs(
		&self,
		room_id: &RoomId,
		create_event: &PduEvent,
		incoming_pdu: &PduEvent,
		origin: &ServerName,
		first_ts_in_room: MilliSecondsSinceUnixEpoch,
	) -> conduwuit::Result<()> {
		let start = Instant::now();
		let mut missing = incoming_pdu
			.prev_events()
			.stream()
			.broad_filter_map(|event_id| async move {
				self.services
					.timeline
					.get_non_outlier_pdu_json(event_id)
					.await
					.is_ok()
					.or(|| event_id.to_owned())
			})
			.collect::<Vec<_>>()
			.await;
		if missing.is_empty() {
			debug!(elapsed=?start.elapsed(), event_id=%incoming_pdu.event_id(), "No missing prev events.");
			return Ok(());
		}
		debug!(elapsed=?start.elapsed(), %room_id, event_id=%incoming_pdu.event_id(), ?missing, "Fetching previous events");
		let tail = self
			.services
			.state
			.get_forward_extremities(room_id)
			.collect::<Vec<_>>()
			.await;

		let mut gapfilled = self
			.get_missing_events(
				room_id,
				incoming_pdu,
				tail,
				origin,
				self.services
					.metadata
					.get_mindepth(room_id)
					.await
					.saturating_sub(
						u8::try_from(incoming_pdu.prev_events.len())
							.unwrap()
							.saturating_mul(2)
							.into(),
					),
			)
			.await?;
		debug_info!(elapsed=?start.elapsed(), "Fetched {} missing events", gapfilled.len());
		missing.retain(|eid| !gapfilled.contains_key(eid));
		if !missing.is_empty() {
			debug_warn!(elapsed=?start.elapsed(), "Still missing {} events, falling back to atomic fetch.", missing.len());
			gapfilled.extend(
				self.fetch_prev_events(origin, missing, create_event, room_id)
					.await,
			);
		}

		// Persist all fetched events
		let mapped = gapfilled
			.iter()
			.map(|(eid, evt)| {
				let mut obj = evt.to_canonical_object();
				obj.remove("event_id"); // event_id is inserted by backfill_missing_events
				(eid.clone(), obj)
			})
			.collect::<HashMap<_, _>>();

		let to_persist = build_local_dag(&mapped, DagBuilderTree::PrevEvents).await?;

		let job_start = Instant::now();
		trace!("Starting to persist {} prev events", to_persist.len());
		for (i, event_id) in to_persist.iter().enumerate() {
			debug!(
				elapsed=?start.elapsed(),
				"Persisting fetched prev event: {event_id} ({}/{})",
				i.saturating_add(1),
				to_persist.len(),
			);
			let obj = mapped.get(event_id).cloned().unwrap();
			let persist_start = Instant::now();
			match self
				.handle_outlier_pdu(origin, create_event, event_id, room_id, obj)
				.await
			{
				| Ok((pdu, val)) if pdu.origin_server_ts() >= first_ts_in_room => {
					Box::pin(self.upgrade_outlier_to_timeline_pdu(
						pdu,
						val,
						create_event,
						origin,
						room_id,
					))
					.await
					.inspect_err(|e| {
						debug_warn!(
							total_elapsed=?start.elapsed(),
							job_elapsed=?job_start.elapsed(),
							task_elapsed=?persist_start.elapsed(),
							"Failed to upgrade prev event {event_id}: {e}",
						);
					})
					.inspect(|_| {
						debug_info!(
							total_elapsed=?start.elapsed(),
							job_elapsed=?job_start.elapsed(),
							task_elapsed=?persist_start.elapsed(),
							"Upgraded prev event {event_id}",
						);
					})
					.ok();
				},
				| Err(e) => debug_warn!(
					total_elapsed=?start.elapsed(),
					job_elapsed=?job_start.elapsed(),
					task_elapsed=?persist_start.elapsed(),
					"Failed to persist prev event {event_id}: {e}",
				),
				| _ => {},
			}
		}

		// NOTE because i keep forgetting: the caller persists incoming_pdu.
		// we only care about its prev events
		trace!(
			total_elapsed=?start.elapsed(),
			persist_elapsed=?job_start.elapsed(),
		);
		Ok(())
	}
}
