//! Track pending consumer deliveries for resolver implementations.
//!
//! Resolvers often need the same delivery lifecycle: keep a fetch alive while
//! `Consumer::deliver` validates a response, abort that validation if the fetch
//! is pruned, and reuse a cached response for subscribers that were added
//! while validation was in progress. This module owns that lifecycle without
//! making assumptions about how data is fetched.

use crate::{Consumer, Delivery, Outcome};
use ahash::AHashMap;
use commonware_macros::select;
use commonware_utils::{channel::oneshot, futures::Pool};
use futures::future::Aborted;
use std::collections::hash_map::Entry as HashMapEntry;
use tracing::debug;

/// Completed consumer validation for a delivery.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Completion<K, S, Context = ()> {
    /// Resolver-specific context associated with the delivery attempt.
    pub context: Context,

    /// Key and subscribers that were passed to the consumer.
    pub delivery: Delivery<K, S>,

    /// Consumer disposition for the delivered response, or `None` if the
    /// consumer dropped the sender without reporting one.
    pub outcome: Option<Outcome>,

    /// Whether the tracked fetch was removed before this completion was polled.
    /// Its verdict applies only to the source, not to any current fetch for the key.
    pub canceled: bool,

    /// Whether the cached response was already accepted when this delivery began.
    pub previously_accepted: bool,
}

// Delivery futures handle cancellation so verdicts sent before receiver closure can still
// reach source accounting. Canceled deliveries without a verdict return `Aborted`.
type DeliveryPool<K, S, Context> = Pool<'static, Result<Completion<K, S, Context>, Aborted>>;

// Cached response that can be redelivered while it is accepted or still unjudged.
struct Response<Context, V> {
    context: Context,
    value: V,
    accepted: bool,
}

// Per-key delivery state retained while a resolver fetch is active.
struct Entry<Context, V, State> {
    delivery: Option<oneshot::Sender<()>>,
    response: Option<Response<Context, V>>,
    state: Option<State>,
}

impl<Context, V, State> Entry<Context, V, State> {
    const fn new(state: State) -> Self {
        Self {
            delivery: None,
            response: None,
            state: Some(state),
        }
    }
}

/// Tracks in-flight consumer deliveries keyed by resolver key.
///
/// `Context` carries resolver-specific metadata back to the caller when
/// validation completes.
///
/// `State` holds per-key resolver state that is dropped when the key is pruned.
/// While the key stays tracked, the resolver can read, update, or take it.
pub struct Tracker<Con, Context = (), State = ()>
where
    Con: Consumer,
    Con::Value: Clone + Send + 'static,
    Context: Clone + Send + 'static,
{
    entries: AHashMap<Con::Key, Entry<Context, Con::Value, State>>,
    deliveries: DeliveryPool<Con::Key, Con::Subscriber, Context>,
    consumer: Con,
}

