#![allow(dead_code)]

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sha2::{Digest, Sha256};
use tracing;
use validator::Validate;

use athena_auth_core::adapters::DatabaseAdapter;
use athena_auth_core::entity::{AuthSession, AuthUser};
use athena_auth_core::utils::cookie_utils::create_clear_session_cookie;
use athena_auth_core::{
    AuthContext, AuthError, AuthPlugin, AuthRequest, AuthResponse, AuthResult, AuthRoute,
    ExtractedSessionToken, GrantDescriptor, HttpMethod, SessionManager, SessionOps,
    SuccessResponse, resolve_avatar_response_urls_in_json,
};

use super::StatusResponse;

use crate::plugins::grants::{merged_grant_catalog, merged_native_right_catalog};
#[cfg(feature = "axum")]
use crate::plugins::helpers::SessionWithServerVersion;
use crate::plugins::helpers::{
    load_effective_grant_state, session_with_server_version, sessions_with_server_version,
};

/// Session management plugin for handling session operations
pub struct SessionManagementPlugin {
    config: SessionManagementConfig,
}

#[derive(Debug, Clone, athena_auth_core::PluginConfig)]
#[plugin(name = "SessionManagementPlugin")]
pub struct SessionManagementConfig {
    #[config(default = true)]
    pub enable_session_listing: bool,
    #[config(default = true)]
    pub enable_session_revocation: bool,
    #[config(default = true)]
    pub require_authentication: bool,
    #[config(default = Vec::new())]
    pub grant_catalog: Vec<GrantDescriptor>,
    #[config(default = Vec::new(), skip)]
    pub native_right_catalog: Vec<athena_rights::AthenaRightDescriptor>,
}

impl SessionManagementConfig {
    pub fn effective_grant_catalog(&self) -> Vec<GrantDescriptor> {
        merged_grant_catalog(&self.grant_catalog, &self.native_right_catalog)
    }

    pub fn merged_native_right_catalog(&self) -> Vec<athena_rights::AthenaRightDescriptor> {
        merged_native_right_catalog(&self.native_right_catalog)
    }
}

// Request structures for session endpoints
#[derive(Debug, Deserialize, Validate)]
struct RevokeSessionRequest {
    #[validate(length(min = 1, message = "Token is required"))]
    token: String,
}

#[cfg(feature = "axum")]
#[derive(Debug, Serialize)]
struct GetSessionResponse<S: Serialize, U: Serialize> {
    session: SessionWithServerVersion<S>,
    user: U,
}

#[derive(Debug, Serialize)]
struct GetSessionUnauthorizedResponse {
    message: String,
    code: &'static str,
    reason: &'static str,
    session: Option<Value>,
    user: Option<Value>,
    diagnostics: GetSessionDiagnostics,
}

#[derive(Debug, Serialize)]
struct GetSessionDiagnostics {
    method: &'static str,
    path: String,
    #[serde(rename = "tokenSource")]
    token_source: Option<&'static str>,
    #[serde(rename = "authorizationHeaderPresent")]
    authorization_header_present: bool,
    #[serde(rename = "cookieHeaderPresent")]
    cookie_header_present: bool,
    #[serde(rename = "sessionCookieName")]
    session_cookie_name: String,
    #[serde(rename = "sessionCookiePresent")]
    session_cookie_present: bool,
    #[serde(rename = "disableCookieCache")]
    disable_cookie_cache: Option<String>,
}

#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
impl<DB: DatabaseAdapter> AuthPlugin<DB> for SessionManagementPlugin {
    fn name(&self) -> &'static str {
        "session-management"
    }

    fn routes(&self) -> Vec<AuthRoute> {
        vec![
            AuthRoute::get("/get-session", "get_session"),
            AuthRoute::post("/get-session", "get_session_post"),
            AuthRoute::post("/sign-out", "sign_out"),
            AuthRoute::get("/list-sessions", "list_sessions"),
            AuthRoute::post("/revoke-session", "revoke_session"),
            AuthRoute::post("/revoke-sessions", "revoke_sessions"),
            AuthRoute::post("/revoke-other-sessions", "revoke_other_sessions"),
        ]
    }

    async fn on_request(
        &self,
        req: &AuthRequest,
        ctx: &AuthContext<DB>,
    ) -> AuthResult<Option<AuthResponse>> {
        match (req.method(), req.path()) {
            (HttpMethod::Get | HttpMethod::Post, "/get-session") => {
                Ok(Some(self.handle_get_session(req, ctx).await?))
            }
            (HttpMethod::Post, "/sign-out") => Ok(Some(self.handle_sign_out(req, ctx).await?)),
            (HttpMethod::Get, "/list-sessions") if self.config.enable_session_listing => {
                Ok(Some(self.handle_list_sessions(req, ctx).await?))
            }
            (HttpMethod::Post, "/revoke-session") if self.config.enable_session_revocation => {
                Ok(Some(self.handle_revoke_session(req, ctx).await?))
            }
            (HttpMethod::Post, "/revoke-sessions") if self.config.enable_session_revocation => {
                Ok(Some(self.handle_revoke_sessions(req, ctx).await?))
            }
            (HttpMethod::Post, "/revoke-other-sessions")
                if self.config.enable_session_revocation =>
            {
                Ok(Some(self.handle_revoke_other_sessions(req, ctx).await?))
            }
            _ => Ok(None),
        }
    }
}

