//! Admission control middleware (The Shield).
//!
//! This middleware enforces proof-of-work requirements for new/untrusted agents.
//! It extracts the agent ID from the `X-Agent-Id` header, checks admission status,
//! and verifies PoW proofs when required.
//!
//! # Request Flow
//!
//! 1. Extract `X-Agent-Id` header (hex-encoded 32-byte public key)
//! 2. Get admission status (tier, PoW requirement)
//! 3. If PoW required:
//!    - Extract `X-PoW-Nonce` and `X-PoW-Timestamp` headers
//!    - Verify proof meets difficulty requirement
//!    - Return 428 if invalid/missing
//! 4. Store admission status in request extensions (for MeterLayer)
//! 5. Add response headers (`X-Trust-Tier`, `X-PoW-Required`, etc.)
//!
//! # Headers
//!
//! | Header | Direction | Description |
//! |--------|-----------|-------------|
//! | `X-Agent-Id` | Request | Agent's Ed25519 public key (hex, 64 chars) |
//! | `X-PoW-Nonce` | Request | PoW solution nonce (decimal) |
//! | `X-PoW-Timestamp` | Request | PoW solution timestamp (Unix seconds) |
//! | `X-Trust-Tier` | Response | Agent's trust tier name |
//! | `X-PoW-Required` | Response | "true" or "false" |
//! | `X-PoW-Difficulty` | Response | Required difficulty in bits |
//! | `X-Quota-Multiplier` | Response | Tier quota multiplier |

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_core::types::PowProof;
use stemedb_storage::{AdmissionCheck, AdmissionStatus, AdmissionStatusResult};
use tower::{Layer, Service};
use tracing::{debug, warn};

/// Header name for agent identification (shared with MeterLayer).
pub const AGENT_ID_HEADER: &str = "x-agent-id";

/// Header name for PoW nonce.
pub const POW_NONCE_HEADER: &str = "x-pow-nonce";

/// Header name for PoW timestamp.
pub const POW_TIMESTAMP_HEADER: &str = "x-pow-timestamp";

/// Response header for trust tier.
pub const TRUST_TIER_HEADER: &str = "x-trust-tier";

/// Response header indicating whether PoW is required.
pub const POW_REQUIRED_HEADER: &str = "x-pow-required";

/// Response header for PoW difficulty.
pub const POW_DIFFICULTY_HEADER: &str = "x-pow-difficulty";

/// Response header for quota multiplier.
pub const QUOTA_MULTIPLIER_HEADER: &str = "x-quota-multiplier";

/// HTTP 428 Precondition Required - PoW needed.
const HTTP_PRECONDITION_REQUIRED: u16 = 428;

/// Error response for PoW required.
#[derive(Debug, Serialize)]
struct PowRequiredError {
    /// Human-readable error message.
    error: String,
    /// Error code for programmatic handling.
    code: String,
    /// Required PoW difficulty in bits.
    required_difficulty: u8,
    /// Whether PoW is required.
    pow_required: bool,
    /// Number of assertions the agent has made.
    agent_assertions: u64,
    /// Agent's current trust score.
    agent_trust_score: f32,
    /// Optional detailed error message (for failed verification).
    #[serde(skip_serializing_if = "Option::is_none")]
    verification_error: Option<String>,
}

/// Tower Layer for admission control.
///
/// Wrap your router with this layer to enable PoW-based admission control.
/// This layer should be applied BEFORE the MeterLayer so that PoW is checked
/// before quota is consumed.
///
/// # Example
///
/// ```ignore
/// let admission_layer = AdmissionLayer::new(admission_store);
/// let meter_layer = MeterLayer::new(quota_store);
///
/// let app = Router::new()
///     .route("/v1/assert", post(create_assertion))
///     .layer(meter_layer)      // Inner: runs second
///     .layer(admission_layer)  // Outer: runs first
/// ```
#[derive(Clone)]
pub struct AdmissionLayer<A> {
    admission_store: Arc<A>,
    /// Paths that bypass admission control (e.g., health checks, status endpoint).
    bypass_paths: Vec<String>,
}

