use athena_auth_core::entity::{AuthAccount, AuthSession, AuthUser, AuthVerification};
use athena_auth_core::{
    AuthContext, AuthError, AuthRequest, AuthResponse, AuthResult, CreateAccount, CreateUser,
    CreateVerification, DatabaseAdapter, UpdateAccount, UpdateUser,
};
use base64::Engine;
use chrono::{Duration, Utc};
use rand::distributions::Alphanumeric;
use rand::{Rng, thread_rng};
use serde::Serialize;
use sha2::{Digest, Sha256};
use std::time::Duration as stdDuration;

use crate::plugins::helpers::{
    extract_session_client_info, resolve_optional_avatar_value_for_sign_in,
    trusted_public_callback_url, write_sign_in_audit_in_background,
};

use super::apple_id_token::verify_apple_id_token;
use super::providers::{OAuthConfig, OAuthUserInfo};
use super::types::{
    AccessTokenResponse, GetAccessTokenRequest, LinkSocialRequest, OAuthCallbackResponse,
    RefreshTokenRequest, RefreshTokenResponse, SocialSignInRedirectResponse, SocialSignInRequest,
    SocialSignInResponse, SocialSignInSessionResponse,
};

use super::encryption::{encrypt_token_set, maybe_decrypt};

const OAUTH_OUTBOUND_USER_AGENT: &str = "athena-auth";
const OAUTH_HTTP_TIMEOUT: stdDuration = stdDuration::from_secs(15);

fn oauth_http_client() -> reqwest::Client {
    reqwest::Client::builder()
        .timeout(OAUTH_HTTP_TIMEOUT)
        .build()
        .unwrap_or_else(|_| reqwest::Client::new())
}

fn oauth_api_get<'a>(
    client: &'a reqwest::Client,
    url: &str,
    access_token: &str,
) -> reqwest::RequestBuilder {
    client
        .get(url)
        .bearer_auth(access_token)
        .header("Accept", "application/json")
        .header("User-Agent", OAUTH_OUTBOUND_USER_AGENT)
}

// ---------------------------------------------------------------------------
// Shared helpers (DRY)
// ---------------------------------------------------------------------------

/// Authenticate the current request and return the validated session.
async fn require_session<DB: DatabaseAdapter>(
    req: &AuthRequest,
    ctx: &AuthContext<DB>,
) -> Result<DB::Session, AuthError> {
    let session_manager: athena_auth_core::SessionManager<DB> = ctx.session_manager();
    let token: String = session_manager
        .extract_session_token(req)
        .ok_or(AuthError::Unauthenticated)?;
    session_manager
        .get_session(&token)
        .await?
        .ok_or(AuthError::Unauthenticated)
}

/// Create session and return the `(OAuthCallbackResponse, token)` tuple for the core function.
async fn create_oauth_session_tuple<DB: DatabaseAdapter>(
    user: DB::User,
    provider_name: &str,
    ip_address: Option<String>,
    user_agent: Option<String>,
    ctx: &AuthContext<DB>,
) -> AuthResult<(OAuthCallbackResponse<DB::User>, String)> {
    let audit_ip: Option<String> = ip_address.clone();
    let audit_user_agent = user_agent.clone();
    let session = ctx
        .session_manager()
        .create_session(&user, ip_address, user_agent)
        .await?;
    let token: String = session.token().to_string();

    write_sign_in_audit_in_background(
        ctx.clone(),
        "auth.sign_in.social",
        Some(user.id().to_string()),
        Some(token.clone()),
        Some("provider".to_string()),
        Some(provider_name.to_string()),
        true,
        None,
        audit_ip,
        audit_user_agent,
        Some(serde_json::json!({
            "email": user.email(),
        })),
    );

    let response = OAuthCallbackResponse {
        token: token.clone(),
        user,
    };

    Ok((response, token))
}

/// Find the account for a specific provider among a user's linked accounts.
fn find_account_for_provider<'a, A: AuthAccount>(
    accounts: &'a [A],
    provider_id: &str,
) -> Result<&'a A, AuthError> {
    accounts
        .iter()
        .find(|a| a.provider_id() == provider_id)
        .ok_or_else(|| {
            AuthError::not_found(format!(
                "No linked account found for provider: {}",
                provider_id
            ))
        })
}

fn generate_pkce() -> (String, String) {
    let verifier: String = thread_rng()
        .sample_iter(&Alphanumeric)
        .take(43)
        .map(char::from)
        .collect();
    let mut hasher = Sha256::new();
    hasher.update(verifier.as_bytes());
    let challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(hasher.finalize());
    (verifier, challenge)
}