impl<Con, Context, State> Tracker<Con, Context, State>
where
    Con: Consumer,
    Con::Value: Clone + Send + 'static,
    Context: Clone + Send + 'static,
{
    /// Create an empty tracker backed by the provided consumer.
    pub fn new(consumer: Con) -> Self {
        Self {
            entries: AHashMap::new(),
            deliveries: Pool::default(),
            consumer,
        }
    }

    /// Returns true when the key has an active tracked fetch.
    pub fn contains(&self, key: &Con::Key) -> bool {
        self.entries.contains_key(key)
    }

    /// Start tracking a key.
    ///
    /// Returns `true` when the key was inserted. If the key is already tracked,
    /// this leaves the existing entry untouched and returns `false`.
    pub(crate) fn insert_with_state(&mut self, key: Con::Key, state: State) -> bool {
        match self.entries.entry(key) {
            HashMapEntry::Vacant(entry) => {
                entry.insert(Entry::new(state));
                true
            }
            HashMapEntry::Occupied(_) => false,
        }
    }

    /// Remove a key and abort any in-progress delivery for it.
    ///
    /// Returns true if the key was present. A verdict accepted before cancellation
    /// closes its receiver is returned as a canceled [`Completion`].
    pub fn remove(&mut self, key: &Con::Key) -> bool {
        self.entries.remove(key).is_some()
    }

    /// Remove a key, aborting any in-progress delivery, and return its state.
    ///
    /// Returns `None` if the key was absent. The inner `Option` is `None` if the
    /// state was already taken while the key stayed active.
    pub(crate) fn remove_with_state(&mut self, key: &Con::Key) -> Option<Option<State>> {
        self.entries.remove(key).map(|entry| entry.state)
    }

    /// Take the key's state without removing the tracked key.
    ///
    /// Returns `None` when the key is absent or the state was already taken.
    pub(crate) fn take_state(&mut self, key: &Con::Key) -> Option<State> {
        self.entries
            .get_mut(key)
            .and_then(|entry| entry.state.take())
    }

    /// Mutably borrow the key's state without taking it.
    ///
    /// Returns `None` when the key is absent or the state was already taken.
    pub(crate) fn state_mut(&mut self, key: &Con::Key) -> Option<&mut State> {
        self.entries
            .get_mut(key)
            .and_then(|entry| entry.state.as_mut())
    }

    /// Retain only entries for which the predicate returns true.
    ///
    /// Dropped entries abort in-progress deliveries. Returns the number of
    /// removed entries.
    pub fn retain<F: FnMut(&Con::Key) -> bool>(&mut self, mut predicate: F) -> usize {
        self.entries.extract_if(|key, _| !predicate(key)).count()
    }

    /// Remove all entries and abort all in-progress deliveries.
    ///
    /// Returns the number of entries removed.
    pub fn drain(&mut self) -> usize {
        let count = self.entries.len();
        self.entries.clear();
        count
    }

    /// Deliver a newly received response to the consumer.
    ///
    /// The response is cached so that later retained subscribers can be
    /// redelivered the same bytes with [`redeliver`](Self::redeliver) once the
    /// consumer accepts it or drops its verdict. Panics if the key is not tracked.
    pub fn deliver(
        &mut self,
        delivery: Delivery<Con::Key, Con::Subscriber>,
        context: Context,
        value: Con::Value,
    ) {
        let entry = self.entries.get_mut(&delivery.key).expect("delivery entry");
        entry.response = Some(Response {
            context: context.clone(),
            value: value.clone(),
            accepted: false,
        });
        self.push_delivery(delivery, context, value);
    }

    /// Deliver the cached response to another set of subscribers.
    ///
    /// This is intended for subscribers added while an earlier validation was
    /// still pending. The cached response is either accepted or still unjudged
    /// because the consumer dropped the earlier verdict. Panics if the key is not
    /// tracked or no response is cached.
    pub fn redeliver(&mut self, delivery: Delivery<Con::Key, Con::Subscriber>) {
        let (context, value) = {
            let entry = self.entries.get(&delivery.key).expect("delivery entry");
            let response = entry.response.as_ref().expect("response");
            (response.context.clone(), response.value.clone())
        };
        self.push_delivery(delivery, context, value);
    }

    /// Returns true if the cached response for this key has been accepted.
    pub fn response_accepted(&self, key: &Con::Key) -> bool {
        self.entries
            .get(key)
            .and_then(|entry| entry.response.as_ref())
            .is_some_and(|response| response.accepted)
    }

    /// Mark the cached response accepted by the consumer.
    ///
    /// Panics if the key is not tracked or no response is cached.
    pub fn accept_response(&mut self, key: &Con::Key) {
        let entry = self.entries.get_mut(key).expect("delivery entry");
        let response = entry.response.as_mut().expect("response");
        response.accepted = true;
    }

    /// Drop the cached response without removing the tracked key.
    ///
    /// Use this when a response is invalid or does not satisfy every delivered
    /// subscriber and the resolver wants to retry the key.
    pub fn discard_response(&mut self, key: &Con::Key) {
        if let Some(entry) = self.entries.get_mut(key) {
            entry.response = None;
        }
    }

    /// Wait for the next consumer validation result.
    ///
    /// Returns [`Aborted`] when cancellation closes the receiver without an accepted
    /// verdict. Otherwise, canceled deliveries return their verdict for source
    /// accounting without affecting any current fetch for the key. Active deliveries
    /// clear their slot so the key can be retried or redelivered. A verdict sender
    /// dropped by the consumer completes with no outcome.
    pub async fn next_completion(
        &mut self,
    ) -> Result<Completion<Con::Key, Con::Subscriber, Context>, Aborted> {
        let completed = self.deliveries.next_completed().await?;

        // Cancellation belongs to the removed entry. Only an active completion can
        // clear a delivery slot; a canceled key may already have a replacement entry.
        if !completed.canceled {
            let entry = self
                .entries
                .get_mut(&completed.delivery.key)
                .expect("delivery entry");
            entry.delivery = None;
        }
        Ok(completed)
    }

    // Start a consumer validation attempt whose cancellation is owned by its entry.
    fn push_delivery(
        &mut self,
        delivery: Delivery<Con::Key, Con::Subscriber>,
        context: Context,
        value: Con::Value,
    ) {
        let entry = self.entries.get_mut(&delivery.key).expect("delivery entry");
        let previously_accepted = entry.response.as_ref().expect("response").accepted;
        let completed = delivery.clone();
        let mut consumer = self.consumer.clone();
        let mut receiver = consumer.deliver(delivery, value);
        let (cancel, canceled) = oneshot::channel::<()>();
        self.deliveries.push(async move {
            // Cancellation takes precedence over a ready verdict so feedback from a
            // removed entry can never complete a replacement. Closing before reading
            // preserves successful sends without waiting for an unfinished verdict.
            let (outcome, canceled) = select! {
                _ = canceled => {
                    receiver.close();
                    (receiver.try_recv().ok().map(Into::into), true)
                },
                outcome = &mut receiver => {
                    let outcome = match outcome {
                        Ok(outcome) => Some(outcome.into()),
                        Err(_) => {
                            debug!(key = ?completed.key, "consumer dropped delivery without a verdict");
                            None
                        }
                    };
                    (outcome, false)
                },
            };
            if canceled && outcome.is_none() {
                return Err(Aborted);
            }
            Ok(Completion {
                context,
                delivery: completed,
                outcome,
                canceled,
                previously_accepted,
            })
        });
        assert!(entry.delivery.replace(cancel).is_none());
    }
}

