use axum::extract::{FromRequestParts, Request, State};
use axum::http::StatusCode;
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use contracts::auth_middleware::AuthUser;
use std::future::Future;
use uuid::Uuid;

use crate::ai::{credits, plans};
use crate::AppState;

/// Extractor that ensures the user has an active AI subscription and is not a
/// customer. It can be combined with `AuthUser` in handlers that need it.
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct AiAccess {
    pub user_id: Uuid,
    pub role_code: Option<String>,
    pub subscription: db::models::ai::UserAiSubscription,
    pub plan: db::models::ai::AiPlan,
}

#[allow(dead_code)]
impl AiAccess {
    /// Remaining credits from the user's subscription.
    pub fn remaining_credits(&self) -> i32 {
        credits::remaining_credits(&self.subscription)
    }

    /// Remaining daily AI actions from the user's subscription.
    pub fn remaining_daily_actions(&self) -> i32 {
        credits::remaining_daily_actions(&self.subscription, &self.plan)
    }

    /// Whether a feature is allowed by the current plan.
    pub fn feature_allowed(&self, feature_code: &str) -> bool {
        plans::is_feature_allowed(&self.plan, feature_code)
    }
}

impl<S> FromRequestParts<S> for AiAccess
where
    S: Send + Sync,
{
    type Rejection = AiAccessError;

    fn from_request_parts(
        parts: &mut axum::http::request::Parts,
        _state: &S,
    ) -> impl Future<Output = Result<Self, Self::Rejection>> + Send {
        let auth_header = parts
            .headers
            .get("Authorization")
            .and_then(|v| v.to_str().ok())
            .map(str::to_string);

        async move {
            let auth_header = auth_header.ok_or(AiAccessError::MissingToken)?;
            let token = auth_header
                .strip_prefix("Bearer ")
                .ok_or(AiAccessError::InvalidToken)?;

            let jwt_secret = std::env::var("JWT_SECRET")
                .expect("JWT_SECRET must be set — refusing to start with insecure default");

            let token_data = jsonwebtoken::decode::<contracts::auth_middleware::Claims>(
                token,
                &jsonwebtoken::DecodingKey::from_secret(jwt_secret.as_bytes()),
                &jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::HS256),
            )
            .map_err(|_| AiAccessError::InvalidToken)?;

            let _user_id = Uuid::parse_str(&token_data.claims.sub)
                .map_err(|_| AiAccessError::InvalidToken)?;

            let _role_code = token_data.claims.active_role.clone();

            // State is not available via FromRequestParts, so we cannot load
            // the subscription here. Use the Layer middleware below for full checks.
            Err(AiAccessError::IncompleteContext)
        }
    }
}

#[derive(Debug)]
#[allow(dead_code)]
pub enum AiAccessError {
    MissingToken,
    InvalidToken,
    IncompleteContext,
    Plan(plans::PlanError),
}

impl From<plans::PlanError> for AiAccessError {
    fn from(e: plans::PlanError) -> Self {
        AiAccessError::Plan(e)
    }
}

impl IntoResponse for AiAccessError {
    fn into_response(self) -> Response {
        match self {
            AiAccessError::Plan(p) => (p.status_code(), axum::Json(p.error_body())).into_response(),
            AiAccessError::MissingToken => (
                StatusCode::UNAUTHORIZED,
                axum::Json(serde_json::json!({
                    "error": "Authorization header required",
                    "code": "MISSING_TOKEN"
                })),
            )
                .into_response(),
            AiAccessError::InvalidToken => (
                StatusCode::UNAUTHORIZED,
                axum::Json(serde_json::json!({
                    "error": "Token is invalid or expired",
                    "code": "INVALID_TOKEN"
                })),
            )
                .into_response(),
            AiAccessError::IncompleteContext => (
                StatusCode::INTERNAL_SERVER_ERROR,
                axum::Json(serde_json::json!({
                    "error": "AI access extractor requires middleware context",
                    "code": "INCOMPLETE_CONTEXT"
                })),
            )
                .into_response(),
        }
    }
}

/// Middleware that ensures a free subscription exists and validates the user's
/// AI access. Applies to all routes under `/api/ai/*`. Non-customer roles that
/// have no explicit subscription will automatically get a Free plan.
pub async fn ai_access_middleware(
    State(state): State<AppState>,
    auth: AuthUser,
    request: Request,
    next: Next,
) -> Result<Response, AiAccessError> {
    let role_code = auth.claims.active_role.clone();

    let (_sub, _plan) = plans::ensure_free_subscription(
        &state.pool,
        auth.user_id,
        Some(&role_code),
    )
    .await?;

    // TODO: if stricter enforcement is required, reject here instead of
    // only inside individual handlers. For now we allow the request through and
    // let handlers charge per-feature.

    Ok(next.run(request).await)
}