fn build_authorization_url(
    config: &OAuthConfig,
    provider_name: &str,
    callback_url: &str,
    scopes: Option<&[String]>,
    state: &str,
    code_challenge: &str,
) -> AuthResult<String> {
    let provider: &super::OAuthProvider = config
        .providers
        .get(provider_name)
        .ok_or_else(|| AuthError::bad_request(format!("Unknown provider: {}", provider_name)))?;

    let effective_scopes: Vec<&str> = scopes
        .map(|s| s.iter().map(|s| s.as_str()).collect())
        .unwrap_or_else(|| provider.scopes.iter().map(|s| s.as_str()).collect());
    let scope_str = effective_scopes.join(" ");

    let mut url = format!(
        "{}?client_id={}&redirect_uri={}&response_type=code&scope={}&state={}&code_challenge={}&code_challenge_method=S256",
        provider.auth_url,
        urlencoding::encode(&provider.client_id),
        urlencoding::encode(callback_url),
        urlencoding::encode(&scope_str),
        urlencoding::encode(state),
        urlencoding::encode(code_challenge),
    );

    for (key, value) in &provider.extra_auth_params {
        url.push('&');
        url.push_str(&urlencoding::encode(key));
        url.push('=');
        url.push_str(&urlencoding::encode(value));
    }

    Ok(url)
}

fn oauth_callback_urls_from_env(raw: &str) -> Vec<String> {
    raw.split(',')
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(str::to_string)
        .collect()
}

fn configured_oauth_callback_urls<DB: DatabaseAdapter>(
    ctx: &AuthContext<DB>,
    provider_name: &str,
) -> Vec<String> {
    let mut urls = Vec::new();
    let env_key = format!("{}_CALLBACK_URL", provider_name.to_uppercase());
    if let Ok(raw) = std::env::var(&env_key) {
        urls.extend(oauth_callback_urls_from_env(raw.as_str()));
    }

    let callback_path = format!("callback/{provider_name}");
    urls.push(ctx.config.public_auth_url(&callback_path));
    urls
}

fn callback_urls_match(requested: &str, allowed: &str) -> bool {
    if requested == allowed {
        return true;
    }

    requested.trim_end_matches('/') == allowed.trim_end_matches('/')
}

async fn validate_registered_callback_url<DB: DatabaseAdapter>(
    ctx: &AuthContext<DB>,
    provider_name: &str,
    callback_url: Option<&str>,
) -> AuthResult<Option<String>> {
    let Some(callback_url) = callback_url
        .map(str::trim)
        .filter(|value| !value.is_empty())
    else {
        return Ok(None);
    };

    if configured_oauth_callback_urls(ctx, provider_name)
        .iter()
        .any(|allowed| callback_urls_match(callback_url, allowed))
    {
        return Ok(Some(callback_url.to_string()));
    }

    let client = ctx
        .database
        .find_active_oauth_client_by_provider_and_redirect_uri(provider_name, callback_url)
        .await?;
    if client.is_none() {
        return Ok(None);
    }

    Ok(Some(callback_url.to_string()))
}

/// Resolve OAuth provider `redirect_uri` and optional app post-auth callback.
///
/// Better Auth-style clients send `callbackURL` as the page to return to after sign-in
/// (for example `https://xbp.app/session`). The provider redirect URI stays on Athena's
/// `/callback/{provider}` route unless the caller explicitly passes a registered OAuth
/// redirect URI.
async fn resolve_social_oauth_urls<DB: DatabaseAdapter>(
    ctx: &AuthContext<DB>,
    provider_name: &str,
    requested_callback_url: Option<&str>,
) -> AuthResult<(String, Option<String>)> {
    let default_redirect_uri = default_oauth_callback_url(ctx, provider_name);

    let Some(requested) = requested_callback_url
        .map(str::trim)
        .filter(|value| !value.is_empty())
    else {
        return Ok((default_redirect_uri, None));
    };

    if configured_oauth_callback_urls(ctx, provider_name)
        .iter()
        .any(|allowed| callback_urls_match(requested, allowed))
    {
        return Ok((requested.to_string(), None));
    }

    if validate_registered_callback_url(ctx, provider_name, Some(requested))
        .await?
        .is_some()
    {
        return Ok((requested.to_string(), None));
    }

    if let Some(app_callback_url) = trusted_public_callback_url(ctx, Some(requested)) {
        return Ok((default_redirect_uri, Some(app_callback_url.to_string())));
    }

    Err(AuthError::bad_request(format!(
        "callbackURL is not registered for provider: {provider_name}"
    )))
}

fn first_oauth_callback_url_from_env(raw: &str) -> Option<String> {
    oauth_callback_urls_from_env(raw).into_iter().next()
}

