apcore 0.22.0

// APCore Protocol — Middleware manager
// Spec reference: Middleware pipeline execution

use std::collections::HashMap;
use std::sync::Arc;

use parking_lot::Mutex;

use super::base::Middleware;
use crate::context::Context;
use crate::errors::{ErrorCode, ModuleError};

/// Options for registering a middleware with identity-based duplicate detection.
pub struct MiddlewareRegistration<M: Middleware + 'static> {
    /// The middleware instance to register.
    pub middleware: M,
    /// When `true`, skips the duplicate-identity warning entirely.
    /// Default: `false`.
    pub allow_duplicate: bool,
    /// Override the identity key used for duplicate detection.
    /// Default: `std::any::type_name::<M>()`.
    pub identity_key: Option<String>,
}

impl<M: Middleware + 'static> MiddlewareRegistration<M> {
    /// Create a new registration with default options (duplicate warning enabled,
    /// identity derived from the type name).
    pub fn new(middleware: M) -> Self {
        Self {
            middleware,
            allow_duplicate: false,
            identity_key: None,
        }
    }

    /// Set whether duplicate-identity warnings are suppressed.
    #[must_use]
    pub fn allow_duplicate(mut self, v: bool) -> Self {
        self.allow_duplicate = v;
        self
    }

    /// Override the identity key used for duplicate detection.
    #[must_use]
    pub fn identity_key(mut self, k: impl Into<String>) -> Self {
        self.identity_key = Some(k.into());
        self
    }
}

/// Manages an ordered pipeline of middleware.
///
/// Thread-safe: internal list is protected by a Mutex. All execution methods
/// clone the Arc pointers out of the mutex before iterating so the lock is
/// never held across `.await` points.
#[derive(Debug)]
pub struct MiddlewareManager {
    middlewares: Mutex<Vec<Arc<dyn Middleware>>>,
    /// Tracks identity -> first registration location hint for duplicate detection.
    registered_identities: Mutex<HashMap<String, String>>,
}

impl MiddlewareManager {
    /// Create a new empty middleware manager.
    #[must_use]
    pub fn new() -> Self {
        Self {
            middlewares: Mutex::new(vec![]),
            registered_identities: Mutex::new(HashMap::new()),
        }
    }

    /// Add a middleware to the pipeline.
    ///
    /// Middlewares are maintained in sorted order by priority (higher first).
    /// Among middlewares with the same priority, registration order is preserved.
    ///
    /// Returns an error if the middleware's priority exceeds 1000 (the spec-defined
    /// maximum from section 11.2).
    ///
    /// Takes `&self` — the internal list is protected by a `Mutex`, so mutation
    /// is possible through a shared reference. This allows `MiddlewareManager`
    /// to be held as `Arc<MiddlewareManager>` and mutated without `Arc::get_mut`
    /// hacks, even after the `Arc` has been cloned into pipeline contexts.
    pub fn add(&self, middleware: Box<dyn Middleware>) -> Result<(), ModuleError> {
        let priority = middleware.priority();
        if priority > 1000 {
            tracing::warn!(
                middleware = middleware.name(),
                priority = priority,
                "Middleware rejected: priority {} exceeds maximum 1000",
                priority,
            );
            return Err(ModuleError::new(
                ErrorCode::GeneralInvalidInput,
                format!(
                    "Middleware '{}' has priority {} which exceeds the maximum allowed value of 1000",
                    middleware.name(),
                    priority,
                ),
            ));
        }
        let mut mws = self.middlewares.lock();
        let arc: Arc<dyn Middleware> = Arc::from(middleware);
        // Find the first position where existing priority is strictly less than
        // the new priority. Insert before that position to maintain stable
        // ordering (later registrations go after earlier ones at same priority).
        let pos = mws
            .iter()
            .position(|m| m.priority() < priority)
            .unwrap_or(mws.len());
        mws.insert(pos, arc);
        Ok(())
    }

