use std::net::SocketAddr;

use axum::Json;
use axum::body::Body;
use axum::extract::{ConnectInfo, State};
use axum::http::{HeaderMap, HeaderValue, Request, StatusCode, header};
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use base64::Engine as _;
use models::users::user_access_keys::AccessKeyScope;
use rmcp::ErrorData;
use rmcp::RoleServer;
use rmcp::service::RequestContext;
use serde_json::json;
use uuid::Uuid;

use crate::{McpRuntime, net};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PrincipalSource {
    UserKey,
    InternalKey,
}

#[derive(Debug, Clone)]
pub struct Principal {
    pub uid: Uuid,
    pub scope: AccessKeyScope,
    pub full_access: bool,
    pub source: PrincipalSource,
}

impl Principal {
    pub fn can_write(&self) -> bool {
        self.full_access || self.scope.allows(true)
    }

    pub(crate) fn from_ctx(ctx: &RequestContext<RoleServer>) -> Result<Self, ErrorData> {
        let parts = ctx
            .extensions
            .get::<axum::http::request::Parts>()
            .ok_or_else(|| ErrorData::internal_error("missing http request parts", None))?;
        parts
            .extensions
            .get::<Principal>()
            .cloned()
            .ok_or_else(|| ErrorData::invalid_request("unauthenticated", None))
    }
}

pub(crate) async fn auth_middleware(
    State(runtime): State<McpRuntime>,
    ConnectInfo(peer): ConnectInfo<SocketAddr>,
    mut request: Request<Body>,
    next: Next,
) -> Response {
    let Some(token) = extract_token(request.headers()) else {
        return unauthorized("缺少访问密钥");
    };
    let peer = Some(peer);
    let client_ip = net::resolve_client_ip(peer, request.headers(), runtime.trusted_proxies());

    let principal = if let Some(name) = runtime.match_internal_key(&token) {
        if !runtime.is_internal_ip(client_ip) {
            tracing::warn!(key = %name, "internal mcp key rejected from a non-internal source");
            return forbidden("内部访问密钥仅允许内网来源使用");
        }
        let Some(uid) = runtime.system_uid() else {
            return unavailable("内部系统账号尚未就绪");
        };
        Principal {
            uid,
            scope: AccessKeyScope::Write,
            full_access: true,
            source: PrincipalSource::InternalKey,
        }
    } else {
        match runtime.services().auth.verify_access_token(&token).await {
            Ok(Some(verified)) => Principal {
                uid: verified.user_id,
                scope: verified.scope,
                full_access: false,
                source: PrincipalSource::UserKey,
            },
            Ok(None) => return unauthorized("访问密钥无效或已过期"),
            Err(error) => {
                tracing::warn!(%error, "mcp access key verification failed");
                return unavailable("鉴权服务暂不可用");
            }
        }
    };

    request.extensions_mut().insert(principal);
    next.run(request).await
}

fn extract_token(headers: &HeaderMap) -> Option<String> {
    let value = headers.get(header::AUTHORIZATION)?.to_str().ok()?;
    if let Some(rest) = value
        .strip_prefix("Bearer ")
        .or_else(|| value.strip_prefix("bearer "))
    {
        let token = rest.trim();
        return (!token.is_empty()).then(|| token.to_string());
    }
    if let Some(rest) = value.strip_prefix("Basic ") {
        let decoded = base64::engine::general_purpose::STANDARD
            .decode(rest.trim())
            .ok()?;
        let credentials = String::from_utf8(decoded).ok()?;
        return credentials
            .split_once(':')
            .map(|(_, password)| password.to_string());
    }
    None
}

fn error_response(status: StatusCode, message: &str) -> Response {
    let body = Json(json!({
        "error": status.canonical_reason().unwrap_or("error").to_lowercase(),
        "message": message,
    }));
    (status, body).into_response()
}

fn unauthorized(message: &str) -> Response {
    let mut response = error_response(StatusCode::UNAUTHORIZED, message);
    response
        .headers_mut()
        .insert(header::WWW_AUTHENTICATE, HeaderValue::from_static("Bearer"));
    response
}

fn forbidden(message: &str) -> Response {
    error_response(StatusCode::FORBIDDEN, message)
}

fn unavailable(message: &str) -> Response {
    error_response(StatusCode::SERVICE_UNAVAILABLE, message)
}