impl<Con, Context> Tracker<Con, Context>
where
    Con: Consumer,
    Con::Value: Clone + Send + 'static,
    Context: Clone + Send + 'static,
{
    /// Start tracking a key without any resolver-specific state.
    ///
    /// Returns `true` when the key was inserted. If the key is already tracked,
    /// this leaves the existing entry untouched and returns `false`.
    pub fn insert(&mut self, key: Con::Key) -> bool {
        self.insert_with_state(key, ())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        Resolution,
        p2p::mocks::{Consumer as MockConsumer, Key as MockKey},
    };
    use bytes::Bytes;
    use commonware_runtime::{Runner as _, deterministic::Runner};
    use commonware_utils::{
        channel::{fallible::FallibleExt, mpsc, oneshot},
        non_empty_vec,
    };
    use futures::FutureExt as _;

    type TestTracker = Tracker<MockConsumer<MockKey, Bytes>, u8>;
    type StateTracker = Tracker<MockConsumer<MockKey, Bytes>, u8, u8>;

    fn delivery(key: MockKey) -> Delivery<MockKey, ()> {
        Delivery {
            key,
            subscribers: non_empty_vec![((), tracing::Span::none())],
        }
    }

    #[derive(Clone)]
    struct PendingConsumer {
        sender: mpsc::UnboundedSender<oneshot::Sender<bool>>,
    }

    impl PendingConsumer {
        fn new() -> (Self, mpsc::UnboundedReceiver<oneshot::Sender<bool>>) {
            let (sender, receiver) = mpsc::unbounded_channel();
            (Self { sender }, receiver)
        }
    }

    impl Consumer for PendingConsumer {
        type Key = MockKey;
        type Value = Bytes;
        type Subscriber = ();
        type Outcome = bool;

        fn deliver(
            &mut self,
            _delivery: Delivery<Self::Key, Self::Subscriber>,
            _value: Self::Value,
        ) -> oneshot::Receiver<bool> {
            let (sender, receiver) = oneshot::channel();
            self.sender.send_lossy(sender);
            receiver
        }
    }

    #[test]
    fn test_insert_contains_remove_round_trip() {
        let runner = Runner::default();
        runner.start(|_| async move {
            let mut tracker = TestTracker::new(MockConsumer::dummy());

            assert!(!tracker.contains(&MockKey(1)));
            assert!(tracker.insert(MockKey(1)));
            assert!(tracker.contains(&MockKey(1)));

            assert!(!tracker.insert(MockKey(1)));
            assert!(tracker.remove(&MockKey(1)));
            assert!(!tracker.contains(&MockKey(1)));
            assert!(!tracker.remove(&MockKey(1)));
        });
    }

    #[test]
    fn test_state_round_trip() {
        let runner = Runner::default();
        runner.start(|_| async move {
            // An untracked key has no state.
            let mut tracker = StateTracker::new(MockConsumer::dummy());
            assert_eq!(tracker.state_mut(&MockKey(1)).map(|state| *state), None);

            // Inserted state can be read and updated in place.
            assert!(tracker.insert_with_state(MockKey(1), 5));
            assert_eq!(tracker.state_mut(&MockKey(1)).map(|state| *state), Some(5));
            *tracker
                .state_mut(&MockKey(1))
                .expect("inserted key must have state") = 7;
            assert_eq!(tracker.state_mut(&MockKey(1)).map(|state| *state), Some(7));

            // Taking the state leaves the key tracked without state, so removing
            // the key returns no state.
            assert_eq!(tracker.take_state(&MockKey(1)), Some(7));
            assert!(tracker.contains(&MockKey(1)));
            assert_eq!(tracker.state_mut(&MockKey(1)), None);
            assert_eq!(tracker.remove_with_state(&MockKey(1)), Some(None));
            assert_eq!(tracker.state_mut(&MockKey(1)).map(|state| *state), None);
        });
    }

    #[test]
    fn test_deliver_completes_with_context_and_consumer_result() {
        let runner = Runner::default();
        runner.start(|_| async move {
            let (consumer, mut events) = MockConsumer::<MockKey, Bytes>::new();
            let mut tracker = TestTracker::new(consumer);
            let key = MockKey(7);
            let value = Bytes::from("data");

            tracker.insert(key.clone());
            tracker.deliver(delivery(key.clone()), 9, value.clone());

            let completed = tracker
                .next_completion()
                .await
                .expect("delivery should complete");
            assert_eq!(completed.context, 9);
            assert_eq!(completed.delivery.key, key);
            assert!(!completed.canceled);
            assert!(!completed.previously_accepted);
            assert_eq!(
                completed.outcome,
                Some(Outcome::Resolved(Resolution::Complete))
            );

            let (delivered_key, delivered_value) = events.recv().await.unwrap();
            assert_eq!(delivered_key, key);
            assert_eq!(delivered_value, value);
        });
    }

    #[test]
    fn test_remove_aborts_in_flight_delivery() {
        let runner = Runner::default();
        runner.start(|_| async move {
            let (consumer, mut senders) = PendingConsumer::new();
            let mut tracker = Tracker::<_, u8>::new(consumer);
            let key = MockKey(1);

            tracker.insert(key.clone());
            tracker.deliver(delivery(key.clone()), 2, Bytes::from("v"));
            let sender = senders.recv().await.unwrap();
            assert!(tracker.remove(&key));

            assert!(matches!(tracker.next_completion().await, Err(Aborted)));
            assert!(sender.is_closed());
        });
    }

    #[test]
    fn test_canceled_delivery_cannot_complete_reinserted_key() {
        #[derive(Clone, Copy)]
        enum VerdictTiming {
            BeforeRemoval,
            AfterRemoval,
            AfterClose,
        }

        // Every removal path must preserve accepted feedback and isolate the replacement.
        for removal in 0..4 {
            for timing in [
                VerdictTiming::BeforeRemoval,
                VerdictTiming::AfterRemoval,
                VerdictTiming::AfterClose,
            ] {
                for polled in [false, true] {
                    Runner::default().start(move |_| async move {
                        // Keep the old verdict pending or make it ready without consuming it.
                        let (consumer, mut senders) = PendingConsumer::new();
                        let mut tracker = Tracker::<PendingConsumer, u8, u8>::new(consumer);
                        let key = MockKey(1);
                        assert!(tracker.insert_with_state(key.clone(), 10));
                        tracker.deliver(delivery(key.clone()), 1, Bytes::from("old"));
                        let mut old_sender = Some(senders.recv().await.unwrap());
                        if polled {
                            assert!(tracker.next_completion().now_or_never().is_none());
                        }
                        if matches!(timing, VerdictTiming::BeforeRemoval) {
                            old_sender.take().unwrap().send(true).unwrap();
                        }

                        // Reuse the key before polling the canceled delivery.
                        match removal {
                            0 => assert!(tracker.remove(&key)),
                            1 => assert_eq!(tracker.remove_with_state(&key), Some(Some(10))),
                            2 => assert_eq!(tracker.retain(|_| false), 1),
                            _ => assert_eq!(tracker.drain(), 1),
                        }
                        if matches!(timing, VerdictTiming::AfterRemoval) {
                            old_sender.take().unwrap().send(true).unwrap();
                        }
                        assert!(tracker.insert_with_state(key.clone(), 20));
                        tracker.deliver(delivery(key.clone()), 2, Bytes::from("new"));
                        let new_sender = senders.recv().await.unwrap();

                        // Cancellation closes pending verdicts and preserves accepted
                        // ones without touching the replacement's state or delivery.
                        let result = tracker.next_completion().await;
                        if matches!(timing, VerdictTiming::AfterClose) {
                            assert!(matches!(result, Err(Aborted)));
                            assert!(old_sender.take().unwrap().send(true).is_err());
                        } else {
                            let completed = result.unwrap();
                            assert!(completed.canceled);
                            assert!(!completed.previously_accepted);
                            assert_eq!(completed.context, 1);
                            assert_eq!(
                                completed.outcome,
                                Some(Outcome::Resolved(Resolution::Complete))
                            );
                        }
                        assert_eq!(tracker.state_mut(&key).map(|state| *state), Some(20));
                        assert!(tracker.entries[&key].delivery.is_some());
                        assert!(tracker.next_completion().now_or_never().is_none());

                        // Only the replacement's verdict may complete the new entry.
                        new_sender.send(true).unwrap();
                        let completed = tracker.next_completion().await.unwrap();
                        assert!(!completed.canceled);
                        assert_eq!(completed.context, 2);
                        assert_eq!(completed.delivery.key, key);
                        assert_eq!(
                            completed.outcome,
                            Some(Outcome::Resolved(Resolution::Complete))
                        );
                        assert!(tracker.deliveries.is_empty());
                    });
                }
            }
        }
    }

    #[test]
    fn test_dropped_verdict_completes_without_outcome_and_redelivers() {
        let runner = Runner::default();
        runner.start(|_| async move {
            let (consumer, mut senders) = PendingConsumer::new();
            let mut tracker = Tracker::<PendingConsumer, u8>::new(consumer);
            let key = MockKey(3);

            tracker.insert(key.clone());
            tracker.deliver(delivery(key.clone()), 4, Bytes::from("unjudged"));
            drop(senders.recv().await.unwrap());

            let completed = tracker
                .next_completion()
                .await
                .expect("dropped verdict should complete");
            assert_eq!(completed.context, 4);
            assert_eq!(completed.delivery.key, key);
            assert_eq!(completed.outcome, None);
            assert!(!tracker.response_accepted(&key));

            // The unjudged response can still be handed to other subscribers.
            tracker.redeliver(delivery(key.clone()));
            senders.recv().await.unwrap().send(true).unwrap();
            let judged = tracker
                .next_completion()
                .await
                .expect("redelivery should complete");
            assert_eq!(judged.context, 4);
            assert_eq!(
                judged.outcome,
                Some(Outcome::Resolved(Resolution::Complete))
            );
        });
    }

    #[test]
    fn test_redeliver_reuses_accepted_response_for_new_subscribers() {
        let runner = Runner::default();
        runner.start(|_| async move {
            let (consumer, mut events) = MockConsumer::<MockKey, Bytes>::new();
            let mut tracker = TestTracker::new(consumer);
            let key = MockKey(5);
            let value = Bytes::from("first");

            tracker.insert(key.clone());
            tracker.deliver(delivery(key.clone()), 3, value.clone());

            let completed = tracker
                .next_completion()
                .await
                .expect("first delivery should complete");
            assert!(!completed.previously_accepted);
            assert_eq!(
                completed.outcome,
                Some(Outcome::Resolved(Resolution::Complete))
            );
            tracker.accept_response(&key);
            assert!(tracker.response_accepted(&key));

            tracker.redeliver(delivery(key.clone()));
            let redelivered = tracker
                .next_completion()
                .await
                .expect("redelivery should complete");
            assert_eq!(redelivered.context, 3);
            assert_eq!(redelivered.delivery.key, key);
            assert!(redelivered.previously_accepted);
            assert_eq!(
                redelivered.outcome,
                Some(Outcome::Resolved(Resolution::Complete))
            );

            let first = events.recv().await.unwrap();
            let second = events.recv().await.unwrap();
            assert_eq!(first, (key.clone(), value.clone()));
            assert_eq!(second, (key, value));
        });
    }

    #[test]
    #[should_panic(expected = "response")]
    fn test_redeliver_requires_cached_response() {
        let runner = Runner::default();
        runner.start(|_| async move {
            let (consumer, _events) = MockConsumer::<MockKey, Bytes>::new();
            let mut tracker = TestTracker::new(consumer);
            let key = MockKey(7);

            tracker.insert(key.clone());
            tracker.deliver(delivery(key.clone()), 3, Bytes::from("first"));
            let completed = tracker
                .next_completion()
                .await
                .expect("first delivery should complete");
            assert_eq!(
                completed.outcome,
                Some(Outcome::Resolved(Resolution::Complete))
            );

            tracker.discard_response(&key);
            tracker.redeliver(delivery(key));
        });
    }

    #[test]
    fn test_rejected_response_can_be_discarded_and_replaced() {
        let runner = Runner::default();
        runner.start(|_| async move {
            let (mut consumer, _events) = MockConsumer::<MockKey, Bytes>::new();
            let key = MockKey(8);
            consumer.add_expected(key.clone(), Bytes::from("good"));
            let mut tracker = TestTracker::new(consumer);

            tracker.insert(key.clone());
            tracker.deliver(delivery(key.clone()), 1, Bytes::from("bad"));
            let rejected = tracker
                .next_completion()
                .await
                .expect("rejected delivery should complete");
            assert_eq!(rejected.outcome, Some(Outcome::Invalid));

            tracker.discard_response(&key);
            assert!(!tracker.response_accepted(&key));
            tracker.deliver(delivery(key.clone()), 2, Bytes::from("good"));

            let accepted = tracker
                .next_completion()
                .await
                .expect("accepted delivery should complete");
            assert_eq!(accepted.context, 2);
            assert_eq!(
                accepted.outcome,
                Some(Outcome::Resolved(Resolution::Complete))
            );
        });
    }
}