use axum::{
    extract::State,
    http::StatusCode,
    response::Json,
};
use base64::Engine;
use nanoid::nanoid;
use serde::{Deserialize, Serialize};
use serde_json::json;
use sha2::{Digest, Sha256};
use std::sync::Arc;

use crate::state::AppState;
use crate::server::{AccessToken, TokenClaims};

/// Token request parameters
#[derive(Debug, Deserialize)]
pub struct TokenRequest {
    pub grant_type: String,
    /// Authorization code (required for authorization_code grant)
    #[serde(default)]
    pub code: String,
    pub client_id: String,
    pub redirect_uri: String,
    /// PKCE code verifier (required for authorization_code grant)
    #[serde(default)]
    pub code_verifier: String,
    /// Refresh token (required for refresh_token grant)
    #[serde(default)]
    pub refresh_token: String,
}

/// Token response
#[derive(Debug, Serialize)]
pub struct TokenResponse {
    pub access_token: String,
    pub token_type: String,
    pub me: String,
    pub scope: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub profile: Option<serde_json::Value>,
    pub expires_in: i64,
    /// OAuth 2.0 refresh token for obtaining new access tokens
    #[serde(skip_serializing_if = "Option::is_none")]
    pub refresh_token: Option<String>,
    /// Lifetime of the refresh token in seconds (optional, typically longer than access token)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub refresh_token_expires_in: Option<i64>,
}

/// Error response
#[derive(Debug, Serialize)]
pub struct ErrorResponse {
    pub error: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error_description: Option<String>,
}

/// POST handler for token endpoint
/// Supports both authorization_code and refresh_token grant types
pub async fn token_handler(
    State(state): State<Arc<AppState>>,
    axum::Form(request): axum::Form<TokenRequest>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
    match request.grant_type.as_str() {
        "authorization_code" => handle_authorization_code_grant(&state, &request).await,
        "refresh_token" => handle_refresh_token_grant(&state, &request).await,
        _ => Err((
            StatusCode::BAD_REQUEST,
            Json(ErrorResponse {
                error: "unsupported_grant_type".to_string(),
                error_description: Some("Only authorization_code and refresh_token grant types are supported".to_string()),
            }),
        )),
    }
}

