use futures_util::TryFutureExt;
use rocket::{
    Request,
    http::Status,
    outcome::{Outcome, try_outcome},
    request::{self, FromRequest},
    response::status::{self},
    serde::json::Json,
};
use rocket_db_pools::Connection;
use rocket_dyn_templates::{Template, context};
use serde::{Deserialize, Serialize};
use totp_rs::{Algorithm, Secret, TOTP};

use crate::{
    auth::{
        account::AuthResponse,
        profile::PasswordAnd2fa,
        session::{Claims, Session, TokenScope},
    },
    db::Postgres,
    user::User,
};

// Utility methods

pub fn totp_gen(user_id: usize, secret: &[u8]) -> Result<TOTP, String> {
    TOTP::new(
        Algorithm::SHA1,
        6,
        1,
        30,
        secret.to_owned(),
        Some("chat.zxq5.dev".to_string()),
        format!("{}", user_id),
    )
    .map_err(|_| String::from("Invalid Secret"))
}

// pages

#[get("/totp")]
pub async fn mfa_page(_session: Session) -> Template {
    Template::render("2fa", context!())
}

#[post("/totp", data = "<form>")]
pub async fn confirm_totp(
    mfa: TOTPSecret,
    form: Json<TOTPSixDigitCode>,
    mut db: Connection<Postgres>,
) -> Result<(), status::Custom<&'static str>> {
    if form.code.len() != 6 || form.code.parse::<u32>().is_err() {
        return Err(status::Custom(Status::BadRequest, "Invalid 6-digit code"));
    }

    println!("valid");

    let totp = totp_gen(mfa.user_id, mfa.secret.as_bytes())
        .map_err(|_| status::Custom(Status::InternalServerError, "TOTP Error"))?;
    if !totp.check_current(&form.code).unwrap_or(false) {
        return Err(status::Custom(Status::BadRequest, "Incorrect code"));
    }
    println!("correct");

    if sqlx::query!(
        "UPDATE users SET twofa_enabled = true WHERE id = $1",
        mfa.user_id as i32
    )
    .execute(&mut **db)
    .await
    .is_err()
    {
        return Err(status::Custom(
            Status::InternalServerError,
            "unable to enable 2fa",
        ));
    };

    println!("enabled");

    Ok(())
}

#[derive(Deserialize)]
pub struct PasswordConfirmation {
    password: String,
}

#[post("/totp.jpg", data = "<form>")]
pub async fn get_totp(
    mfa: TOTPSecret,
    form: Json<PasswordConfirmation>,
) -> Option<Json<QrResponse>> {
    let qr_b64 = totp_gen(mfa.user_id, mfa.secret.as_bytes())
        .expect("Invalid TOTP")
        .get_qr_base64()
        .unwrap();

    Some(Json(QrResponse {
        qr_code: format!("data:image/png;base64,{}", qr_b64),
    }))
}

#[derive(Debug, Deserialize)]
pub struct TOTPSixDigitCode {
    code: String,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TotpStatus {
    Enabled,
    Disabled,
}

pub struct TOTPSecret {
    user_id: usize,
    secret: String,
}

#[derive(Serialize)]
pub struct QrResponse {
    qr_code: String,
}

#[rocket::async_trait]
impl<'r> FromRequest<'r> for TOTPSecret {
    type Error = ();

