//! Economic throttling middleware (The Meter).
//!
//! This middleware enforces per-agent per-hour quotas using a Token Bucket algorithm.
//! It extracts the agent ID from the `X-Agent-Id` header, determines the operation
//! type from the request path, and checks/records costs against the quota store.
//!
//! # Request Flow
//!
//! 1. Extract `X-Agent-Id` header (hex-encoded 32-byte public key)
//! 2. Determine operation type from request path (`/v1/assert`, `/v1/vote`, `/v1/query`)
//! 3. Calculate cost based on operation + payload size
//! 4. Check quota and reject with 429 if exceeded
//! 5. Record cost and forward request
//! 6. Add `X-Quota-Remaining` header to response
//!
//! # Headers
//!
//! | Header | Direction | Description |
//! |--------|-----------|-------------|
//! | `X-Agent-Id` | Request | Agent's Ed25519 public key (hex, 64 chars) |
//! | `X-Quota-Remaining` | Response | Remaining tokens in current window |
//! | `X-Quota-Limit` | Response | Total tokens allowed per hour |
//! | `X-Quota-Reset` | Response | Unix timestamp when quota resets |

use axum::{
    body::Body,
    http::{Request, Response, StatusCode},
    response::IntoResponse,
    Json,
};
use futures::future::BoxFuture;
use serde::Serialize;
use std::sync::Arc;
use std::task::{Context, Poll};
use stemedb_storage::{OperationType, QuotaStore};
use tower::{Layer, Service};
use tracing::{debug, warn};

/// Header name for agent identification.
pub const AGENT_ID_HEADER: &str = "x-agent-id";

/// Header name for remaining quota.
pub const QUOTA_REMAINING_HEADER: &str = "x-quota-remaining";

/// Header name for quota limit.
pub const QUOTA_LIMIT_HEADER: &str = "x-quota-limit";

/// Header name for quota reset timestamp.
pub const QUOTA_RESET_HEADER: &str = "x-quota-reset";

/// Error response for quota exceeded.
#[derive(Debug, Serialize)]
struct QuotaExceededError {
    error: String,
    code: String,
    remaining: u64,
    limit: u64,
    reset_at: u64,
}

/// Tower Layer for economic throttling.
///
/// Wrap your router with this layer to enable quota enforcement.
/// The cost configuration is determined by the `QuotaStore`.
///
/// # Example
///
/// ```ignore
/// let meter_layer = MeterLayer::new(quota_store);
/// let app = Router::new()
///     .route("/v1/assert", post(create_assertion))
///     .layer(meter_layer);
/// ```
#[derive(Clone)]
pub struct MeterLayer<Q> {
    quota_store: Arc<Q>,
    /// Paths that bypass metering (e.g., health checks).
    bypass_paths: Vec<String>,
}

impl<Q> MeterLayer<Q> {
    /// Create a new MeterLayer.
    ///
    /// The cost configuration is determined by the `QuotaStore`.
    pub fn new(quota_store: Arc<Q>) -> Self {
        Self {
            quota_store,
            bypass_paths: vec![
                "/v1/health".to_string(),
                "/swagger-ui".to_string(),
                "/api-docs".to_string(),
            ],
        }
    }

    /// Add a path to bypass metering.
    pub fn bypass_path(mut self, path: impl Into<String>) -> Self {
        self.bypass_paths.push(path.into());
        self
    }
}

impl<S, Q> Layer<S> for MeterLayer<Q>
where
    Q: Clone,
{
    type Service = MeterService<S, Q>;

    fn layer(&self, inner: S) -> Self::Service {
        MeterService {
            inner,
            quota_store: Arc::clone(&self.quota_store),
            bypass_paths: self.bypass_paths.clone(),
        }
    }
}

/// Tower Service for economic throttling.
#[derive(Clone)]
pub struct MeterService<S, Q> {
    inner: S,
    quota_store: Arc<Q>,
    bypass_paths: Vec<String>,
}

impl<S, Q> MeterService<S, Q> {
    /// Determine operation type from request path.
    fn operation_from_path(path: &str) -> Option<OperationType> {
        if path.starts_with("/v1/assert") {
            Some(OperationType::Assert)
        } else if path.starts_with("/v1/vote") {
            Some(OperationType::Vote)
        } else if path.starts_with("/v1/query") || path.starts_with("/v1/skeptic") {
            // Count lenses: if path has "lens" param or it's skeptic (which uses Skeptic lens)
            let lens_count = if path.starts_with("/v1/skeptic") { 1 } else { 0 };
            Some(OperationType::Query { lens_count })
        } else if path.starts_with("/v1/audit") {
            // Audit queries are free (internal debugging)
            None
        } else {
            None
        }
    }

    /// Check if path should bypass metering.
    #[allow(dead_code)] // Used in tests, may be used directly in the future
    fn should_bypass(&self, path: &str) -> bool {
        self.bypass_paths.iter().any(|p| path.starts_with(p))
    }

    /// Extract agent ID from request headers.
    fn extract_agent_id(req: &Request<Body>) -> Option<[u8; 32]> {
        req.headers().get(AGENT_ID_HEADER).and_then(|v| v.to_str().ok()).and_then(|s| {
            let bytes = hex::decode(s).ok()?;
            if bytes.len() == 32 {
                let mut arr = [0u8; 32];
                arr.copy_from_slice(&bytes);
                Some(arr)
            } else {
                None
            }
        })
    }
}

