use super::{
    Producer,
    config::Config,
    fetcher::{Config as FetcherConfig, Fetcher},
    inflight::Inflight,
    ingress::{FetchKey, Mailbox, Message},
    metrics, wire,
};
use crate::{Consumer, Delivery, Outcome, Resolution, delivery::Completion, subscribers};
use bytes::Bytes;
use commonware_actor::mailbox;
use commonware_cryptography::PublicKey;
use commonware_macros::select_loop;
use commonware_p2p::{
    Blocker, Provider, Receiver, Recipients, Sender,
    utils::codec::{WrappedSender, wrap},
};
use commonware_runtime::{
    Clock, ContextCell, Handle, Metrics, Spawner, spawn_cell,
    telemetry::metrics::{GaugeExt, histogram, status::Status},
};
use commonware_utils::{Span, channel::oneshot, futures::Pool as FuturesPool};
use futures::future::{self, Either};
use rand_core::Rng;
use tracing::{debug, error, trace, warn};

/// Represents a pending serve operation.
struct Serve<P: PublicKey> {
    timer: histogram::Timer,
    peer: P,
    id: u64,
    result: Result<Bytes, oneshot::error::RecvError>,
}

/// Manages incoming and outgoing P2P requests, coordinating fetch and serve operations.
pub struct Engine<E, P, D, B, Key, Con, Pro>
where
    E: Clock + Spawner + Rng + Metrics,
    P: PublicKey,
    D: Provider<PublicKey = P>,
    B: Blocker<PublicKey = P>,
    Key: Span,
    Con: Consumer<Key = Key, Value = Bytes>,
    Pro: Producer<Key = Key>,
{
    /// Context used to spawn tasks, manage time, etc.
    context: ContextCell<E>,

    /// Produces data for incoming requests
    producer: Pro,

    /// Manages the list of peers that can be used to fetch data
    peer_provider: D,

    /// The blocker that will be used to block peers that send invalid responses
    blocker: B,

    /// Used to detect changes in the peer set
    last_peer_set_id: Option<u64>,

    /// Mailbox that makes and prunes fetches
    mailbox: mailbox::Receiver<Message<Key, P, Con::Subscriber>>,

    /// Manages outgoing fetch requests
    fetcher: Fetcher<E, P, Key>,

    /// Tracks all in-flight fetch state
    inflight: Inflight<Con, P>,

    /// Subscribers that keep each fetch alive.
    subscribers: subscribers::Tracker<Key, Con::Subscriber>,

    /// Holds futures that resolve once the `Producer` has produced the data.
    /// Once the future is resolved, the data (or an error) is sent to the peer.
    /// Has unbounded size; the number of concurrent requests should be limited
    /// by the `Producer` which may drop requests.
    serves: FuturesPool<'static, Serve<P>>,

    /// Whether responses are sent with priority over other network messages
    priority_responses: bool,

    /// Metrics for the peer actor
    metrics: metrics::Metrics,
}