fn default_oauth_callback_url<DB: DatabaseAdapter>(
    ctx: &AuthContext<DB>,
    provider_name: &str,
) -> String {
    let env_key = format!("{}_CALLBACK_URL", provider_name.to_uppercase());
    if let Ok(url) = std::env::var(&env_key) {
        if let Some(first) = first_oauth_callback_url_from_env(url.as_str()) {
            return first;
        }
    }
    let callback_path = format!("callback/{provider_name}");
    ctx.config.public_auth_url(&callback_path)
}

async fn resolve_github_email_if_needed(
    client: &reqwest::Client,
    access_token: &str,
    provider_name: &str,
    user_info: OAuthUserInfo,
) -> AuthResult<OAuthUserInfo> {
    if provider_name != "github" || !user_info.email.trim().is_empty() {
        return Ok(user_info);
    }

    let emails_resp: reqwest::Response =
        oauth_api_get(client, "https://api.github.com/user/emails", access_token)
            .send()
            .await
            .map_err(|e| AuthError::internal(format!("Failed to fetch GitHub emails: {}", e)))?;

    if !emails_resp.status().is_success() {
        let error_body = emails_resp
            .text()
            .await
            .unwrap_or_else(|_| "Unknown error".to_string());
        return Err(AuthError::internal(format!(
            "GitHub emails request failed: {}",
            error_body
        )));
    }

    let emails: Vec<serde_json::Value> = emails_resp
        .json()
        .await
        .map_err(|e| AuthError::internal(format!("Failed to parse GitHub emails: {}", e)))?;

    let email = emails
        .iter()
        .find(|entry| entry["primary"].as_bool() == Some(true))
        .or_else(|| {
            emails
                .iter()
                .find(|entry| entry["verified"].as_bool() == Some(true))
        })
        .and_then(|entry| entry["email"].as_str())
        .map(String::from)
        .ok_or_else(|| AuthError::internal("GitHub account has no accessible email"))?;

    Ok(OAuthUserInfo { email, ..user_info })
}

// ---------------------------------------------------------------------------
// Core functions
// ---------------------------------------------------------------------------

/// Internal result used to attach a session cookie when sign-in completes inline.
pub(crate) enum SocialSignInCoreResult<U: Serialize> {
    Redirect(SocialSignInRedirectResponse),
    Session(OAuthCallbackResponse<U>, String),
}

pub(crate) async fn social_sign_in_core<DB: DatabaseAdapter>(
    body: &SocialSignInRequest,
    config: &OAuthConfig,
    ctx: &AuthContext<DB>,
    ip_address: Option<String>,
    user_agent: Option<String>,
) -> AuthResult<SocialSignInCoreResult<DB::User>> {
    let provider_name = &body.provider;

    if let Some(id_token) = body.id_token.as_deref().filter(|value| !value.is_empty()) {
        return id_token_sign_in_core(body, config, ctx, id_token, ip_address, user_agent).await;
    }

    let (redirect_uri, app_callback_url) =
        resolve_social_oauth_urls(ctx, provider_name, body.callback_url.as_deref()).await?;

    let redirect = initiate_oauth_flow_core(
        config,
        ctx,
        provider_name,
        &redirect_uri,
        app_callback_url.as_deref(),
        body.scopes.as_deref(),
        None,
    )
    .await?;

    Ok(SocialSignInCoreResult::Redirect(redirect))
}

async fn id_token_sign_in_core<DB: DatabaseAdapter>(
    body: &SocialSignInRequest,
    config: &OAuthConfig,
    ctx: &AuthContext<DB>,
    id_token: &str,
    ip_address: Option<String>,
    user_agent: Option<String>,
) -> AuthResult<SocialSignInCoreResult<DB::User>> {
    let provider_name = &body.provider;
    let provider = config
        .providers
        .get(provider_name)
        .ok_or_else(|| AuthError::bad_request(format!("Unknown provider: {provider_name}")))?;

    let user_info = match provider_name.as_str() {
        "apple" => verify_apple_id_token(id_token, &provider.client_id).await?,
        _ => {
            return Err(AuthError::bad_request(format!(
                "ID token sign-in is not supported for provider: {provider_name}"
            )));
        }
    };

    let (response, token) = sign_in_with_oauth_user_info(
        ctx,
        provider_name,
        user_info,
        None,
        None,
        Some(id_token.to_string()),
        None,
        None,
        None,
        ip_address,
        user_agent,
    )
    .await?;

    Ok(SocialSignInCoreResult::Session(response, token))
}

pub(crate) fn social_sign_in_core_result_to_response<U: Serialize>(
    result: SocialSignInCoreResult<U>,
) -> (SocialSignInResponse<U>, Option<String>) {
    match result {
        SocialSignInCoreResult::Redirect(redirect) => {
            (SocialSignInResponse::Redirect(redirect), None)
        }
        SocialSignInCoreResult::Session(response, token) => {
            let session = SocialSignInSessionResponse {
                redirect: false,
                token: response.token.clone(),
                user: response.user,
            };
            (SocialSignInResponse::Session(session), Some(token))
        }
    }
}

