use axum::{
    extract::State,
    http::StatusCode,
    response::Json,
};
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
use hmac::{Hmac, Mac};
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,
    pub code: String,
    pub client_id: String,
    pub redirect_uri: String,
    pub code_verifier: 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,
}

/// 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
/// Validates authorization codes and exchanges them for access tokens
pub async fn token_handler(
    State(state): State<Arc<AppState>>,
    axum::Form(request): axum::Form<TokenRequest>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
    // Validate grant type
    if request.grant_type != "authorization_code" {
        return Err((
            StatusCode::BAD_REQUEST,
            Json(ErrorResponse {
                error: "unsupported_grant_type".to_string(),
                error_description: Some("Only authorization_code grant type is supported".to_string()),
            }),
        ));
    }

    // 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 access token
    let access_token = generate_access_token();
    let expires_in = 3600i64; // 1 hour

    // 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
    let token_data = AccessToken::new(
        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),
    );
    state.store_token(signed_token.clone(), token_data);

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

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

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

/// Verify PKCE code challenge
fn verify_pkce(code_challenge: &str, code_challenge_method: &str, code_verifier: &str) -> bool {
    match code_challenge_method {
        "S256" => {
            let mut hasher = Sha256::new();
            hasher.update(code_verifier.as_bytes());
            let hash = hasher.finalize();
            let computed_challenge = URL_SAFE_NO_PAD.encode(&hash);
            
            // Constant-time comparison to prevent timing attacks
            use subtle::ConstantTimeEq;
            computed_challenge.as_bytes().ct_eq(code_challenge.as_bytes()).into()
        }
        "plain" => {
            code_challenge == code_verifier
        }
        _ => false,
    }
}

/// Generate a secure access token
fn generate_access_token() -> String {
    format!("token_{}", nanoid!(32))
}

/// Sign a token using HMAC-SHA256
fn sign_token(key: &[u8; 32], claims: &TokenClaims) -> Result<String, Box<dyn std::error::Error>> {
    // Create a simple JWT-like token structure
    let header = json!({
        "alg": "HS256",
        "typ": "JWT"
    });
    
    let payload = json!({
        "me": claims.me.to_string(),
        "scope": claims.scope.to_string(),
        "client_id": claims.client_id.to_string(),
        "iat": claims.issued_at,
        "exp": claims.expires_at,
    });

    // Encode header and payload
    let header_encoded = URL_SAFE_NO_PAD.encode(serde_json::to_string(&header)?.as_bytes());
    let payload_encoded = URL_SAFE_NO_PAD.encode(serde_json::to_string(&payload)?.as_bytes());
    
    // Create signature
    let message = format!("{}.{}", header_encoded, payload_encoded);
    let mut mac = Hmac::<Sha256>::new_from_slice(key).map_err(|_| "Invalid key")?;
    mac.update(message.as_bytes());
    let signature = mac.finalize().into_bytes();
    let signature_encoded = URL_SAFE_NO_PAD.encode(&signature);
    
    Ok(format!("{}.{}", message, signature_encoded))
}

/// Build profile response based on granted scopes
fn build_profile_response(profile: &crate::state::Profile, scopes: &indieweb::standards::indieauth::Scopes) -> Option<serde_json::Value> {
    if scopes.is_empty() {
        return None;
    }

    let mut profile_data = serde_json::Map::new();
    
    // Always include basic profile info if any scope is granted
    profile_data.insert("name".to_string(), json!(profile.name));
    
    if let Some(photo) = &profile.photo {
        profile_data.insert("photo".to_string(), json!(photo.to_string()));
    }
    
    if let Some(url) = &profile.url {
        profile_data.insert("url".to_string(), json!(url.to_string()));
    } else {
        profile_data.insert("url".to_string(), json!(profile.me.to_string()));
    }

    // Include email if email scope is granted
    if scopes.has("email") {
        if let Some(email) = &profile.email {
            profile_data.insert("email".to_string(), json!(email));
        }
    }

    Some(json!(profile_data))
}