/// Handle authorization_code grant type
async fn handle_authorization_code_grant(
    state: &Arc<AppState>,
    request: &TokenRequest,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
    // Consume the authorization code
    let code_data = match state.consume_code(&request.code) {
        Some(data) => data,
        None => {
            return Err((
                StatusCode::BAD_REQUEST,
                Json(ErrorResponse {
                    error: "invalid_grant".to_string(),
                    error_description: Some("Invalid or expired authorization code".to_string()),
                }),
            ));
        }
    };

    // Verify code hasn't expired
    if code_data.is_expired() {
        return Err((
            StatusCode::BAD_REQUEST,
            Json(ErrorResponse {
                error: "invalid_grant".to_string(),
                error_description: Some("Authorization code has expired".to_string()),
            }),
        ));
    }

    // Verify client_id matches
    if code_data.client_id.as_str() != request.client_id {
        return Err((
            StatusCode::BAD_REQUEST,
            Json(ErrorResponse {
                error: "invalid_client".to_string(),
                error_description: Some("Client ID mismatch".to_string()),
            }),
        ));
    }

    // Verify redirect_uri matches
    if code_data.redirect_uri.to_string() != request.redirect_uri {
        return Err((
            StatusCode::BAD_REQUEST,
            Json(ErrorResponse {
                error: "invalid_grant".to_string(),
                error_description: Some("Redirect URI mismatch".to_string()),
            }),
        ));
    }

    // Verify PKCE code_challenge/code_verifier
    if !verify_pkce(&code_data.code_challenge, &code_data.code_challenge_method, &request.code_verifier) {
        return Err((
            StatusCode::BAD_REQUEST,
            Json(ErrorResponse {
                error: "invalid_grant".to_string(),
                error_description: Some("PKCE verification failed".to_string()),
            }),
        ));
    }

    // Get the user profile
    let user = if let Some(me_url) = &code_data.me {
        state.users.values().find(|u| u.profile.me == *me_url)
    } else {
        None
    };

    let user = match user {
        Some(u) => u,
        None => {
            // Fallback: find user by matching the me URL domain/path
            state.users.values()
                .find(|u| {
                    code_data.me.as_ref()
                        .map(|me| me.host_str() == u.profile.me.host_str())
                        .unwrap_or(false)
                })
                .ok_or_else(|| {
                    (
                        StatusCode::INTERNAL_SERVER_ERROR,
                        Json(ErrorResponse {
                            error: "server_error".to_string(),
                            error_description: Some("User profile not found".to_string()),
                        }),
                    )
                })?
        }
    };

    // Generate tokens
    let refresh_token = generate_refresh_token();
    let expires_in = 3600i64; // 1 hour
    let refresh_token_expires_in = 86400i64 * 30; // 30 days

    // Build token claims
    let claims = TokenClaims {
        me: user.profile.me.clone(),
        scope: code_data.scope.clone(),
        client_id: code_data.client_id.clone(),
        issued_at: std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs(),
        expires_at: std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs() + expires_in as u64,
    };

    // Sign the token
    let signed_token = sign_token(&state.signing_key, &claims)
        .map_err(|_| {
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(ErrorResponse {
                    error: "server_error".to_string(),
                    error_description: Some("Failed to sign token".to_string()),
                }),
            )
        })?;

    // Store the access token with refresh token
    let token_data = AccessToken::with_refresh_token(
        signed_token.clone(),
        user.profile.me.clone(),
        code_data.scope.clone(),
        code_data.client_id.clone(),
        std::time::Duration::from_secs(expires_in as u64),
        refresh_token.clone(),
        std::time::Duration::from_secs(refresh_token_expires_in as u64),
    );
    state.store_token(signed_token.clone(), token_data);
    
    // Store refresh token -> access token mapping
    state.store_refresh_token(refresh_token.clone(), signed_token.clone());

    // Build profile data based on scopes
    let profile = build_profile_response(&user.profile, &code_data.scope);

    // Build response
    let response = TokenResponse {
        access_token: signed_token,
        token_type: "Bearer".to_string(),
        me: user.profile.me.to_string(),
        scope: code_data.scope.to_string(),
        profile,
        expires_in,
        refresh_token: Some(refresh_token),
        refresh_token_expires_in: Some(refresh_token_expires_in),
    };

    Ok(Json(json!(response)))
}

/// Handle refresh token grant type
async fn handle_refresh_token_grant(
    state: &Arc<AppState>,
    request: &TokenRequest,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
    // Validate refresh token
    let access_token = match state.get_access_token_by_refresh_token(&request.refresh_token) {
        Some(token) => token,
        None => {
            return Err((
                StatusCode::BAD_REQUEST,
                Json(ErrorResponse {
                    error: "invalid_grant".to_string(),
                    error_description: Some("Invalid or expired refresh token".to_string()),
                }),
            ));
        }
    };

    // Verify the refresh token hasn't expired
    let token_data = match state.get_token(&access_token) {
        Some(data) => data,
        None => {
            return Err((
                StatusCode::BAD_REQUEST,
                Json(ErrorResponse {
                    error: "invalid_grant".to_string(),
                    error_description: Some("Invalid refresh token".to_string()),
                }),
            ));
        }
    };

    if token_data.is_expired() {
        return Err((
            StatusCode::BAD_REQUEST,
            Json(ErrorResponse {
                error: "invalid_grant".to_string(),
                error_description: Some("Refresh token has expired".to_string()),
            }),
        ));
    }

    // Generate new tokens
    let new_refresh_token = generate_refresh_token();
    
    // Clone values for different use cases
    let me_for_claims = token_data.me.clone();
    let scope_for_claims = token_data.scope.clone();
    let client_id_for_claims = token_data.client_id.clone();
    
    let me_for_token = token_data.me.clone();
    let scope_for_token = token_data.scope.clone();
    let client_id_for_token = token_data.client_id.clone();
    
    let me_for_response = token_data.me.clone();
    let scope_for_response = token_data.scope.clone();
    
    // Create new claims for the refreshed token
    let new_claims = TokenClaims {
        me: me_for_claims,
        scope: scope_for_claims,
        client_id: client_id_for_claims,
        issued_at: std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs(),
        expires_at: std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs() + 3600, // 1 hour
    };
    
    let new_access_token = sign_token(&state.signing_key, &new_claims).unwrap();
    let expires_in = 3600i64; // 1 hour
    let refresh_token_expires_in = 86400i64 * 30; // 30 days

    // Store new access token
    let new_token_data = AccessToken::with_refresh_token(
        new_access_token.clone(),
        me_for_token,
        scope_for_token,
        client_id_for_token,
        std::time::Duration::from_secs(expires_in as u64),
        new_refresh_token.clone(),
        std::time::Duration::from_secs(refresh_token_expires_in as u64),
    );
    state.store_token(new_access_token.clone(), new_token_data);

    // Update refresh token mapping
    state.store_refresh_token(new_refresh_token.clone(), new_access_token.clone());

    // Build response
    let response = TokenResponse {
        access_token: new_access_token,
        token_type: "Bearer".to_string(),
        me: me_for_response.to_string(),
        scope: scope_for_response.to_string(),
        profile: None,
        expires_in,
        refresh_token: Some(new_refresh_token),
        refresh_token_expires_in: Some(refresh_token_expires_in),
    };

    Ok(Json(json!(response)))
}