pub(crate) async fn link_social_core<DB: DatabaseAdapter>(
    body: &LinkSocialRequest,
    session: &DB::Session,
    config: &OAuthConfig,
    ctx: &AuthContext<DB>,
) -> AuthResult<SocialSignInRedirectResponse> {
    let provider_name = &body.provider;

    let (redirect_uri, app_callback_url) =
        resolve_social_oauth_urls(ctx, provider_name, body.callback_url.as_deref()).await?;

    initiate_oauth_flow_core(
        config,
        ctx,
        provider_name,
        &redirect_uri,
        app_callback_url.as_deref(),
        body.scopes.as_deref(),
        Some(session.user_id()),
    )
    .await
}

pub(crate) async fn get_access_token_core<DB: DatabaseAdapter>(
    body: &GetAccessTokenRequest,
    session: &DB::Session,
    ctx: &AuthContext<DB>,
) -> AuthResult<AccessTokenResponse> {
    let accounts = ctx.database.get_user_accounts(session.user_id()).await?;
    let account = find_account_for_provider(&accounts, &body.provider_id)?;

    let encrypt: bool = ctx.config.account.encrypt_oauth_tokens;
    let secret = &ctx.config.secret;

    Ok(AccessTokenResponse {
        access_token: maybe_decrypt(account.access_token(), encrypt, secret)?,
        access_token_expires_at: account.access_token_expires_at().map(|dt| dt.to_rfc3339()),
        scope: account.scope().map(String::from),
    })
}

pub(crate) async fn refresh_token_core<DB: DatabaseAdapter>(
    body: &RefreshTokenRequest,
    session: &DB::Session,
    config: &OAuthConfig,
    ctx: &AuthContext<DB>,
) -> AuthResult<RefreshTokenResponse> {
    let provider_name = &body.provider_id;

    let provider = config
        .providers
        .get(provider_name)
        .ok_or_else(|| AuthError::bad_request(format!("Unknown provider: {}", provider_name)))?;

    let accounts = ctx.database.get_user_accounts(session.user_id()).await?;
    let account = find_account_for_provider(&accounts, provider_name)?;

    let encrypt = ctx.config.account.encrypt_oauth_tokens;
    let secret = &ctx.config.secret;

    let current_refresh_token = maybe_decrypt(account.refresh_token(), encrypt, secret)?
        .ok_or_else(|| AuthError::bad_request("No refresh token available for this provider"))?;

    let client = oauth_http_client();
    let token_resp = client
        .post(&provider.token_url)
        .header("Accept", "application/json")
        .form(&[
            ("grant_type", "refresh_token"),
            ("refresh_token", &current_refresh_token),
            ("client_id", &provider.client_id),
            ("client_secret", &provider.client_secret),
        ])
        .send()
        .await
        .map_err(|e| AuthError::internal(format!("Token refresh failed: {}", e)))?;

    if !token_resp.status().is_success() {
        let error_body = token_resp
            .text()
            .await
            .unwrap_or_else(|_| "Unknown error".to_string());
        return Err(AuthError::internal(format!(
            "Token refresh returned error: {}",
            error_body
        )));
    }

    let token_data: serde_json::Value = token_resp
        .json()
        .await
        .map_err(|e| AuthError::internal(format!("Failed to parse refresh response: {}", e)))?;

    let new_access_token = token_data["access_token"]
        .as_str()
        .ok_or_else(|| AuthError::internal("Missing access_token in refresh response"))?;

    let new_refresh_token = token_data["refresh_token"].as_str().map(String::from);
    let expires_in = token_data["expires_in"].as_i64();
    let new_scope = token_data["scope"].as_str().map(String::from);

    let access_token_expires_at = expires_in.map(|secs| Utc::now() + Duration::seconds(secs));

    let tokens = encrypt_token_set(
        ctx,
        Some(new_access_token.to_string()),
        new_refresh_token.clone(),
        None,
    )?;
    ctx.database
        .update_account(
            account.id(),
            UpdateAccount {
                access_token: tokens.access_token,
                refresh_token: tokens.refresh_token,
                access_token_expires_at,
                scope: new_scope.clone(),
                ..Default::default()
            },
        )
        .await?;

    Ok(RefreshTokenResponse {
        access_token: Some(new_access_token.to_string()),
        access_token_expires_at: access_token_expires_at.map(|dt| dt.to_rfc3339()),
        refresh_token: new_refresh_token,
        scope: new_scope,
    })
}

