use crate::db;
use crate::db::users::User;
use crate::errors::CognitoError;
use crate::server::AppState;
use axum::response::{IntoResponse, Response};
use std::collections::HashMap;
use std::sync::Arc;

fn issue_tokens(
    state: &Arc<AppState>,
    pool_id: &str,
    user: &User,
    client_id: &str,
) -> Result<serde_json::Value, CognitoError> {
    let groups = db::groups::list_groups_for_user(&state.db, pool_id, &user.username)
        .map_err(|e| CognitoError::InternalError(e.to_string()))?;
    let group_names: Vec<String> = groups.iter().map(|g| g.group_name.clone()).collect();

    let issuer = if state.config.issuer_base.is_empty() {
        format!("http://localhost:{}/{}", state.config.port, pool_id)
    } else {
        format!("{}/{}", state.config.issuer_base, pool_id)
    };

    let access_token = state
        .crypto
        .sign_access_token(
            &user.sub,
            &user.username,
            pool_id,
            client_id,
            &issuer,
            group_names.clone(),
        )
        .map_err(|e| CognitoError::InternalError(e.to_string()))?;

    let mut attributes: HashMap<String, serde_json::Value> = HashMap::new();
    for attr in &user.attributes {
        attributes.insert(attr.name.clone(), serde_json::Value::String(attr.value.clone()));
    }

    let id_token = state
        .crypto
        .sign_id_token(
            &user.sub,
            &user.username,
            pool_id,
            client_id,
            &issuer,
            group_names,
            attributes,
        )
        .map_err(|e| CognitoError::InternalError(e.to_string()))?;

    let refresh_token = db::tokens::create_refresh_token(&state.db, pool_id, &user.username)
        .map_err(|e| CognitoError::InternalError(e.to_string()))?;

    Ok(serde_json::json!({
        "AccessToken": access_token,
        "IdToken": id_token,
        "RefreshToken": refresh_token,
        "ExpiresIn": 3600,
        "TokenType": "Bearer",
    }))
}

/// Resolve pool_id from the request body: prefer explicit UserPoolId,
/// fall back to the first pool in the database.
fn resolve_pool_id(
    body: &serde_json::Value,
    state: &Arc<AppState>,
) -> Result<String, CognitoError> {
    if let Some(pid) = body["UserPoolId"].as_str() {
        return Ok(pid.to_string());
    }
    let pools = db::pools::list(&state.db, 1)
        .map_err(|e| CognitoError::InternalError(e.to_string()))?;
    pools
        .first()
        .map(|p| p.id.clone())
        .ok_or_else(|| CognitoError::ResourceNotFound("No user pools exist".into()))
}

pub async fn sign_up(
    body: serde_json::Value,
    state: &Arc<AppState>,
) -> Result<Response, CognitoError> {
    let client_id = body["ClientId"]
        .as_str()
        .ok_or_else(|| CognitoError::InvalidParameter("ClientId is required".into()))?;
    let _ = client_id; // ClientId is required but not stored in this implementation

    let username = body["Username"]
        .as_str()
        .ok_or_else(|| CognitoError::InvalidParameter("Username is required".into()))?;

    let password = body["Password"]
        .as_str()
        .ok_or_else(|| CognitoError::InvalidParameter("Password is required".into()))?;

    let pool_id = resolve_pool_id(&body, state)?;

    let password_hash = db::users::hash_password(password)
        .map_err(|e| CognitoError::InternalError(e.to_string()))?;

    let user = db::users::create(
        &state.db,
        &pool_id,
        username,
        Some(&password_hash),
        "UNCONFIRMED",
        &[],
    )
    .map_err(|e| {
        if e.to_string().contains("UNIQUE constraint failed")
            || e.to_string().contains("PRIMARY KEY")
        {
            CognitoError::UsernameExists(format!("User account already exists for {username}"))
        } else {
            CognitoError::InternalError(e.to_string())
        }
    })?;

    tracing::info!(
        username = username,
        "user signed up — confirmation code: 123456"
    );

    let response = serde_json::json!({
        "UserConfirmed": false,
        "UserSub": user.sub,
    });
    Ok(axum::Json(response).into_response())
}