// ---------------------------------------------------------------------------
// Core functions — framework-agnostic business logic
// ---------------------------------------------------------------------------

pub(crate) async fn sign_out_core<DB: DatabaseAdapter>(
    session: &DB::Session,
    ctx: &AuthContext<DB>,
) -> AuthResult<SuccessResponse> {
    ctx.database.delete_session(session.token()).await?;
    Ok(SuccessResponse { success: true })
}

pub(crate) async fn list_sessions_core<DB: DatabaseAdapter>(
    user_id: &str,
    ctx: &AuthContext<DB>,
) -> AuthResult<Vec<DB::Session>> {
    ctx.session_manager().list_user_sessions(user_id).await
}

pub(crate) async fn revoke_session_core<DB: DatabaseAdapter>(
    user: &DB::User,
    token: &str,
    ctx: &AuthContext<DB>,
) -> AuthResult<StatusResponse> {
    // Verify the session belongs to the current user before revoking
    let session_manager: SessionManager<DB> = ctx.session_manager();
    if let Some(session_to_revoke) = session_manager.get_session(token).await?
        && session_to_revoke.user_id() != user.id()
    {
        return Err(AuthError::forbidden(
            "Cannot revoke session that belongs to another user",
        ));
    }

    ctx.database.delete_session(token).await?;
    Ok(StatusResponse { status: true })
}

pub(crate) async fn revoke_sessions_core<DB: DatabaseAdapter>(
    user_id: &str,
    ctx: &AuthContext<DB>,
) -> AuthResult<StatusResponse> {
    ctx.database.delete_user_sessions(user_id).await?;
    Ok(StatusResponse { status: true })
}

pub(crate) async fn revoke_other_sessions_core<DB: DatabaseAdapter>(
    user_id: &str,
    current_session: &DB::Session,
    ctx: &AuthContext<DB>,
) -> AuthResult<StatusResponse> {
    let all_sessions: Vec<DB::Session> = ctx.database.get_user_sessions(user_id).await?;
    for session in all_sessions {
        if session.token() != current_session.token() {
            ctx.database.delete_session(session.token()).await?;
        }
    }
    Ok(StatusResponse { status: true })
}