/// Verify PKCE code_challenge and code_verifier
fn verify_pkce(
    code_challenge: &str,
    code_challenge_method: &str,
    code_verifier: &str,
) -> bool {
    if code_challenge_method != "S256" {
        // Only S256 is supported
        return false;
    }

    // Generate code challenge from code verifier
    let mut hasher = Sha256::new();
    hasher.update(code_verifier.as_bytes());
    let hash = hasher.finalize();
    let computed_challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(hash);

    // Compare with stored challenge
    computed_challenge == code_challenge
}

/// Generate a new refresh token
fn generate_refresh_token() -> String {
    nanoid!(32)
}

/// Sign a JWT token with the provided key
fn sign_token(signing_key: &[u8], claims: &TokenClaims) -> Result<String, String> {
    // This is a simplified token signing implementation
    // In a production environment, you'd use a proper JWT library
    
    let header = json!({
        "typ": "JWT",
        "alg": "HS256"
    });

    let header_str = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(serde_json::to_string(&header).unwrap());
    let claims_str = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(serde_json::to_string(claims).unwrap());
    
    let signature_input = format!("{}.{}", header_str, claims_str);
    let mut hasher = Sha256::new();
    hasher.update(signature_input.as_bytes());
    hasher.update(signing_key);
    let signature = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(hasher.finalize());
    
    Ok(format!("{}.{}.{}", header_str, claims_str, signature))
}

/// Build profile response based on scopes
fn build_profile_response(
    user_profile: &crate::state::Profile,
    scope: &indieweb::standards::indieauth::Scopes,
) -> Option<serde_json::Value> {
    let mut profile = serde_json::Map::new();
    
    // Always include 'me' field
    profile.insert("me".to_string(), serde_json::Value::String(user_profile.me.to_string()));
    
    // Include profile information based on scopes
    if scope.has("profile") {
        let name = user_profile.name.clone();
            profile.insert("name".to_string(), serde_json::Value::String(name));
        
        if let Some(photo) = &user_profile.photo {
            profile.insert("photo".to_string(), serde_json::Value::String(photo.to_string()));
        }
        
        if let Some(url) = &user_profile.url {
            profile.insert("url".to_string(), serde_json::Value::String(url.to_string()));
        }
    }
    
    if scope.has("email") {
        if let Some(email) = &user_profile.email {
            profile.insert("email".to_string(), serde_json::Value::String(email.to_string()));
        }
    }
    
    Some(serde_json::Value::Object(profile))
}

/// POST handler for token revocation
pub async fn revoke_token_handler(
    State(state): State<Arc<AppState>>,
    axum::Form(form): axum::Form<serde_json::Value>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
    // Extract token from form data
    let token = form.get("token")
        .and_then(|v| v.as_str())
        .ok_or_else(|| {
            (
                StatusCode::BAD_REQUEST,
                Json(ErrorResponse {
                    error: "invalid_request".to_string(),
                    error_description: Some("Missing token parameter".to_string()),
                }),
            )
        })?;

    // Remove the token
    state.revoke_token(token);
    
    // Return empty 200 OK response as per OAuth 2.0 spec
    Ok(Json(json!({})))
}