impl<E, P, D, B, Key, Con, Pro> Engine<E, P, D, B, Key, Con, Pro>
where
    E: Clock + Spawner + Rng + Metrics,
    P: PublicKey,
    D: Provider<PublicKey = P>,
    B: Blocker<PublicKey = P>,
    Key: Span,
    Con: Consumer<Key = Key, Value = Bytes>,
    Pro: Producer<Key = Key>,
    Con::Subscriber: Ord,
{
    /// Creates a new `Actor` with the given configuration.
    ///
    /// Returns the actor and a mailbox to send messages to it.
    pub fn new(
        context: E,
        cfg: Config<P, D, B, Key, Con, Pro>,
    ) -> (Self, Mailbox<Key, P, Con::Subscriber>) {
        let (sender, receiver) = mailbox::new(context.child("mailbox"), cfg.mailbox_size);

        let metrics = metrics::Metrics::init(&context);
        let fetcher = Fetcher::new(
            context.child("fetcher"),
            FetcherConfig {
                me: cfg.me,
                fetch_timeout: cfg.fetch_timeout,
                retry_delay: cfg.fetch_retry_delay,
                priority_requests: cfg.priority_requests,
            },
        );
        (
            Self {
                context: ContextCell::new(context),
                producer: cfg.producer,
                peer_provider: cfg.peer_provider,
                blocker: cfg.blocker,
                last_peer_set_id: None,
                mailbox: receiver,
                fetcher,
                inflight: Inflight::new(cfg.consumer),
                subscribers: subscribers::Tracker::new(),
                serves: FuturesPool::default(),
                priority_responses: cfg.priority_responses,
                metrics,
            },
            Mailbox::new(sender),
        )
    }

    /// Runs the actor until the context is stopped.
    ///
    /// The actor will handle:
    /// - Fetching data from other peers and notifying the `Consumer`
    /// - Serving data to other peers by requesting it from the `Producer`
    pub fn start<NetS: Sender<PublicKey = P>, NetR: Receiver<PublicKey = P>>(
        mut self,
        network: (NetS, NetR),
    ) -> Handle<()> {
        spawn_cell!(self.context, self.run(network))
    }

    /// Inner run loop called by `start`.
    async fn run<NetS: Sender<PublicKey = P>, NetR: Receiver<PublicKey = P>>(
        mut self,
        network: (NetS, NetR),
    ) {
        // Wrap channel
        let (mut sender, mut receiver) = wrap((), network.0, network.1);
        let mut peer_set_subscription = self.peer_provider.subscribe().await;
        let mut blocked_subscription = Some(self.blocker.blocked());

        select_loop! {
            self.context,
            on_start => {
                // Wait for the next blocked-set update, or forever once the
                // network stops publishing them.
                let blocked_update = blocked_subscription.as_mut().map_or_else(
                    || Either::Right(future::pending()),
                    |subscription| Either::Left(subscription.recv()),
                );

                // Update metrics
                let _ = self
                    .metrics
                    .fetch_pending
                    .try_set(self.fetcher.len_pending());
                let _ = self.metrics.fetch_active.try_set(self.fetcher.len_active());
                let _ = self.metrics.serve_processing.try_set(self.serves.len());

                // Get retry deadline (if any)
                let deadline_pending = match self.fetcher.get_pending_deadline() {
                    Some(deadline) => Either::Left(self.context.sleep_until(deadline)),
                    None => Either::Right(future::pending()),
                };

                // Get requester timeout (if any)
                let deadline_active = match self.fetcher.get_active_deadline() {
                    Some(deadline) => Either::Left(self.context.sleep_until(deadline)),
                    None => Either::Right(future::pending()),
                };
            },
            on_stopped => {
                debug!("shutdown");
                self.inflight.drain();
                self.subscribers.clear();
                self.serves.cancel_all();
            },
            // Handle peer set updates
            Some(update) = peer_set_subscription.recv() else {
                debug!("peer set subscription closed");
                return;
            } => {
                if self.last_peer_set_id < Some(update.index) {
                    self.last_peer_set_id = Some(update.index);
                    self.fetcher.reconcile(update.latest.union().as_ref());
                }
            },
            // Handle blocked-set updates
            blocked = blocked_update => {
                match blocked {
                    Some(blocked) => self.fetcher.set_blocked(blocked),
                    None => {
                        debug!("blocked subscription closed");
                        blocked_subscription = None;
                    }
                }
            },
            // Handle active deadline
            _ = deadline_active => {
                if let Some(key) = self.fetcher.pop_active() {
                    debug!(?key, "requester timeout");
                    self.metrics.fetch.inc(Status::Failure);
                    self.fetcher.add_retry(key);
                }
            },
            // Poll delivery completions before mailbox work so a ready completion
            // can retire its subscribers before a queued fetch is coalesced with them.
            delivery = self.inflight.next_delivery() => {
                if let Ok(completed) = delivery {
                    self.handle_delivery(completed);
                }
            },
            // Handle mailbox messages
            Some(msg) = self.mailbox.recv() else {
                error!("mailbox closed");
                return;
            } => {
                match msg {
                    Message::Fetch(keys) => {
                        for FetchKey {
                            key,
                            subscribers,
                            metadata: targets,
                        } in keys
                        {
                            trace!(?key, "mailbox: fetch");

                            // Check if the fetch is already in progress
                            let is_new = !self.inflight.contains(&key);
                            self.subscribers.insert(key.clone(), subscribers);

                            // Update targets
                            match targets {
                                Some(targets) => {
                                    // Only add targets if this is a new fetch OR the existing
                                    // fetch already has targets. Don't restrict an "all" fetch
                                    // (no targets) to specific targets.
                                    if is_new || self.fetcher.has_targets(&key) {
                                        self.fetcher.add_targets(key.clone(), targets);
                                    }
                                }
                                None => self.fetcher.clear_targets(&key),
                            }

                            // Only start new fetch if not already in progress
                            if is_new {
                                self.inflight.insert(
                                    key.clone(),
                                    self.metrics.fetch_duration.timer(self.context.as_ref()),
                                );
                                self.fetcher.add_ready(key);
                            } else {
                                trace!(?key, "updated targets for existing fetch");
                            }
                        }
                    }
                    Message::Retain { predicate } => {
                        trace!("mailbox: retain");

                        self.subscribers
                            .retain(|key, subscriber| predicate(key, subscriber));
                        let subscribers = &self.subscribers;
                        self.fetcher.retain(|key| subscribers.contains(key));
                        let count = self.inflight.retain(|key| subscribers.contains(key)) as u64;
                        self.record_cancellations(count);
                    }
                }
            },
            // Wake the loop when pending work becomes ready. The send is
            // performed in `on_end` after the selected event is handled.
            _ = deadline_pending => {},
            // Handle completed server requests
            serve = self.serves.next_completed() => {
                let Serve {
                    timer,
                    peer,
                    id,
                    result,
                } = serve;

                // Metrics and logs
                match result {
                    Ok(_) => {
                        timer.observe(self.context.as_ref());
                        self.metrics.serve.inc(Status::Success);
                    }
                    Err(ref err) => {
                        debug!(?err, ?peer, ?id, "serve failed");
                        self.metrics.serve.inc(Status::Failure);
                    }
                }

                // Send response to peer
                self.handle_serve(&mut sender, peer, id, result, self.priority_responses);
            },
            // Handle network messages
            msg = receiver.recv() => {
                // Break if the receiver is closed
                let (peer, msg) = match msg {
                    Ok(msg) => msg,
                    Err(err) => {
                        error!(?err, "receiver closed");
                        return;
                    }
                };

                match msg {
                    Ok(msg) => match msg.payload {
                        wire::Payload::Request(key) => {
                            self.handle_network_request(peer, msg.id, key)
                        }
                        wire::Payload::Response(response) => {
                            self.handle_network_response(peer, msg.id, response)
                        }
                        wire::Payload::Error => self.handle_network_error_response(peer, msg.id),
                    },
                    Err(err) => {
                        trace!(?err, ?peer, "decode failed");
                    }
                };
            },
            on_end => {
                // Attempt at most one due outbound request after each selected
                // event so sustained event traffic cannot starve pending work.
                if self
                    .fetcher
                    .get_pending_deadline()
                    .is_some_and(|deadline| deadline <= self.context.current())
                {
                    self.fetcher.fetch(&mut sender);
                }
            },
        }
    }

    /// Record cancellation metrics for a retain-style operation.
    fn record_cancellations(&mut self, count: u64) {
        if count == 0 {
            self.metrics.cancel.inc(Status::Dropped);
        } else {
            self.metrics.cancel.inc_by(Status::Success, count);
        }
    }

    /// Handles the case where the application responds to a request from an external peer.
    fn handle_serve<NetS: Sender<PublicKey = P>>(
        &mut self,
        sender: &mut WrappedSender<NetS, wire::Message<Key>>,
        peer: P,
        id: u64,
        response: Result<Bytes, oneshot::error::RecvError>,
        priority: bool,
    ) {
        // Encode message
        let payload: wire::Payload<Key> = response.map_or_else(
            |_| wire::Payload::Error,
            |data| wire::Payload::Response(data),
        );
        let msg = wire::Message { id, payload };

        // Send message to peer
        let result = sender.send(Recipients::One(peer.clone()), msg, priority);

        // Log result, but do not handle errors.
        if result.is_empty() {
            warn!(?peer, ?id, "serve send failed");
        } else {
            trace!(?peer, ?id, "serve sent");
        };
    }

    /// Handle a network request from a peer.
    fn handle_network_request(&mut self, peer: P, id: u64, key: Key) {
        // Serve the request
        trace!(?peer, ?id, "peer request");
        let mut producer = self.producer.clone();
        let timer = self.metrics.serve_duration.timer(self.context.as_ref());
        let receiver = producer.produce(key);
        self.serves.push(async move {
            let result = receiver.await;
            Serve {
                timer,
                peer,
                id,
                result,
            }
        });
    }

    /// Handle a network response from a peer.
    fn handle_network_response(&mut self, peer: P, id: u64, response: Bytes) {
        trace!(?peer, ?id, "peer response: data");

        // Get the key associated with the response, if any
        let Some((key, elapsed)) = self.fetcher.pop_response(id, &peer) else {
            // It's possible that the key does not exist if the request was pruned.
            return;
        };

        let Some(subscribers) = self.subscribers.pending(&key) else {
            warn!(?key, "response for fetch with no subscribers");
            self.inflight.cancel(&key);
            return;
        };
        let delivery = Delivery { key, subscribers };

        // The peer had the data, so deliver it to the consumer without blocking the engine.
        self.inflight.deliver(delivery, peer, elapsed, response);
    }

    /// Handle completed delivery to the consumer.
    fn handle_delivery(
        &mut self,
        completed: Completion<Key, Con::Subscriber, (P, std::time::Duration, usize)>,
    ) {
        let Completion {
            context: (peer, elapsed, bytes),
            delivery:
                Delivery {
                    key,
                    subscribers: delivered,
                },
            outcome,
            canceled,
            previously_accepted: already_accepted,
        } = completed;

        // Source feedback belongs to the delivered response even if its fetch was
        // canceled. Local redelivery of an accepted response cannot earn more credit.
        match outcome {
            Some(Outcome::Resolved(_)) if !already_accepted => {
                self.fetcher.record_response(&peer, elapsed, bytes);
                self.metrics.fetch.inc(Status::Success);
            }
            Some(Outcome::Unresolved) if !already_accepted => {
                self.metrics.fetch.inc(Status::Ambiguous);
            }
            Some(Outcome::Invalid) => {
                // Conflicting verdicts on an accepted response reflect local consumer
                // state. Only the first validation can identify invalid peer data.
                if already_accepted {
                    warn!(
                        ?key,
                        "previously accepted response was rejected during local redelivery"
                    );
                } else {
                    commonware_p2p::block!(self.blocker, peer, "invalid data received");
                }
                self.metrics.fetch.inc(Status::Failure);
            }
            _ => {}
        }

        // Canceled feedback cannot retire, retry, or redeliver a replacement fetch.
        if canceled {
            return;
        }

        // A dropped verdict says nothing about the response, only that the consumer
        // did not judge it for these subscribers. Hand the response to the
        // remaining subscribers, or retire the key when none remain.
        let Some(outcome) = outcome else {
            let remaining = self
                .subscribers
                .remove_delivered(&key, delivered.map_into(|(subscriber, _)| subscriber));
            if let Some(subscribers) = remaining {
                self.inflight.redeliver(Delivery { key, subscribers });
                return;
            }
            if !already_accepted {
                self.metrics.fetch.inc(Status::Dropped);
            }
            self.inflight.cancel(&key);
            self.fetcher.clear(&key);
            return;
        };

        // Resolution reports demand within the delivered subscriber snapshot.
        // Ignored cancels demand for the entire key, including later subscribers.
        match outcome {
            Outcome::Resolved(Resolution::Complete) => {
                // Remove only the subscribers that accepted this response. If other
                // subscribers still need the key, deliver the same accepted response
                // locally with the remaining annotations.
                let remaining = self
                    .subscribers
                    .remove_delivered(&key, delivered.map_into(|(subscriber, _)| subscriber));

                if let Some(subscribers) = remaining {
                    if !already_accepted {
                        self.inflight.accept_response(&key, self.context.as_ref());
                    }
                    self.inflight.redeliver(Delivery { key, subscribers });
                } else {
                    // All subscribers observed a valid response, so drop the key's targets
                    // and peer-selection history.
                    self.inflight.complete(self.context.as_ref(), &key);
                    self.fetcher.clear(&key);
                }
            }
            Outcome::Resolved(Resolution::Partial) | Outcome::Unresolved => {
                // Valid data can leave subscribers unresolved. Discard it and retry
                // without penalizing the peer or recounting an accepted response.
                self.inflight.discard_response(&key);
                self.fetcher.add_retry(key);
            }
            Outcome::Invalid => {
                // An accepted response remains valid peer data. Retire conflicting
                // local demand without fetching the same response again.
                if already_accepted {
                    self.inflight.complete(self.context.as_ref(), &key);
                    self.subscribers.remove(&key);
                    self.fetcher.clear(&key);
                    return;
                }

                // Subscribers still need valid data, so discard the rejected response
                // and schedule another attempt.
                self.inflight.discard_response(&key);
                self.fetcher.add_retry(key);
            }
            Outcome::Ignored => {
                // The consumer no longer needs the key. Retire the entire fetch without
                // scoring or blocking the response's source.
                self.metrics.fetch.inc(Status::Dropped);
                self.inflight.cancel(&key);
                self.subscribers.remove(&key);
                self.fetcher.clear(&key);
            }
        }
    }

    /// Handle a network response from a peer that did not have the data.
    fn handle_network_error_response(&mut self, peer: P, id: u64) {
        trace!(?peer, ?id, "peer response: error");

        // Get the key associated with the response, if any
        let Some(key) = self.fetcher.pop_missing(id, &peer) else {
            // It's possible that the key does not exist if the request was pruned.
            return;
        };

        // The peer did not have the data, so we need to try again
        self.metrics.fetch.inc(Status::Failure);
        self.fetcher.add_retry(key);
    }
}