    async fn from_request(request: &'r Request<'_>) -> request::Outcome<Self, Self::Error> {
        let auth_header = request.headers().get_one("Authorization");
        println!(
            "TOTPSecret guard - Auth header present: {}",
            auth_header.is_some()
        );

        let user = try_outcome!(request.guard::<Claims>().await);
        println!(
            "TOTPSecret guard - Claims ok, user: {}, scope: {:?}",
            user.sub, user.scope
        );

        // only allow full tokens for TOTP setup
        if user.scope != TokenScope::Full {
            println!("TOTPSecret guard - rejected, scope is {:?}", user.scope);
            return Outcome::Error((Status::Forbidden, ()));
        }

        let user = try_outcome!(request.guard::<Session>().await);
        let mut pool = match request.guard::<Connection<Postgres>>().await {
            Outcome::Success(pool) => pool,
            _ => return Outcome::Error((Status::Unauthorized, ())),
        };

        let row = sqlx::query!(
            "SELECT twofa_enabled, totp_secret FROM users WHERE id = $1",
            user.user_id as i32
        )
        .fetch_one(&mut **pool)
        .await;

        let (enabled, mut secret) = match row {
            Ok(r) => (r.twofa_enabled, r.totp_secret),
            Err(_) => return Outcome::Error((Status::Unauthorized, ())),
        };

        if secret.is_none() {
            let new_secret = Secret::generate_secret().to_encoded().to_string();
            sqlx::query!(
                "UPDATE users SET totp_secret = $1 WHERE id = $2",
                new_secret,
                user.user_id as i32
            )
            .execute(&mut **pool)
            .await
            .ok();
            secret = Some(new_secret);
        }

        Outcome::Success(TOTPSecret {
            user_id: user.user_id,
            secret: secret.unwrap(),
        })
    }
}

impl TOTPSecret {
    pub async fn enable(&self, db: &mut Connection<Postgres>) -> Result<(), ()> {
        match sqlx::query!(
            "UPDATE users SET twofa_enabled = true WHERE id = $1",
            self.user_id as i32,
        )
        .execute(&mut ***db)
        .await
        {
            Ok(_) => Ok(()),
            Err(_) => Err(()),
        }
    }
}

#[derive(Deserialize)]
pub struct TotpVerifyRequest {
    pub code: String,
}

#[get("/totp/status")]
pub async fn get_totp_status(
    user: Session,
    mut db: Connection<Postgres>,
) -> Result<Json<TotpStatus>, Status> {
    Ok(Json(
        if sqlx::query!(
            "SELECT twofa_enabled FROM users WHERE id = $1",
            user.user_id as i32,
        )
        .fetch_one(&mut **db)
        .await
        .map_err(|_| Status::NotFound)?
        .twofa_enabled
        {
            TotpStatus::Enabled
        } else {
            TotpStatus::Disabled
        },
    ))
}

#[delete("/totp", data = "<form>")]
pub async fn disable_totp(
    user: Session,
    mut db: Connection<Postgres>,
    form: Json<PasswordAnd2fa>,
) -> Result<Json<AuthResponse>, Status> {
    let totp_code = form.totp_code.clone().ok_or(Status::BadRequest)?;
    let mut user = User::get_by_id(user.user_id, &mut db)
        .await
        .ok_or(Status::NotFound)?;

    user.verify_password(&form.password)?;
    user.verify_2fa(&totp_code)?;
    user.set_twofa_enabled(false, &mut db)
        .await
        .map_err(|_| Status::InternalServerError)?;

    Ok(Json(AuthResponse {
        token: Claims::new(user.id as usize, TokenScope::Full).encode(),
        totp_required: false,
    }))
}

#[post("/totp/verify", data = "<body>")]
pub async fn verify_totp(
    claims: Claims, // request guard checks token validity
    mut db: Connection<Postgres>,
    body: Json<TotpVerifyRequest>,
) -> Result<Json<AuthResponse>, Status> {
    println!("reached 1");

    // reject if they somehow got here with a full token
    if claims.scope != TokenScope::TotpPending {
        return Err(Status::Forbidden);
    }

    println!("reached 2");

    let row = sqlx::query!(
        "SELECT totp_secret FROM users WHERE id = $1 AND twofa_enabled = TRUE",
        claims.sub
    )
    .fetch_one(&mut **db)
    .await
    .map_err(|_| Status::Unauthorized)?;

    println!("reached 3");

    let totp = totp_gen(
        claims.sub as usize,
        row.totp_secret
            .expect("user with 2fa enabled has no totp secret")
            .as_bytes(),
    )
    .map_err(|_| Status::InternalServerError)?;

    if !totp
        .check_current(&body.code)
        .map_err(|_| Status::InternalServerError)?
    {
        return Err(Status::Unauthorized);
    }

    println!("reached 5");

    let claims = Claims::new(claims.sub as usize, TokenScope::Full);

    Ok(Json(AuthResponse {
        token: claims.encode(),
        totp_required: false,
    }))
}
