use rocket::{
    Request,
    http::Status,
    outcome::{Outcome, try_outcome},
    request::{self, FromRequest},
    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::session::Session, db::DbConn};

// 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!())
}

// api

#[post("/totp", data = "<totp>")]
pub async fn confirm_totp(
    mfa: TOTPSecret,
    totp: Json<TOTPSixDigitCode>,
    mut db: Connection<DbConn>,
) -> Status {
    if totp.code.len() == 6
        && let Ok(code) = totp.code.parse::<usize>()
    {
        let totp = totp_gen(mfa.user_id, mfa.secret.as_bytes()).unwrap();

        println!("input: {} {}", code, totp.generate_current().unwrap());

        if totp.check_current(&format!("{}", mfa.user_id)).unwrap() {
            if sqlx::query!(
                "UPDATE users SET twofa_enabled = true WHERE id = $1",
                mfa.user_id as i32
            )
            .execute(&mut **db)
            .await
            .is_err()
            {
                return Status::InternalServerError;
            };
        }
    }

    println!("ok!");

    return Status::Ok;
}

#[get("/totp.jpg")]
pub async fn get_totp(mfa: TOTPSecret) -> 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,
}

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 user = try_outcome!(request.guard::<Session>().await);
        let mut pool = match request.guard::<Connection<DbConn>>().await {
            Outcome::Success(pool) => pool,
            _ => return Outcome::Error((Status::Unauthorized, ())),
        };

        let (enabled, mut secret) = match sqlx::query!(
            "SELECT twofa_enabled, totp_secret FROM users WHERE id = $1",
            user.user_id as i32,
        )
        .fetch_one(&mut **pool)
        .await
        {
            Ok(row) => (row.twofa_enabled, row.totp_secret),
            Err(_) => return Outcome::Error((Status::Unauthorized, ())),
        };

        if !enabled || secret.is_none() {
            secret = Some(Secret::generate_secret().to_string());

            match sqlx::query!(
                "UPDATE users SET totp_secret = $1 WHERE id = $2",
                secret.as_ref().unwrap(),
                user.user_id as i32,
            )
            .execute(&mut **pool)
            .await
            {
                Ok(_) => (),
                Err(_) => return Outcome::Error((Status::InternalServerError, ())),
            }
        }

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

impl TOTPSecret {
    pub async fn enable(&self, db: &mut Connection<DbConn>) -> 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(()),
        }
    }
}