    /// Register a middleware with duplicate-detection options.
    ///
    /// If the same identity (type name or explicit key) is registered twice and
    /// `allow_duplicate` is `false`, a `tracing::warn!` is emitted. Registration
    /// always succeeds; the chain is not otherwise modified.
    ///
    /// The identity is captured generically via `std::any::type_name::<M>()` so
    /// two distinct instances of the same type are treated as duplicates by default.
    /// Use `identity_key` to distinguish them (e.g. `"auth-primary"` vs
    /// `"auth-secondary"`).
    #[track_caller]
    pub fn add_with_opts<M: Middleware + 'static>(
        &self,
        opts: MiddlewareRegistration<M>,
    ) -> Result<(), ModuleError> {
        let location = std::panic::Location::caller().to_string();
        let identity = opts
            .identity_key
            .clone()
            .unwrap_or_else(|| std::any::type_name::<M>().to_string());

        if !opts.allow_duplicate {
            let mut ids = self.registered_identities.lock();
            if let Some(first_site) = ids.get(&identity) {
                tracing::warn!(
                    identity = %identity,
                    first_registration = %first_site,
                    duplicate_registration = %location,
                    "duplicate middleware registration detected"
                );
            } else {
                ids.insert(identity, location);
            }
        }

        self.add(Box::new(opts.middleware))
    }

    /// Remove the first middleware whose `name()` matches `name`.
    ///
    /// **Name-based, NOT identity-based.** Unlike Python's `id()` /
    /// reference-equality removal, this method matches on the
    /// `Middleware::name()` string. Two middlewares registered with the
    /// same name **cannot be removed independently** — disambiguate them
    /// at registration time by giving each a unique name (sync finding
    /// A-D-401).
    ///
    /// Removes only the first match in pipeline order; returns `false` if
    /// no middleware with that name is registered.
    ///
    /// Takes `&self` — mutation goes through the internal `Mutex`.
    pub fn remove(&self, name: &str) -> bool {
        let mut mws = self.middlewares.lock();
        let pos = mws.iter().position(|m| m.name() == name);
        if let Some(i) = pos {
            mws.remove(i);
            true
        } else {
            false
        }
    }

    /// Return a snapshot of middleware names in pipeline order.
    pub fn snapshot(&self) -> Vec<String> {
        let mws = self.middlewares.lock();
        mws.iter().map(|m| m.name().to_string()).collect()
    }

    /// Run the before hooks for all middlewares in order.
    ///
    /// Returns the (possibly modified) input and the list of indices of
    /// middlewares that were successfully executed (used by `execute_on_error`
    /// for onion-model unwinding).
    ///
    /// If a middleware's `before()` fails, wraps the error as a
    /// `MiddlewareChainError` and returns it. The `executed` list allows
    /// callers to run `on_error` only on middlewares that actually ran.
    pub async fn execute_before(
        &self,
        module_id: &str,
        mut inputs: serde_json::Value,
        ctx: &Context<serde_json::Value>,
    ) -> Result<(serde_json::Value, Vec<usize>), ModuleError> {
        // Clone Arc pointers so the mutex is not held across .await points.
        let mws: Vec<Arc<dyn Middleware>> = {
            let guard = self.middlewares.lock();
            guard.iter().map(Arc::clone).collect()
        };
        let mut executed: Vec<usize> = Vec::new();
        for (i, mw) in mws.iter().enumerate() {
            // Push BEFORE calling before() so the failing middleware is included in
            // executed — enabling on_error self-heal (mirrors Python/TS behaviour).
            executed.push(i);
            match mw.before(module_id, inputs.clone(), ctx).await {
                Ok(Some(modified)) => {
                    inputs = modified;
                }
                Ok(None) => {
                    // No modification — keep current inputs
                }
                Err(e) => {
                    // Wrap as MiddlewareChainError but preserve the original
                    // typed error in details["inner_error"] so callers can
                    // recover it via ModuleError::unwrap_middleware_chain_error()
                    // (sync finding A-D-015 — match Python's
                    // `MiddlewareChainError.original` and TypeScript's
                    // `MiddlewareChainError.original`).
                    let mut details = std::collections::HashMap::new();
                    if let Ok(inner_json) = serde_json::to_value(&e) {
                        details.insert("inner_error".to_string(), inner_json);
                    }
                    return Err(ModuleError::new(
                        ErrorCode::MiddlewareChainError,
                        e.message.clone(),
                    )
                    .with_details(details)
                    .with_cause(format!(
                        "Middleware '{}' before() failed: {}",
                        mw.name(),
                        e
                    )));
                }
            }
        }
        Ok((inputs, executed))
    }

    /// Run the after hooks for all middlewares in reverse order (onion model).
    ///
    /// If a middleware returns `Some(value)`, the output is replaced.
    /// If it returns `None`, the current output is kept unchanged.
    pub async fn execute_after(
        &self,
        module_id: &str,
        inputs: serde_json::Value,
        mut output: serde_json::Value,
        ctx: &Context<serde_json::Value>,
    ) -> Result<serde_json::Value, ModuleError> {
        // Clone Arc pointers so the mutex is not held across .await points.
        let mws: Vec<Arc<dyn Middleware>> = {
            let guard = self.middlewares.lock();
            guard.iter().map(Arc::clone).collect()
        };
        for mw in mws.iter().rev() {
            if let Some(modified) = mw
                .after(module_id, inputs.clone(), output.clone(), ctx)
                .await?
            {
                output = modified;
            } else {
                // No modification — keep current output
            }
        }
        Ok(output)
    }

    /// Run the `on_error` hooks in reverse order over the middlewares that
    /// were executed during `execute_before`.
    ///
    /// Returns `Some(recovery)` if any middleware provides a recovery value.
    /// Returns `None` if no handler recovers.
    ///
    /// Early-returns on the first recovery value — subsequent middlewares'
    /// `on_error` hooks are NOT called after recovery (mirrors Python/TS).
    pub async fn execute_on_error(
        &self,
        module_id: &str,
        inputs: serde_json::Value,
        error: &ModuleError,
        ctx: &Context<serde_json::Value>,
        executed: &[usize],
    ) -> Option<serde_json::Value> {
        // Backward-compatible variant: only returns recovery values, drops
        // any RetrySignal request. Use [`Self::execute_on_error_outcome`] to
        // observe RetrySignals from middleware (sync finding A-D-017).
        match self
            .execute_on_error_outcome(module_id, inputs, error, ctx, executed)
            .await
        {
            Some(crate::middleware::base::OnErrorOutcome::Recovery(v)) => Some(v),
            Some(crate::middleware::base::OnErrorOutcome::Retry(_)) | None => None,
        }
    }

    /// Run the `on_error_outcome` hooks in reverse order. Returns the first
    /// non-None outcome (Recovery or Retry) — Recovery becomes the call's
    /// return value; Retry tells the executor to re-run the pipeline.
    pub async fn execute_on_error_outcome(
        &self,
        module_id: &str,
        inputs: serde_json::Value,
        error: &ModuleError,
        ctx: &Context<serde_json::Value>,
        executed: &[usize],
    ) -> Option<crate::middleware::base::OnErrorOutcome> {
        // Clone Arc pointers so the mutex is not held across .await points.
        let mws: Vec<Arc<dyn Middleware>> = {
            let guard = self.middlewares.lock();
            guard.iter().map(Arc::clone).collect()
        };

        for &i in executed.iter().rev() {
            if let Some(mw) = mws.get(i) {
                match mw
                    .on_error_outcome(module_id, inputs.clone(), error, ctx)
                    .await
                {
                    Ok(Some(outcome)) => {
                        // Early-return on first non-None outcome (matches Python/TS).
                        return Some(outcome);
                    }
                    Ok(None) => {
                        // No recovery / retry from this middleware — continue
                    }
                    Err(e) => {
                        tracing::error!("Middleware '{}' on_error failed: {}", mw.name(), e);
                    }
                }
            }
        }

        None
    }
}

