use async_trait::async_trait;

use athena_auth_core::AuthResult;
use athena_auth_core::adapters::DatabaseAdapter;
use athena_auth_core::{AuthContext, AuthPlugin, AuthRoute};
use athena_auth_core::{AuthRequest, AuthResponse, HttpMethod};

#[cfg(feature = "axum")]
use athena_auth_core::plugin::{AuthState, AxumPlugin};

mod apple_id_token;
pub mod encryption;
mod handlers;
mod providers;
mod types;

pub use providers::{OAuthConfig, OAuthProvider, OAuthStateStrategy, OAuthUserInfo};

pub struct OAuthPlugin {
    config: OAuthConfig,
}

impl OAuthPlugin {
    pub fn new() -> Self {
        Self {
            config: OAuthConfig::default(),
        }
    }

    pub fn with_config(config: OAuthConfig) -> Self {
        Self { config }
    }

    pub fn add_provider(mut self, name: &str, provider: OAuthProvider) -> Self {
        self.config.providers.insert(name.to_string(), provider);
        self
    }
}

impl Default for OAuthPlugin {
    fn default() -> Self {
        Self::new()
    }
}

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

    fn routes(&self) -> Vec<AuthRoute> {
        vec![
            AuthRoute::post("/sign-in/social", "social_sign_in"),
            AuthRoute::get("/callback/{provider}", "oauth_callback"),
            AuthRoute::post("/link-social", "link_social"),
            AuthRoute::post("/get-access-token", "get_access_token"),
            AuthRoute::post("/refresh-token", "refresh_token"),
        ]
    }

    async fn on_request(
        &self,
        req: &AuthRequest,
        ctx: &AuthContext<DB>,
    ) -> AuthResult<Option<AuthResponse>> {
        match (req.method(), req.path()) {
            (HttpMethod::Post, "/sign-in/social") => Ok(Some(
                handlers::handle_social_sign_in(&self.config, req, ctx).await?,
            )),
            (HttpMethod::Get, path) if path_matches_callback(path) => {
                let provider = extract_provider_from_callback(path);
                Ok(Some(
                    handlers::handle_callback(&self.config, &provider, req, ctx).await?,
                ))
            }
            (HttpMethod::Post, "/link-social") => Ok(Some(
                handlers::handle_link_social(&self.config, req, ctx).await?,
            )),
            (HttpMethod::Post, "/get-access-token") => Ok(Some(
                handlers::handle_get_access_token(&self.config, req, ctx).await?,
            )),
            (HttpMethod::Post, "/refresh-token") => Ok(Some(
                handlers::handle_refresh_token(&self.config, req, ctx).await?,
            )),
            _ => Ok(None),
        }
    }
}

/// Check if the path matches `/callback/{provider}` (with optional query string).
fn path_matches_callback(path: &str) -> bool {
    let path_without_query = path.split('?').next().unwrap_or(path);
    path_without_query.starts_with("/callback/") && path_without_query.len() > "/callback/".len()
}

/// Extract the provider name from `/callback/{provider}?...`.
fn extract_provider_from_callback(path: &str) -> String {
    let path_without_query = path.split('?').next().unwrap_or(path);
    path_without_query["/callback/".len()..].to_string()
}