async fn get_session_core_response<DB: DatabaseAdapter>(
    req: &AuthRequest,
    ctx: &AuthContext<DB>,
    config: &SessionManagementConfig,
) -> AuthResult<AuthResponse> {
    let session_manager: SessionManager<DB> = ctx.session_manager();
    let Some(extracted_token) = session_manager.extract_session_token_with_source(req) else {
        tracing::warn!(
            path = %req.path(),
            method = ?req.method(),
            authorization_header_present = request_header_present(req, "authorization"),
            cookie_header_present = request_header_present(req, "cookie"),
            session_cookie_name = %ctx.config.session.cookie_name,
            disable_cookie_cache = ?req.query.get("disableCookieCache"),
            "get-session rejected: no session credential found"
        );
        return get_session_unauthorized_response(
            req,
            ctx,
            None,
            "SESSION_TOKEN_MISSING",
            "missing_session_token",
            "Authentication required",
        );
    };

    let token_hash: String = token_hash_prefix(&extracted_token.token);
    let Some(session) = session_manager.get_session(&extracted_token.token).await? else {
        tracing::warn!(
            path = %req.path(),
            method = ?req.method(),
            token_source = extracted_token.source.as_str(),
            session_token_hash = %token_hash,
            token_len = extracted_token.token.len(),
            session_cookie_present = request_has_configured_session_cookie(req, &ctx.config.session.cookie_name),
            disable_cookie_cache = ?req.query.get("disableCookieCache"),
            "get-session rejected: session token was not found or is expired"
        );
        return get_session_unauthorized_response(
            req,
            ctx,
            Some(&extracted_token),
            "SESSION_NOT_FOUND",
            "session_not_found_or_expired",
            "Session not found or expired",
        );
    };

    tracing::debug!(
        path = %req.path(),
        method = ?req.method(),
        token_source = extracted_token.source.as_str(),
        session_token_hash = %token_hash,
        session_id = %session.id(),
        user_id = %session.user_id(),
        "get-session resolved active session"
    );

    let Some(user) = ctx.database.get_user_by_id(session.user_id()).await? else {
        tracing::warn!(
            path = %req.path(),
            method = ?req.method(),
            token_source = extracted_token.source.as_str(),
            session_token_hash = %token_hash,
            session_id = %session.id(),
            user_id = %session.user_id(),
            "get-session rejected: session user was not found"
        );
        return get_session_unauthorized_response(
            req,
            ctx,
            Some(&extracted_token),
            "SESSION_USER_NOT_FOUND",
            "session_user_not_found",
            "Session user not found",
        );
    };

    let mut user_json: Value = serde_json::to_value(&user).map_err(|error| {
        AuthError::internal(format!(
            "Failed to serialize get-session user response: {error}"
        ))
    })?;
    ctx.session_user_cache
        .insert(user.id().to_string(), user_json.clone());
    resolve_avatar_response_urls_in_json(&ctx.config, &mut user_json).await?;

    tracing::trace!(user_json = ?user_json, "Serialized user JSON for session response");
    let effective_grant_catalog = config.effective_grant_catalog();
    let native_right_catalog = config.merged_native_right_catalog();
    let grants = load_effective_grant_state(
        ctx,
        &effective_grant_catalog,
        Some(&native_right_catalog),
        "user",
        user.id(),
        None,
    )
    .await?;

    let response: Value = serde_json::json!({
        "session": session_with_server_version(session),
        "user": user_json,
        "grants": grants,
        "rights": grants,
    });
    tracing::trace!(response = ?response, "Serialized response of the session: {response}");

    Ok(AuthResponse::json(200, &response)?)
}

// ---------------------------------------------------------------------------
// Old handler methods — delegate to core functions
// ---------------------------------------------------------------------------

impl SessionManagementPlugin {
    pub fn with_grant_catalog(mut self, grant_catalog: Vec<GrantDescriptor>) -> Self {
        self.config.grant_catalog = grant_catalog;
        self
    }

    pub fn with_native_right_catalog(
        mut self,
        native_right_catalog: Vec<athena_rights::AthenaRightDescriptor>,
    ) -> Self {
        self.config.native_right_catalog = native_right_catalog;
        self
    }

    async fn handle_get_session<DB: DatabaseAdapter>(
        &self,
        req: &AuthRequest,
        ctx: &AuthContext<DB>,
    ) -> AuthResult<AuthResponse> {
        get_session_core_response(req, ctx, &self.config).await
    }

    async fn handle_sign_out<DB: DatabaseAdapter>(
        &self,
        req: &AuthRequest,
        ctx: &AuthContext<DB>,
    ) -> AuthResult<AuthResponse> {
        let (user, session) = ctx.require_session(req).await?;
        let response: SuccessResponse = sign_out_core(&session, ctx).await?;
        tracing::info!(
            user_id = %user.id(),
            email = %user.email().unwrap_or(""),
            "user signed out"
        );
        let clear_cookie: String = create_clear_session_cookie(&ctx.config);
        Ok(AuthResponse::json(200, &response)?.with_header("Set-Cookie", clear_cookie))
    }

    async fn handle_list_sessions<DB: DatabaseAdapter>(
        &self,
        req: &AuthRequest,
        ctx: &AuthContext<DB>,
    ) -> AuthResult<AuthResponse> {
        let (user, _) = ctx.require_session(req).await?;
        let sessions: Vec<super::helpers::SessionWithServerVersion<<DB as SessionOps>::Session>> =
            sessions_with_server_version(list_sessions_core(user.id(), ctx).await?);
        Ok(AuthResponse::json(200, &sessions)?)
    }

    async fn handle_revoke_session<DB: DatabaseAdapter>(
        &self,
        req: &AuthRequest,
        ctx: &AuthContext<DB>,
    ) -> AuthResult<AuthResponse> {
        let (user, current_session) = ctx.require_session(req).await?;

        let revoke_req: RevokeSessionRequest = match athena_auth_core::validate_request_body(req) {
            Ok(v) => v,
            Err(resp) => return Ok(resp),
        };

        let response: StatusResponse = revoke_session_core(&user, &revoke_req.token, ctx).await?;
        let auth_response: AuthResponse = AuthResponse::json(200, &response)?;
        if revoke_req.token == current_session.token() {
            let clear_cookie: String = create_clear_session_cookie(&ctx.config);
            return Ok(auth_response.with_header("Set-Cookie", clear_cookie));
        }
        Ok(auth_response)
    }