impl Default for MiddlewareManager {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::context::Context;
    use crate::errors::ModuleError;
    use async_trait::async_trait;

    /// Simple test middleware with configurable name and priority.
    #[derive(Debug)]
    struct TestMiddleware {
        mw_name: String,
        mw_priority: u16,
    }

    impl TestMiddleware {
        fn new(name: &str, priority: u16) -> Self {
            Self {
                mw_name: name.to_string(),
                mw_priority: priority,
            }
        }

        fn boxed(name: &str, priority: u16) -> Box<Self> {
            Box::new(Self::new(name, priority))
        }
    }

    #[async_trait]
    impl Middleware for TestMiddleware {
        fn name(&self) -> &str {
            &self.mw_name
        }

        fn priority(&self) -> u16 {
            self.mw_priority
        }

        async fn before(
            &self,
            _module_id: &str,
            _inputs: serde_json::Value,
            _ctx: &Context<serde_json::Value>,
        ) -> Result<Option<serde_json::Value>, ModuleError> {
            Ok(None)
        }

        async fn after(
            &self,
            _module_id: &str,
            _inputs: serde_json::Value,
            _output: serde_json::Value,
            _ctx: &Context<serde_json::Value>,
        ) -> Result<Option<serde_json::Value>, ModuleError> {
            Ok(None)
        }