impl<A> AdmissionLayer<A> {
    /// Create a new AdmissionLayer.
    pub fn new(admission_store: Arc<A>) -> Self {
        Self {
            admission_store,
            bypass_paths: vec![
                "/v1/health".to_string(),
                "/v1/admission/status".to_string(),
                "/swagger-ui".to_string(),
                "/api-docs".to_string(),
            ],
        }
    }

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

impl<S, A> Layer<S> for AdmissionLayer<A>
where
    A: Clone,
{
    type Service = AdmissionService<S, A>;

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

/// Tower Service for admission control.
#[derive(Clone)]
pub struct AdmissionService<S, A> {
    inner: S,
    admission_store: Arc<A>,
    bypass_paths: Vec<String>,
}

impl<S, A> AdmissionService<S, A> {
    /// Check if path should bypass admission control.
    #[allow(dead_code)] // Used in tests
    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
            }
        })
    }

    /// Extract PoW proof from request headers.
    fn extract_pow_proof(req: &Request<Body>, agent_id: [u8; 32]) -> Option<PowProof> {
        let nonce_str = req.headers().get(POW_NONCE_HEADER)?.to_str().ok()?;
        let timestamp_str = req.headers().get(POW_TIMESTAMP_HEADER)?.to_str().ok()?;

        let nonce: u64 = nonce_str.parse().ok()?;
        let timestamp: u64 = timestamp_str.parse().ok()?;

        Some(PowProof::new(nonce, agent_id, timestamp))
    }

    /// Add admission headers to response.
    fn add_response_headers(response: &mut Response<Body>, status: &AdmissionStatus) {
        let headers = response.headers_mut();

        if let Ok(v) = status.tier.name().parse() {
            headers.insert(TRUST_TIER_HEADER, v);
        }
        if let Ok(v) = status.pow_required.to_string().parse() {
            headers.insert(POW_REQUIRED_HEADER, v);
        }
        if let Ok(v) = status.pow_difficulty.to_string().parse() {
            headers.insert(POW_DIFFICULTY_HEADER, v);
        }
        if let Ok(v) = format!("{:.1}", status.quota_multiplier).parse() {
            headers.insert(QUOTA_MULTIPLIER_HEADER, v);
        }
    }

    /// Build a 428 response for PoW required.
    fn pow_required_response(
        status: &AdmissionStatus,
        verification_error: Option<String>,
    ) -> Response<Body> {
        let error_message = if verification_error.is_some() {
            "Proof-of-Work verification failed"
        } else {
            "Proof-of-Work required"
        };

        let error = PowRequiredError {
            error: error_message.to_string(),
            code: "POW_REQUIRED".to_string(),
            required_difficulty: status.pow_difficulty,
            pow_required: true,
            agent_assertions: status.assertions_count,
            agent_trust_score: status.trust_score,
            verification_error,
        };

        let mut response = (
            StatusCode::from_u16(HTTP_PRECONDITION_REQUIRED)
                .unwrap_or(StatusCode::PRECONDITION_FAILED),
            Json(error),
        )
            .into_response();

        Self::add_response_headers(&mut response, status);
        response
    }
}