    async fn handle_revoke_sessions<DB: DatabaseAdapter>(
        &self,
        req: &AuthRequest,
        ctx: &AuthContext<DB>,
    ) -> AuthResult<AuthResponse> {
        let (user, _) = ctx.require_session(req).await?;
        let response: StatusResponse = revoke_sessions_core(user.id(), ctx).await?;
        let clear_cookie: String = create_clear_session_cookie(&ctx.config);
        Ok(AuthResponse::json(200, &response)?.with_header("Set-Cookie", clear_cookie))
    }

    async fn handle_revoke_other_sessions<DB: DatabaseAdapter>(
        &self,
        req: &AuthRequest,
        ctx: &AuthContext<DB>,
    ) -> AuthResult<AuthResponse> {
        let (user, current_session) = ctx.require_session(req).await?;
        let response: StatusResponse =
            revoke_other_sessions_core(user.id(), &current_session, ctx).await?;
        Ok(AuthResponse::json(200, &response)?)
    }
}

fn get_session_unauthorized_response<DB: DatabaseAdapter>(
    req: &AuthRequest,
    ctx: &AuthContext<DB>,
    extracted_token: Option<&ExtractedSessionToken>,
    code: &'static str,
    reason: &'static str,
    message: impl Into<String>,
) -> AuthResult<AuthResponse> {
    let diagnostics: GetSessionDiagnostics = GetSessionDiagnostics {
        method: method_label(req.method()),
        path: req.path().to_string(),
        token_source: extracted_token.map(|token| token.source.as_str()),
        authorization_header_present: request_header_present(req, "authorization"),
        cookie_header_present: request_header_present(req, "cookie"),
        session_cookie_name: ctx.config.session.cookie_name.clone(),
        session_cookie_present: request_has_configured_session_cookie(
            req,
            &ctx.config.session.cookie_name,
        ),
        disable_cookie_cache: req.query.get("disableCookieCache").cloned(),
    };
    let should_clear_cookie = diagnostics.session_cookie_present;
    let payload: GetSessionUnauthorizedResponse = GetSessionUnauthorizedResponse {
        message: message.into(),
        code,
        reason,
        session: None,
        user: None,
        diagnostics,
    };
    let mut response: AuthResponse = AuthResponse::json(401, &payload)?;
    if should_clear_cookie {
        response = response.with_header("Set-Cookie", create_clear_session_cookie(&ctx.config));
    }
    Ok(response)
}

fn request_header_present(req: &AuthRequest, name: &str) -> bool {
    req.headers
        .keys()
        .any(|header_name| header_name.eq_ignore_ascii_case(name))
}

fn request_header<'a>(req: &'a AuthRequest, name: &str) -> Option<&'a str> {
    req.headers
        .iter()
        .find(|(header_name, _)| header_name.eq_ignore_ascii_case(name))
        .map(|(_, header_value)| header_value.as_str())
}

fn request_has_configured_session_cookie(req: &AuthRequest, cookie_name: &str) -> bool {
    let Some(cookie_header) = request_header(req, "cookie") else {
        return false;
    };

    cookie::Cookie::split_parse(cookie_header)
        .flatten()
        .any(|cookie| cookie.name() == cookie_name && !cookie.value().is_empty())
}

fn method_label(method: &HttpMethod) -> &'static str {
    match method {
        HttpMethod::Get => "GET",
        HttpMethod::Post => "POST",
        HttpMethod::Put => "PUT",
        HttpMethod::Delete => "DELETE",
        HttpMethod::Patch => "PATCH",
        HttpMethod::Options => "OPTIONS",
        HttpMethod::Head => "HEAD",
    }
}

fn token_hash_prefix(token: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(token.as_bytes());
    let hash = format!("{:x}", hasher.finalize());
    hash.chars().take(16).collect()
}

#[cfg(feature = "axum")]
mod axum_impl {
    use super::*;
    use std::sync::Arc;

    use athena_auth_core::{
        AuthRequestExt, AuthState, AxumAuthResponse, CurrentSession, ValidatedJson,
    };
    use axum::Json;
    use axum::extract::{Extension, State};
    use axum::http::header;
    use axum::response::IntoResponse;

    #[derive(Clone)]
    struct PluginState {
        config: SessionManagementConfig,
    }