        async fn on_error(
            &self,
            _module_id: &str,
            _inputs: serde_json::Value,
            _error: &ModuleError,
            _ctx: &Context<serde_json::Value>,
        ) -> Result<Option<serde_json::Value>, ModuleError> {
            Ok(None)
        }
    }

    #[test]
    fn test_higher_priority_executes_first_in_before() {
        let mgr = MiddlewareManager::new();
        mgr.add(TestMiddleware::boxed("low", 1)).unwrap();
        mgr.add(TestMiddleware::boxed("high", 10)).unwrap();
        mgr.add(TestMiddleware::boxed("mid", 5)).unwrap();

        let names = mgr.snapshot();
        assert_eq!(names, vec!["high", "mid", "low"]);
    }

    #[test]
    fn test_equal_priority_preserves_registration_order() {
        let mgr = MiddlewareManager::new();
        mgr.add(TestMiddleware::boxed("first", 5)).unwrap();
        mgr.add(TestMiddleware::boxed("second", 5)).unwrap();
        mgr.add(TestMiddleware::boxed("third", 5)).unwrap();

        let names = mgr.snapshot();
        assert_eq!(names, vec!["first", "second", "third"]);
    }

    #[test]
    fn test_default_priority_orders_after_explicit_priority() {
        let mgr = MiddlewareManager::new();
        mgr.add(TestMiddleware::boxed("default_a", 0)).unwrap();
        mgr.add(TestMiddleware::boxed("explicit", 1)).unwrap();
        mgr.add(TestMiddleware::boxed("default_b", 0)).unwrap();

        let names = mgr.snapshot();
        assert_eq!(names, vec!["explicit", "default_a", "default_b"]);
    }

    #[test]
    fn test_snapshot_reflects_priority_sorted_order() {
        let mgr = MiddlewareManager::new();
        mgr.add(TestMiddleware::boxed("d", 0)).unwrap();
        mgr.add(TestMiddleware::boxed("a", 100)).unwrap();
        mgr.add(TestMiddleware::boxed("c", 5)).unwrap();
        mgr.add(TestMiddleware::boxed("b", 50)).unwrap();

        let names = mgr.snapshot();
        assert_eq!(names, vec!["a", "b", "c", "d"]);
    }

    #[test]
    fn test_add_rejects_priority_above_1000() {
        let mgr = MiddlewareManager::new();
        let result = mgr.add(TestMiddleware::boxed("over_limit", 1001));
        assert!(result.is_err());
        let err = result.unwrap_err();
        assert!(
            err.message.contains("exceeds the maximum"),
            "Expected error about priority limit, got: {}",
            err.message,
        );
        // Pipeline should be empty — the middleware was not added.
        assert!(mgr.snapshot().is_empty());
    }

    #[test]
    fn test_add_accepts_priority_at_1000() {
        let mgr = MiddlewareManager::new();
        mgr.add(TestMiddleware::boxed("at_limit", 1000)).unwrap();
        assert_eq!(mgr.snapshot(), vec!["at_limit"]);
    }
}