/// Shared logic for social sign-in and link-social flows.
///
/// Both flows build a verification payload, store it, construct the
/// authorization URL, and return a redirect response. The only difference
/// is `link_user_id` (None for sign-in, Some for linking).
async fn initiate_oauth_flow_core<DB: DatabaseAdapter>(
    config: &OAuthConfig,
    ctx: &AuthContext<DB>,
    provider_name: &str,
    redirect_uri: &str,
    app_callback_url: Option<&str>,
    scopes: Option<&[String]>,
    link_user_id: Option<&str>,
) -> AuthResult<SocialSignInRedirectResponse> {
    let provider = config
        .providers
        .get(provider_name)
        .ok_or_else(|| AuthError::bad_request(format!("Unknown provider: {}", provider_name)))?;

    let (code_verifier, code_challenge) = generate_pkce();
    let state: String = uuid::Uuid::new_v4().to_string();

    let effective_scopes: Vec<String> = scopes
        .map(|s| s.to_vec())
        .unwrap_or_else(|| provider.scopes.clone());

    let mut payload = serde_json::json!({
        "provider": provider_name,
        "redirect_uri": redirect_uri,
        "callback_url": redirect_uri,
        "code_verifier": code_verifier,
        "link_user_id": link_user_id,
        "scopes": effective_scopes.join(" "),
    });
    if let Some(app_callback_url) = app_callback_url {
        payload["app_callback_url"] = serde_json::Value::String(app_callback_url.to_string());
    }

    ctx.database
        .create_verification(CreateVerification {
            identifier: format!("oauth:{}", state),
            value: payload.to_string(),
            expires_at: Utc::now() + Duration::minutes(10),
        })
        .await?;

    let url = build_authorization_url(
        config,
        provider_name,
        redirect_uri,
        scopes,
        &state,
        &code_challenge,
    )?;

    Ok(SocialSignInRedirectResponse {
        url,
        redirect: true,
    })
}

async fn sign_in_with_oauth_user_info<DB: DatabaseAdapter>(
    ctx: &AuthContext<DB>,
    provider_name: &str,
    user_info: OAuthUserInfo,
    access_token: Option<String>,
    refresh_token: Option<String>,
    id_token: Option<String>,
    access_token_expires_at: Option<chrono::DateTime<Utc>>,
    scope: Option<String>,
    link_user_id: Option<String>,
    ip_address: Option<String>,
    user_agent: Option<String>,
) -> AuthResult<(OAuthCallbackResponse<DB::User>, String)> {
    if let Some(link_user_id) = link_user_id {
        let user = ctx
            .database
            .get_user_by_id(&link_user_id)
            .await?
            .ok_or(AuthError::UserNotFound)?;

        if ctx
            .database
            .get_account(provider_name, &user_info.id)
            .await?
            .is_some()
        {
            return Err(AuthError::conflict(
                "This social account is already linked to a user",
            ));
        }

        let tokens = encrypt_token_set(ctx, access_token, refresh_token, id_token)?;
        ctx.database
            .create_account(CreateAccount {
                user_id: link_user_id,
                account_id: user_info.id,
                provider_id: provider_name.to_string(),
                access_token: tokens.access_token,
                refresh_token: tokens.refresh_token,
                id_token: tokens.id_token,
                access_token_expires_at,
                refresh_token_expires_at: None,
                scope,
                password: None,
            })
            .await?;

        return create_oauth_session_tuple(user, provider_name, ip_address, user_agent, ctx).await;
    }

    if let Some(existing_account) = ctx
        .database
        .get_account(provider_name, &user_info.id)
        .await?
    {
        if ctx.config.account.update_account_on_sign_in {
            let tokens = encrypt_token_set(
                ctx,
                access_token.clone(),
                refresh_token.clone(),
                id_token.clone(),
            )?;
            ctx.database
                .update_account(
                    existing_account.id(),
                    UpdateAccount {
                        access_token: tokens.access_token,
                        refresh_token: tokens.refresh_token,
                        id_token: tokens.id_token,
                        access_token_expires_at,
                        scope: scope.clone(),
                        ..Default::default()
                    },
                )
                .await?;
        }

        let user = ctx
            .database
            .get_user_by_id(existing_account.user_id())
            .await?
            .ok_or(AuthError::UserNotFound)?;

        return create_oauth_session_tuple(user, provider_name, ip_address, user_agent, ctx).await;
    }

    if user_info.email.trim().is_empty() {
        return Err(AuthError::bad_request(
            "Email is required to create a new account with this provider",
        ));
    }

    let linking_cfg = &ctx.config.account.account_linking;

    let user = if let Some(existing_user) = ctx.database.get_user_by_email(&user_info.email).await?
    {
        if linking_cfg.enabled {
            let provider_trusted = linking_cfg.trusted_providers.is_empty()
                || linking_cfg
                    .trusted_providers
                    .iter()
                    .any(|p| p == provider_name);

            if !provider_trusted {
                return Err(AuthError::bad_request(
                    "Account linking is not allowed for this provider",
                ));
            }

            if linking_cfg.update_user_info_on_link {
                let mut update = UpdateUser::default();
                if let Some(name) = &user_info.name {
                    update.name = Some(name.clone());
                }
                if let Some(image) = resolve_optional_avatar_value_for_sign_in(
                    ctx,
                    user_info.image.as_deref(),
                    athena_auth_core::AvatarTemplateContext::for_user(existing_user.id()),
                )
                .await?
                {
                    update.image = Some(image);
                }
                let updated_user = ctx.database.update_user(existing_user.id(), update).await?;
                if let Err(cache_error) = ctx.cache_session_user(updated_user.id(), &updated_user) {
                    tracing::warn!(
                        user_id = %updated_user.id(),
                        error = %cache_error,
                        "failed to refresh session user cache after oauth link update"
                    );
                }
                updated_user
            } else {
                existing_user
            }
        } else {
            return Err(AuthError::bad_request(
                "Account linking is disabled. Cannot sign in with a new provider for an existing email.",
            ));
        }
    } else {
        let mut create_user = CreateUser::new()
            .with_email(&user_info.email)
            .with_name(user_info.name.as_deref().unwrap_or(&user_info.email))
            .with_email_verified(user_info.email_verified);
        let pending_user_id = create_user
            .id
            .clone()
            .expect("CreateUser::new should always assign a generated user id");
        let image = resolve_optional_avatar_value_for_sign_in(
            ctx,
            user_info.image.as_deref(),
            athena_auth_core::AvatarTemplateContext::for_user(&pending_user_id),
        )
        .await?;
        if let Some(image) = image {
            create_user = create_user.with_image(image);
        }

        ctx.database.create_user(create_user).await?
    };

    let tokens = encrypt_token_set(ctx, access_token, refresh_token, id_token)?;
    ctx.database
        .create_account(CreateAccount {
            user_id: user.id().to_string(),
            account_id: user_info.id,
            provider_id: provider_name.to_string(),
            access_token: tokens.access_token,
            refresh_token: tokens.refresh_token,
            id_token: tokens.id_token,
            access_token_expires_at,
            refresh_token_expires_at: None,
            scope,
            password: None,
        })
        .await?;

    create_oauth_session_tuple(user, provider_name, ip_address, user_agent, ctx).await
}