pub async fn confirm_sign_up(
    body: serde_json::Value,
    state: &Arc<AppState>,
) -> Result<Response, CognitoError> {
    let username = body["Username"]
        .as_str()
        .ok_or_else(|| CognitoError::InvalidParameter("Username is required".into()))?;

    let code = body["ConfirmationCode"]
        .as_str()
        .ok_or_else(|| CognitoError::InvalidParameter("ConfirmationCode is required".into()))?;

    // Accept any 6-digit numeric code
    if code.len() != 6 || !code.chars().all(|c| c.is_ascii_digit()) {
        return Err(CognitoError::CodeMismatch(
            "Invalid confirmation code".into(),
        ));
    }

    let pool_id = resolve_pool_id(&body, state)?;

    let user = db::users::get(&state.db, &pool_id, username)
        .map_err(|e| CognitoError::InternalError(e.to_string()))?
        .ok_or_else(|| {
            CognitoError::NotAuthorized(format!("User {username} not found"))
        })?;

    if user.status != "UNCONFIRMED" {
        return Err(CognitoError::NotAuthorized(format!(
            "User {username} is not in UNCONFIRMED state"
        )));
    }

    db::users::update_status(&state.db, &pool_id, username, "CONFIRMED")
        .map_err(|e| CognitoError::InternalError(e.to_string()))?;

    Ok(axum::Json(serde_json::json!({})).into_response())
}

pub async fn initiate_auth(
    body: serde_json::Value,
    state: &Arc<AppState>,
) -> Result<Response, CognitoError> {
    let auth_flow = body["AuthFlow"]
        .as_str()
        .ok_or_else(|| CognitoError::InvalidParameter("AuthFlow is required".into()))?;

    let client_id = body["ClientId"]
        .as_str()
        .ok_or_else(|| CognitoError::InvalidParameter("ClientId is required".into()))?;

    match auth_flow {
        "USER_PASSWORD_AUTH" => {
            let auth_params = &body["AuthParameters"];
            let username = auth_params["USERNAME"]
                .as_str()
                .ok_or_else(|| {
                    CognitoError::InvalidParameter("USERNAME is required in AuthParameters".into())
                })?;
            let password = auth_params["PASSWORD"]
                .as_str()
                .ok_or_else(|| {
                    CognitoError::InvalidParameter("PASSWORD is required in AuthParameters".into())
                })?;

            let pool_id = resolve_pool_id(&body, state)?;

            let user = db::users::get(&state.db, &pool_id, username)
                .map_err(|e| CognitoError::InternalError(e.to_string()))?
                .ok_or_else(|| {
                    CognitoError::NotAuthorized("Incorrect username or password.".into())
                })?;

            if !user.enabled {
                return Err(CognitoError::NotAuthorized("User is disabled.".into()));
            }

            let pw_hash = user.password_hash.as_deref().ok_or_else(|| {
                CognitoError::NotAuthorized("User has no password set.".into())
            })?;

            let valid = db::users::verify_password(password, pw_hash)
                .map_err(|e| CognitoError::InternalError(e.to_string()))?;
            if !valid {
                return Err(CognitoError::NotAuthorized(
                    "Incorrect username or password.".into(),
                ));
            }

            if user.status == "FORCE_CHANGE_PASSWORD" {
                let response = serde_json::json!({
                    "ChallengeName": "NEW_PASSWORD_REQUIRED",
                    "ChallengeParameters": {
                        "USERNAME": username,
                        "USER_ID_FOR_SRP": username,
                    },
                    "Session": uuid::Uuid::new_v4().to_string(),
                });
                return Ok(axum::Json(response).into_response());
            }

            if user.status != "CONFIRMED" {
                return Err(CognitoError::NotAuthorized(format!(
                    "User {username} is not confirmed."
                )));
            }

            let tokens = issue_tokens(state, &pool_id, &user, client_id)?;
            let response = serde_json::json!({
                "AuthenticationResult": tokens,
            });
            Ok(axum::Json(response).into_response())
        }
        "REFRESH_TOKEN_AUTH" | "REFRESH_TOKEN" => {
            let auth_params = &body["AuthParameters"];
            let refresh_token = auth_params["REFRESH_TOKEN"]
                .as_str()
                .ok_or_else(|| {
                    CognitoError::InvalidParameter(
                        "REFRESH_TOKEN is required in AuthParameters".into(),
                    )
                })?;

            let token_info = db::tokens::validate_refresh_token(&state.db, refresh_token)
                .map_err(|e| CognitoError::InternalError(e.to_string()))?
                .ok_or_else(|| {
                    CognitoError::NotAuthorized("Invalid or expired refresh token.".into())
                })?;

            let user = db::users::get(&state.db, &token_info.pool_id, &token_info.username)
                .map_err(|e| CognitoError::InternalError(e.to_string()))?
                .ok_or_else(|| {
                    CognitoError::NotAuthorized("User no longer exists.".into())
                })?;

            let tokens = issue_tokens(state, &token_info.pool_id, &user, client_id)?;
            let response = serde_json::json!({
                "AuthenticationResult": tokens,
            });
            Ok(axum::Json(response).into_response())
        }
        _ => Err(CognitoError::InvalidParameter(format!(
            "Unsupported auth flow: {auth_flow}"
        ))),
    }
}