impl<S, A> Service<Request<Body>> for AdmissionService<S, A>
where
    S: Service<Request<Body>, Response = Response<Body>> + Clone + Send + 'static,
    S::Future: Send,
    A: AdmissionCheck + '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 admission_store = Arc::clone(&self.admission_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 admission control
            if bypass_paths.iter().any(|p| path.starts_with(p)) {
                debug!(path = %path, "Bypassing admission control for path");
                return inner.call(req).await;
            }

            // Only check admission for write paths
            let is_write_path = path.starts_with("/v1/assert")
                || path.starts_with("/v1/vote")
                || path.starts_with("/v1/supersede");

            if !is_write_path {
                // Read-only paths don't need admission control
                debug!(path = %path, "Skipping admission for read-only 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 (will fail signature verification)
                    debug!(path = %path, "No agent ID, skipping admission");
                    return inner.call(req).await;
                }
            };

            // Extract PoW proof (if provided)
            let proof = Self::extract_pow_proof(&req, agent_id);

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

            // Check admission
            let admission_result =
                match admission_store.check_admission(&agent_id, proof.as_ref(), server_time).await
                {
                    Ok(result) => result,
                    Err(e) => {
                        warn!(error = %e, "Admission check failed, allowing request");
                        // On error, allow the request (fail open for availability)
                        return inner.call(req).await;
                    }
                };

            match admission_result {
                AdmissionStatusResult::Admitted(status) => {
                    debug!(
                        agent = %hex::encode(agent_id),
                        tier = %status.tier,
                        "Agent admitted"
                    );

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

                    // Add admission headers to response
                    Self::add_response_headers(&mut response, &status);

                    Ok(response)
                }
                AdmissionStatusResult::PowRequired(status) => {
                    debug!(
                        agent = %hex::encode(agent_id),
                        difficulty = status.pow_difficulty,
                        "PoW required"
                    );

                    Ok(Self::pow_required_response(&status, None))
                }
                AdmissionStatusResult::PowFailed { status, error } => {
                    debug!(
                        agent = %hex::encode(agent_id),
                        error = %error,
                        "PoW verification failed"
                    );

                    Ok(Self::pow_required_response(&status, Some(error.to_string())))
                }
            }
        })
    }
}

/// Request extension to share admission status with other middleware.
///
/// The MeterLayer can read this to apply tier-based quota multipliers.
#[derive(Debug, Clone)]
pub struct AdmissionExtension {
    /// The agent's admission status.
    pub status: AdmissionStatus,
}

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

    #[test]
    fn test_bypass_paths() {
        let service = AdmissionService::<(), ()> {
            inner: (),
            admission_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"));
    }

    #[test]
    fn test_extract_agent_id() {
        let req = Request::builder()
            .header(
                AGENT_ID_HEADER,
                "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef",
            )
            .body(Body::empty())
            .expect("build request");

        let agent_id = AdmissionService::<(), ()>::extract_agent_id(&req);
        assert!(agent_id.is_some());
        let id = agent_id.expect("id");
        assert_eq!(id[0], 0x01);
        assert_eq!(id[1], 0x23);
    }

    #[test]
    fn test_extract_agent_id_invalid_length() {
        let req = Request::builder()
            .header(AGENT_ID_HEADER, "0123456789abcdef") // Too short
            .body(Body::empty())
            .expect("build request");

        let agent_id = AdmissionService::<(), ()>::extract_agent_id(&req);
        assert!(agent_id.is_none());
    }

    #[test]
    fn test_extract_pow_proof() {
        let agent_id = [0xABu8; 32];
        let req = Request::builder()
            .header(POW_NONCE_HEADER, "12345")
            .header(POW_TIMESTAMP_HEADER, "1700000000")
            .body(Body::empty())
            .expect("build request");

        let proof = AdmissionService::<(), ()>::extract_pow_proof(&req, agent_id);
        assert!(proof.is_some());
        let p = proof.expect("proof");
        assert_eq!(p.nonce, 12345);
        assert_eq!(p.timestamp, 1700000000);
        assert_eq!(p.agent_id, agent_id);
    }

    #[test]
    fn test_extract_pow_proof_missing_headers() {
        let agent_id = [0xABu8; 32];

        // Missing nonce
        let req = Request::builder()
            .header(POW_TIMESTAMP_HEADER, "1700000000")
            .body(Body::empty())
            .expect("build request");

        let proof = AdmissionService::<(), ()>::extract_pow_proof(&req, agent_id);
        assert!(proof.is_none());

        // Missing timestamp
        let req = Request::builder()
            .header(POW_NONCE_HEADER, "12345")
            .body(Body::empty())
            .expect("build request");

        let proof = AdmissionService::<(), ()>::extract_pow_proof(&req, agent_id);
        assert!(proof.is_none());
    }
}