// ---------------------------------------------------------------------------
// Old handlers (rewritten to call core)
// ---------------------------------------------------------------------------

pub async fn handle_social_sign_in<DB: DatabaseAdapter>(
    config: &OAuthConfig,
    req: &AuthRequest,
    ctx: &AuthContext<DB>,
) -> AuthResult<AuthResponse> {
    let body: SocialSignInRequest = match athena_auth_core::validate_request_body(req) {
        Ok(v) => v,
        Err(resp) => return Ok(resp),
    };
    let (ip_address, user_agent) = extract_session_client_info(ctx, req);
    let result = social_sign_in_core(&body, config, ctx, ip_address, user_agent).await?;
    let (response, session_token) = social_sign_in_core_result_to_response(result);
    let mut auth_response = AuthResponse::json(200, &response).map_err(AuthError::from)?;
    if let Some(token) = session_token {
        let cookie_header =
            athena_auth_core::utils::cookie_utils::create_session_cookie(&token, &ctx.config);
        auth_response = auth_response.with_header("Set-Cookie", cookie_header);
    }
    Ok(auth_response)
}

pub(crate) async fn callback_core<DB: DatabaseAdapter>(
    code: &str,
    state_param: &str,
    provider_name: &str,
    config: &OAuthConfig,
    ip_address: Option<String>,
    user_agent: Option<String>,
    ctx: &AuthContext<DB>,
) -> AuthResult<(OAuthCallbackResponse<DB::User>, String, Option<String>)> {
    // Look up state verification
    let verification = ctx
        .database
        .get_verification_by_identifier(&format!("oauth:{}", state_param))
        .await?
        .ok_or_else(|| AuthError::bad_request("Invalid or expired OAuth state"))?;

    let payload: serde_json::Value = serde_json::from_str(verification.value())
        .map_err(|e| AuthError::internal(format!("Invalid state payload: {}", e)))?;

    let stored_provider = payload["provider"]
        .as_str()
        .ok_or_else(|| AuthError::internal("Missing provider in state"))?;

    if stored_provider != provider_name {
        return Err(AuthError::bad_request("Provider mismatch"));
    }

    let redirect_uri = payload["redirect_uri"]
        .as_str()
        .or_else(|| payload["callback_url"].as_str())
        .ok_or_else(|| AuthError::internal("Missing redirect_uri in state"))?;

    let app_callback_url = payload["app_callback_url"].as_str().map(str::to_string);

    let code_verifier = payload["code_verifier"]
        .as_str()
        .ok_or_else(|| AuthError::internal("Missing code_verifier in state"))?;

    let link_user_id = payload["link_user_id"].as_str().map(String::from);

    let scopes = payload["scopes"].as_str().map(String::from);

    // Delete the verification now that we've used it
    ctx.database.delete_verification(verification.id()).await?;

    let provider = config
        .providers
        .get(provider_name)
        .ok_or_else(|| AuthError::bad_request(format!("Unknown provider: {}", provider_name)))?;

    // Exchange code for tokens
    let client = oauth_http_client();
    let token_resp = client
        .post(&provider.token_url)
        .header("Accept", "application/json")
        .form(&[
            ("grant_type", "authorization_code"),
            ("code", code),
            ("redirect_uri", redirect_uri),
            ("client_id", &provider.client_id),
            ("client_secret", &provider.client_secret),
            ("code_verifier", code_verifier),
        ])
        .send()
        .await
        .map_err(|e| AuthError::internal(format!("Token exchange failed: {}", e)))?;

    if !token_resp.status().is_success() {
        let error_body = token_resp
            .text()
            .await
            .unwrap_or_else(|_| "Unknown error".to_string());
        return Err(AuthError::internal(format!(
            "Token exchange returned error: {}",
            error_body
        )));
    }

    let token_data: serde_json::Value = token_resp
        .json()
        .await
        .map_err(|e| AuthError::internal(format!("Failed to parse token response: {}", e)))?;

    let access_token = token_data["access_token"]
        .as_str()
        .ok_or_else(|| AuthError::internal("Missing access_token in token response"))?;

    let refresh_token = token_data["refresh_token"].as_str().map(String::from);
    let id_token = token_data["id_token"].as_str().map(String::from);
    let expires_in = token_data["expires_in"].as_i64();

    // Fetch user info
    let user_info_resp = oauth_api_get(&client, &provider.user_info_url, access_token)
        .send()
        .await
        .map_err(|e| AuthError::internal(format!("Failed to fetch user info: {}", e)))?;

    if !user_info_resp.status().is_success() {
        let error_body = user_info_resp
            .text()
            .await
            .unwrap_or_else(|_| "Unknown error".to_string());
        return Err(AuthError::internal(format!(
            "User info request failed: {}",
            error_body
        )));
    }

    let user_info_json: serde_json::Value = user_info_resp
        .json()
        .await
        .map_err(|e| AuthError::internal(format!("Failed to parse user info: {}", e)))?;

    let user_info = (provider.map_user_info)(user_info_json)
        .map_err(|e| AuthError::internal(format!("Failed to map user info: {}", e)))?;
    let user_info =
        resolve_github_email_if_needed(&client, access_token, provider_name, user_info).await?;

    let access_token_expires_at = expires_in.map(|secs| Utc::now() + Duration::seconds(secs));

    let (response, token) = sign_in_with_oauth_user_info(
        ctx,
        provider_name,
        user_info,
        Some(access_token.to_string()),
        refresh_token,
        id_token,
        access_token_expires_at,
        scopes,
        link_user_id,
        ip_address,
        user_agent,
    )
    .await?;

    Ok((response, token, app_callback_url))
}