    // get_session is trivially simple: just construct the response directly.
    async fn handle_get_session<DB: DatabaseAdapter>(
        State(state): State<AuthState<DB>>,
        Extension(ps): Extension<Arc<PluginState>>,
        AuthRequestExt(req): AuthRequestExt,
    ) -> Result<AxumAuthResponse, AuthError> {
        let ctx = state.to_context();
        Ok(AxumAuthResponse(
            get_session_core_response(&req, &ctx, &ps.config).await?,
        ))
    }

    async fn handle_sign_out<DB: DatabaseAdapter>(
        State(state): State<AuthState<DB>>,
        CurrentSession { session, .. }: CurrentSession<DB>,
    ) -> Result<([(header::HeaderName, String); 1], Json<SuccessResponse>), AuthError> {
        let ctx = state.to_context();
        let response = sign_out_core(&session, &ctx).await?;
        let cookie = state.clear_session_cookie();
        Ok(([(header::SET_COOKIE, cookie)], Json(response)))
    }

    async fn handle_list_sessions<DB: DatabaseAdapter>(
        State(state): State<AuthState<DB>>,
        Extension(ps): Extension<Arc<PluginState>>,
        CurrentSession { user, .. }: CurrentSession<DB>,
    ) -> Result<Json<Vec<SessionWithServerVersion<DB::Session>>>, AuthError> {
        if !ps.config.enable_session_listing {
            return Err(AuthError::not_found("Not found"));
        }
        let ctx = state.to_context();
        let sessions = list_sessions_core(user.id(), &ctx).await?;
        Ok(Json(sessions_with_server_version(sessions)))
    }

    async fn handle_revoke_session<DB: DatabaseAdapter>(
        State(state): State<AuthState<DB>>,
        Extension(ps): Extension<Arc<PluginState>>,
        CurrentSession { user, session }: CurrentSession<DB>,
        ValidatedJson(body): ValidatedJson<RevokeSessionRequest>,
    ) -> Result<axum::response::Response, AuthError> {
        if !ps.config.enable_session_revocation {
            return Err(AuthError::not_found("Not found"));
        }
        let ctx = state.to_context();
        let response = revoke_session_core(&user, &body.token, &ctx).await?;
        if body.token == session.token() {
            let cookie = state.clear_session_cookie();
            return Ok(([(header::SET_COOKIE, cookie)], Json(response)).into_response());
        }
        Ok(Json(response).into_response())
    }

    async fn handle_revoke_sessions<DB: DatabaseAdapter>(
        State(state): State<AuthState<DB>>,
        Extension(ps): Extension<Arc<PluginState>>,
        CurrentSession { user, .. }: CurrentSession<DB>,
    ) -> Result<axum::response::Response, AuthError> {
        if !ps.config.enable_session_revocation {
            return Err(AuthError::not_found("Not found"));
        }
        let ctx = state.to_context();
        let response = revoke_sessions_core(user.id(), &ctx).await?;
        let cookie = state.clear_session_cookie();
        Ok(([(header::SET_COOKIE, cookie)], Json(response)).into_response())
    }

    async fn handle_revoke_other_sessions<DB: DatabaseAdapter>(
        State(state): State<AuthState<DB>>,
        Extension(ps): Extension<Arc<PluginState>>,
        CurrentSession { user, session }: CurrentSession<DB>,
    ) -> Result<Json<StatusResponse>, AuthError> {
        if !ps.config.enable_session_revocation {
            return Err(AuthError::not_found("Not found"));
        }
        let ctx = state.to_context();
        let response = revoke_other_sessions_core(user.id(), &session, &ctx).await?;
        Ok(Json(response))
    }

    impl<DB: DatabaseAdapter> athena_auth_core::AxumPlugin<DB> for SessionManagementPlugin {
        fn name(&self) -> &'static str {
            "session-management"
        }

        fn router(&self) -> axum::Router<AuthState<DB>> {
            use axum::routing::{get, post};

            let plugin_state = Arc::new(PluginState {
                config: self.config.clone(),
            });
            axum::Router::new()
                .route(
                    "/get-session",
                    get(handle_get_session::<DB>).post(handle_get_session::<DB>),
                )
                .route("/sign-out", post(handle_sign_out::<DB>))
                .route("/list-sessions", get(handle_list_sessions::<DB>))
                .route("/revoke-session", post(handle_revoke_session::<DB>))
                .route("/revoke-sessions", post(handle_revoke_sessions::<DB>))
                .route(
                    "/revoke-other-sessions",
                    post(handle_revoke_other_sessions::<DB>),
                )
                .layer(Extension(plugin_state))
        }
    }
}

#[cfg(test)]
#[path = "__tests__/session_management.rs"]
mod tests;