pub async fn admin_initiate_auth(
    body: serde_json::Value,
    state: &Arc<AppState>,
) -> Result<Response, CognitoError> {
    let auth_flow = body["AuthFlow"]
        .as_str()
        .ok_or_else(|| CognitoError::InvalidParameter("AuthFlow is required".into()))?;

    let pool_id = body["UserPoolId"]
        .as_str()
        .ok_or_else(|| CognitoError::InvalidParameter("UserPoolId is required".into()))?;

    let client_id = body["ClientId"]
        .as_str()
        .ok_or_else(|| CognitoError::InvalidParameter("ClientId is required".into()))?;

    match auth_flow {
        "ADMIN_NO_SRP_AUTH" | "ADMIN_USER_PASSWORD_AUTH" => {
            let auth_params = &body["AuthParameters"];
            let username = auth_params["USERNAME"]
                .as_str()
                .ok_or_else(|| {
                    CognitoError::InvalidParameter("USERNAME is required in AuthParameters".into())
                })?;
            let password = auth_params["PASSWORD"]
                .as_str()
                .ok_or_else(|| {
                    CognitoError::InvalidParameter("PASSWORD is required in AuthParameters".into())
                })?;

            let user = db::users::get(&state.db, pool_id, username)
                .map_err(|e| CognitoError::InternalError(e.to_string()))?
                .ok_or_else(|| {
                    CognitoError::NotAuthorized("Incorrect username or password.".into())
                })?;

            if !user.enabled {
                return Err(CognitoError::NotAuthorized("User is disabled.".into()));
            }

            let pw_hash = user.password_hash.as_deref().ok_or_else(|| {
                CognitoError::NotAuthorized("User has no password set.".into())
            })?;

            let valid = db::users::verify_password(password, pw_hash)
                .map_err(|e| CognitoError::InternalError(e.to_string()))?;
            if !valid {
                return Err(CognitoError::NotAuthorized(
                    "Incorrect username or password.".into(),
                ));
            }

            if user.status == "FORCE_CHANGE_PASSWORD" {
                let response = serde_json::json!({
                    "ChallengeName": "NEW_PASSWORD_REQUIRED",
                    "ChallengeParameters": {
                        "USERNAME": username,
                        "USER_ID_FOR_SRP": username,
                    },
                    "Session": uuid::Uuid::new_v4().to_string(),
                });
                return Ok(axum::Json(response).into_response());
            }

            if user.status != "CONFIRMED" {
                return Err(CognitoError::NotAuthorized(format!(
                    "User {username} is not confirmed."
                )));
            }

            let tokens = issue_tokens(state, pool_id, &user, client_id)?;
            let response = serde_json::json!({
                "AuthenticationResult": tokens,
            });
            Ok(axum::Json(response).into_response())
        }
        _ => Err(CognitoError::InvalidParameter(format!(
            "Unsupported auth flow: {auth_flow}"
        ))),
    }
}

pub async fn respond_to_auth_challenge(
    body: serde_json::Value,
    state: &Arc<AppState>,
) -> Result<Response, CognitoError> {
    let challenge_name = body["ChallengeName"]
        .as_str()
        .ok_or_else(|| CognitoError::InvalidParameter("ChallengeName is required".into()))?;

    let client_id = body["ClientId"]
        .as_str()
        .ok_or_else(|| CognitoError::InvalidParameter("ClientId is required".into()))?;

    match challenge_name {
        "NEW_PASSWORD_REQUIRED" => {
            let challenge_responses = &body["ChallengeResponses"];
            let username = challenge_responses["USERNAME"]
                .as_str()
                .ok_or_else(|| {
                    CognitoError::InvalidParameter(
                        "USERNAME is required in ChallengeResponses".into(),
                    )
                })?;
            let new_password = challenge_responses["NEW_PASSWORD"]
                .as_str()
                .ok_or_else(|| {
                    CognitoError::InvalidParameter(
                        "NEW_PASSWORD is required in ChallengeResponses".into(),
                    )
                })?;

            let pool_id = resolve_pool_id(&body, state)?;

            let new_hash = db::users::hash_password(new_password)
                .map_err(|e| CognitoError::InternalError(e.to_string()))?;

            db::users::set_password(&state.db, &pool_id, username, &new_hash, true)
                .map_err(|e| CognitoError::InternalError(e.to_string()))?;

            let user = db::users::get(&state.db, &pool_id, username)
                .map_err(|e| CognitoError::InternalError(e.to_string()))?
                .ok_or_else(|| {
                    CognitoError::NotAuthorized(format!("User {username} not found"))
                })?;

            let tokens = issue_tokens(state, &pool_id, &user, client_id)?;
            let response = serde_json::json!({
                "AuthenticationResult": tokens,
            });
            Ok(axum::Json(response).into_response())
        }
        _ => Err(CognitoError::InvalidParameter(format!(
            "Unsupported challenge: {challenge_name}"
        ))),
    }
}