pub async fn handle_callback<DB: DatabaseAdapter>(
    config: &OAuthConfig,
    provider_name: &str,
    req: &AuthRequest,
    ctx: &AuthContext<DB>,
) -> AuthResult<AuthResponse> {
    let (ip_address, user_agent) = extract_session_client_info(ctx, req);
    let code = match req
        .query
        .get("code")
        .cloned()
        .or_else(|| extract_query_param(req.path(), "code"))
    {
        Some(code) => code,
        None => {
            write_sign_in_audit_in_background(
                ctx.clone(),
                "auth.sign_in.social",
                None,
                None,
                Some("provider".to_string()),
                Some(provider_name.to_string()),
                false,
                Some("Missing code parameter".to_string()),
                ip_address.clone(),
                user_agent.clone(),
                Some(serde_json::json!({
                    "phase": "oauth.callback",
                })),
            );
            return Err(AuthError::bad_request("Missing code parameter"));
        }
    };

    let state = match req
        .query
        .get("state")
        .cloned()
        .or_else(|| extract_query_param(req.path(), "state"))
    {
        Some(state) => state,
        None => {
            write_sign_in_audit_in_background(
                ctx.clone(),
                "auth.sign_in.social",
                None,
                None,
                Some("provider".to_string()),
                Some(provider_name.to_string()),
                false,
                Some("Missing state parameter".to_string()),
                ip_address.clone(),
                user_agent.clone(),
                Some(serde_json::json!({
                    "phase": "oauth.callback",
                })),
            );
            return Err(AuthError::bad_request("Missing state parameter"));
        }
    };

    let audit_ip = ip_address.clone();
    let audit_user_agent = user_agent.clone();
    let (response, token, app_callback_url) = match callback_core(
        &code,
        &state,
        provider_name,
        config,
        ip_address,
        user_agent,
        ctx,
    )
    .await
    {
        Ok(result) => result,
        Err(err) => {
            write_sign_in_audit_in_background(
                ctx.clone(),
                "auth.sign_in.social",
                None,
                None,
                Some("provider".to_string()),
                Some(provider_name.to_string()),
                false,
                Some(err.to_string()),
                audit_ip.clone(),
                audit_user_agent.clone(),
                Some(serde_json::json!({
                    "phase": "oauth.callback",
                })),
            );
            return Err(err);
        }
    };
    let cookie_header =
        athena_auth_core::utils::cookie_utils::create_session_cookie(&token, &ctx.config);

    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_name),
        );
        return Ok(AuthResponse::text(302, "Redirect")
            .with_header("Set-Cookie", cookie_header)
            .with_header("Location", location));
    }

    Ok(AuthResponse::json(200, &response)?.with_header("Set-Cookie", cookie_header))
}