// ---------------------------------------------------------------------------
// Axum-native routing (feature-gated)
// ---------------------------------------------------------------------------

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

    use athena_auth_core::error::AuthError;
    use athena_auth_core::extractors::{CurrentSession, ValidatedJson};
    use axum::Json;
    use axum::extract::{Extension, Path, Query, State};
    use axum::http::HeaderMap;
    use axum::http::header;
    use axum::response::IntoResponse;

    use super::handlers::{
        callback_core, get_access_token_core, link_social_core, refresh_token_core,
        social_sign_in_core, social_sign_in_core_result_to_response,
    };
    use super::types::{
        AccessTokenResponse, GetAccessTokenRequest, LinkSocialRequest, RefreshTokenRequest,
        RefreshTokenResponse, SocialSignInRequest, SocialSignInResponse,
    };

    #[derive(serde::Deserialize)]
    struct CallbackQuery {
        code: String,
        state: String,
    }

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

    async fn handle_social_sign_in<DB: DatabaseAdapter>(
        State(state): State<AuthState<DB>>,
        Extension(ps): Extension<Arc<PluginState>>,
        headers: HeaderMap,
        ValidatedJson(body): ValidatedJson<SocialSignInRequest>,
    ) -> Result<impl IntoResponse, AuthError> {
        let ctx = state.to_context();
        let (ip_address, user_agent) = extract_session_client_info_from_headers(&ctx, &headers);
        let result = social_sign_in_core(&body, &ps.config, &ctx, ip_address, user_agent).await?;
        let (response, session_token) = social_sign_in_core_result_to_response(result);
        if let Some(token) = session_token {
            let cookie = state.session_cookie(&token);
            Ok(([(header::SET_COOKIE, cookie)], Json(response)).into_response())
        } else {
            Ok(Json(response).into_response())
        }
    }

    async fn handle_callback<DB: DatabaseAdapter>(
        State(state): State<AuthState<DB>>,
        Extension(ps): Extension<Arc<PluginState>>,
        headers: HeaderMap,
        Path(provider): Path<String>,
        Query(params): Query<CallbackQuery>,
    ) -> Result<impl IntoResponse, AuthError> {
        let ctx = state.to_context();
        let (ip_address, user_agent) = extract_session_client_info_from_headers(&ctx, &headers);
        let (response, token, app_callback_url) = callback_core(
            &params.code,
            &params.state,
            &provider,
            &ps.config,
            ip_address,
            user_agent,
            &ctx,
        )
        .await?;
        let cookie = state.session_cookie(&token);
        if let Some(app_callback_url) = app_callback_url {
            let separator = if app_callback_url.contains('?') {
                '&'
            } else {
                '?'
            };
            let location = format!(
                "{app_callback_url}{separator}token={}&provider={}",
                urlencoding::encode(&token),
                urlencoding::encode(&provider),
            );
            return Ok((
                axum::http::StatusCode::FOUND,
                [(header::SET_COOKIE, cookie), (header::LOCATION, location)],
            )
                .into_response());
        }
        Ok(([(header::SET_COOKIE, cookie)], Json(response)).into_response())
    }

    async fn handle_link_social<DB: DatabaseAdapter>(
        State(state): State<AuthState<DB>>,
        Extension(ps): Extension<Arc<PluginState>>,
        CurrentSession { session, .. }: CurrentSession<DB>,
        ValidatedJson(body): ValidatedJson<LinkSocialRequest>,
    ) -> Result<Json<SocialSignInResponse<DB::User>>, AuthError> {
        let ctx = state.to_context();
        let result = link_social_core(&body, &session, &ps.config, &ctx).await?;
        Ok(Json(SocialSignInResponse::Redirect(result)))
    }

    async fn handle_get_access_token<DB: DatabaseAdapter>(
        State(state): State<AuthState<DB>>,
        CurrentSession { session, .. }: CurrentSession<DB>,
        ValidatedJson(body): ValidatedJson<GetAccessTokenRequest>,
    ) -> Result<Json<AccessTokenResponse>, AuthError> {
        let ctx = state.to_context();
        let result = get_access_token_core(&body, &session, &ctx).await?;
        Ok(Json(result))
    }

    async fn handle_refresh_token<DB: DatabaseAdapter>(
        State(state): State<AuthState<DB>>,
        Extension(ps): Extension<Arc<PluginState>>,
        CurrentSession { session, .. }: CurrentSession<DB>,
        ValidatedJson(body): ValidatedJson<RefreshTokenRequest>,
    ) -> Result<Json<RefreshTokenResponse>, AuthError> {
        let ctx = state.to_context();
        let result = refresh_token_core(&body, &session, &ps.config, &ctx).await?;
        Ok(Json(result))
    }

    #[async_trait]
    impl<DB: DatabaseAdapter> AxumPlugin<DB> for OAuthPlugin {
        fn name(&self) -> &'static str {
            "oauth"
        }

        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("/sign-in/social", post(handle_social_sign_in::<DB>))
                .route("/callback/:provider", get(handle_callback::<DB>))
                .route("/link-social", post(handle_link_social::<DB>))
                .route("/get-access-token", post(handle_get_access_token::<DB>))
                .route("/refresh-token", post(handle_refresh_token::<DB>))
                .layer(Extension(plugin_state))
        }
    }
}