impl<S, Q> Service<Request<Body>> for MeterService<S, Q>
where
    S: Service<Request<Body>, Response = Response<Body>> + Clone + Send + 'static,
    S::Future: Send,
    Q: QuotaStore + 'static,
{
    type Response = Response<Body>;
    type Error = S::Error;
    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, req: Request<Body>) -> Self::Future {
        let path = req.uri().path().to_string();
        let quota_store = Arc::clone(&self.quota_store);
        let bypass_paths = self.bypass_paths.clone();

        // Clone the inner service for the async block
        let mut inner = self.inner.clone();

        Box::pin(async move {
            // Check if this path should bypass metering
            if bypass_paths.iter().any(|p| path.starts_with(p)) {
                debug!(path = %path, "Bypassing meter for path");
                return inner.call(req).await;
            }

            // Determine operation type
            let operation = match Self::operation_from_path(&path) {
                Some(op) => op,
                None => {
                    // Unknown operation, pass through without metering
                    debug!(path = %path, "No metering for path");
                    return inner.call(req).await;
                }
            };

            // Extract agent ID
            let agent_id = match Self::extract_agent_id(&req) {
                Some(id) => id,
                None => {
                    // No agent ID provided, pass through (anonymous queries allowed)
                    debug!(path = %path, "No agent ID, allowing anonymous request");
                    return inner.call(req).await;
                }
            };

            // Get content length for payload-based costing
            let content_length = req
                .headers()
                .get(axum::http::header::CONTENT_LENGTH)
                .and_then(|v| v.to_str().ok())
                .and_then(|s| s.parse::<usize>().ok())
                .unwrap_or(0);

            // Get current timestamp
            let timestamp = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .map(|d| d.as_secs())
                .unwrap_or(0);

            // Check and record quota
            let quota_result = match quota_store
                .check_and_record(&agent_id, operation, content_length, timestamp)
                .await
            {
                Ok(result) => result,
                Err(e) => {
                    warn!(error = %e, "Quota check failed, allowing request");
                    // On error, allow the request (fail open for availability)
                    return inner.call(req).await;
                }
            };

            if !quota_result.allowed {
                // Quota exceeded - return 429
                let error = QuotaExceededError {
                    error: format!(
                        "Quota exceeded. You have {} tokens remaining, but this operation costs {}.",
                        quota_result.remaining, quota_result.cost
                    ),
                    code: "QUOTA_EXCEEDED".to_string(),
                    remaining: quota_result.remaining,
                    limit: quota_result.limit,
                    reset_at: quota_result.reset_at,
                };

                let mut response = (StatusCode::TOO_MANY_REQUESTS, Json(error)).into_response();

                // Add quota headers
                let headers = response.headers_mut();
                if let Ok(v) = quota_result.remaining.to_string().parse() {
                    headers.insert(QUOTA_REMAINING_HEADER, v);
                }
                if let Ok(v) = quota_result.limit.to_string().parse() {
                    headers.insert(QUOTA_LIMIT_HEADER, v);
                }
                if let Ok(v) = quota_result.reset_at.to_string().parse() {
                    headers.insert(QUOTA_RESET_HEADER, v);
                }

                return Ok(response);
            }

            // Quota OK - call inner service
            let mut response = inner.call(req).await?;

            // Add quota headers to response
            let headers = response.headers_mut();
            if let Ok(v) = quota_result.remaining.to_string().parse() {
                headers.insert(QUOTA_REMAINING_HEADER, v);
            }
            if let Ok(v) = quota_result.limit.to_string().parse() {
                headers.insert(QUOTA_LIMIT_HEADER, v);
            }
            if let Ok(v) = quota_result.reset_at.to_string().parse() {
                headers.insert(QUOTA_RESET_HEADER, v);
            }

            Ok(response)
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_operation_from_path() {
        assert!(matches!(
            MeterService::<(), ()>::operation_from_path("/v1/assert"),
            Some(OperationType::Assert)
        ));

        assert!(matches!(
            MeterService::<(), ()>::operation_from_path("/v1/vote"),
            Some(OperationType::Vote)
        ));

        assert!(matches!(
            MeterService::<(), ()>::operation_from_path("/v1/query"),
            Some(OperationType::Query { lens_count: 0 })
        ));

        assert!(matches!(
            MeterService::<(), ()>::operation_from_path("/v1/skeptic"),
            Some(OperationType::Query { lens_count: 1 })
        ));

        assert!(MeterService::<(), ()>::operation_from_path("/v1/health").is_none());
        assert!(MeterService::<(), ()>::operation_from_path("/v1/audit/queries").is_none());
    }

    #[test]
    fn test_bypass_paths() {
        let service = MeterService::<(), ()> {
            inner: (),
            quota_store: Arc::new(()),
            bypass_paths: vec!["/v1/health".to_string(), "/swagger-ui".to_string()],
        };

        assert!(service.should_bypass("/v1/health"));
        assert!(service.should_bypass("/swagger-ui/index.html"));
        assert!(!service.should_bypass("/v1/assert"));
    }
}