pub async fn handle_link_social<DB: DatabaseAdapter>(
    config: &OAuthConfig,
    req: &AuthRequest,
    ctx: &AuthContext<DB>,
) -> AuthResult<AuthResponse> {
    let session = require_session(req, ctx).await?;
    let body: LinkSocialRequest = match athena_auth_core::validate_request_body(req) {
        Ok(v) => v,
        Err(resp) => return Ok(resp),
    };
    let response = link_social_core(&body, &session, config, ctx).await?;
    AuthResponse::json(200, &response).map_err(AuthError::from)
}

pub async fn handle_get_access_token<DB: DatabaseAdapter>(
    config: &OAuthConfig,
    req: &AuthRequest,
    ctx: &AuthContext<DB>,
) -> AuthResult<AuthResponse> {
    let _ = config;
    let session = require_session(req, ctx).await?;
    let body: GetAccessTokenRequest = match athena_auth_core::validate_request_body(req) {
        Ok(v) => v,
        Err(resp) => return Ok(resp),
    };
    let response = get_access_token_core(&body, &session, ctx).await?;
    AuthResponse::json(200, &response).map_err(AuthError::from)
}

pub async fn handle_refresh_token<DB: DatabaseAdapter>(
    config: &OAuthConfig,
    req: &AuthRequest,
    ctx: &AuthContext<DB>,
) -> AuthResult<AuthResponse> {
    let session = require_session(req, ctx).await?;
    let body: RefreshTokenRequest = match athena_auth_core::validate_request_body(req) {
        Ok(v) => v,
        Err(resp) => return Ok(resp),
    };
    let response: RefreshTokenResponse = refresh_token_core(&body, &session, config, ctx).await?;
    AuthResponse::json(200, &response).map_err(AuthError::from)
}

/// Extract a query parameter from a path string using the `url` crate.
fn extract_query_param(path: &str, key: &str) -> Option<String> {
    // url::Url requires a base; use a dummy scheme+host so relative paths parse.
    let full = if path.starts_with("http") {
        path.to_string()
    } else {
        format!("http://x{}", path)
    };
    let parsed = url::Url::parse(&full).ok()?;
    parsed
        .query_pairs()
        .find(|(k, _)| k == key)
        .map(|(_, v)| v.into_owned())
}

#[cfg(test)]
mod tests {
    use super::{OAUTH_OUTBOUND_USER_AGENT, oauth_api_get};

    #[test]
    fn oauth_api_get_sets_required_headers() {
        let client = reqwest::Client::new();
        let request = oauth_api_get(&client, "https://example.com/user", "test-token")
            .build()
            .expect("request should build");

        assert_eq!(
            request
                .headers()
                .get("user-agent")
                .and_then(|value| value.to_str().ok()),
            Some(OAUTH_OUTBOUND_USER_AGENT)
        );
        assert_eq!(
            request
                .headers()
                .get("accept")
                .and_then(|value| value.to_str().ok()),
            Some("application/json")
        );
        assert_eq!(
            request
                .headers()
                .get("authorization")
                .and_then(|value| value.to_str().ok()),
            Some("Bearer test-token")
        );
    }
}
