#[macro_use]
extern crate rocket;
extern crate bls12_381;
extern crate rand;

use bls12_381::pairing;
use bls12_381::G2Affine;
use rocket::serde::json::Json;
use rocket::Request;
use rocket_okapi::settings::UrlObject;
use rocket_okapi::{openapi, openapi_get_routes, rapidoc::*, swagger_ui::*};
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::RwLock;

use ophe::core;
use ophe::core::{
    RatelimiterRequest, RatelimiterResponse, RatelimiterRotateRequest, RatelimiterRotateResponse,
    SetKeyHelper,
};

use bls12_381::G1Affine;
use bls12_381::Scalar;
use ophe::utils;
use rocket::State;

fn make_public_parameters(key: Scalar) -> CryptoserviceState {
    let generator = pairing(&G1Affine::generator(), &G2Affine::generator());
    let public_key = generator * key;

    let pp = core::PublicParameters {
        ratelimiter_public_key: public_key,
        gt_gen: generator,
    };
    CryptoserviceState { key, pp }
}

struct CryptoserviceState {
    key: Scalar,
    pp: core::PublicParameters,
}

type CryptoserviceStatePointer = Arc<RwLock<CryptoserviceState>>;

#[rocket::main]
async fn main() {
    let c = make_public_parameters(utils::random_scalar());
    let c_p: CryptoserviceStatePointer = Arc::new(RwLock::new(c));
    let figment = rocket::Config::figment().merge(("port", 8081));

    let launch_result = rocket::custom(figment)
        .manage(c_p)
        .mount(
            "/",
            openapi_get_routes![phe_help, get_public_parameters, set_key, rotate_key],
        )
        .mount(
            "/swagger-ui/",
            make_swagger_ui(&SwaggerUIConfig {
                url: "../openapi.json".to_owned(),
                ..Default::default()
            }),
        )
        .mount(
            "/rapidoc/",
            make_rapidoc(&RapiDocConfig {
                general: GeneralConfig {
                    spec_urls: vec![UrlObject::new("General", "../openapi.json")],
                    ..Default::default()
                },
                hide_show: HideShowConfig {
                    allow_spec_url_load: false,
                    allow_spec_file_load: false,
                    ..Default::default()
                },
                ..Default::default()
            }),
        )
        .register("/", catchers![serialize_failed])
        .launch()
        .await;
    match launch_result {
        Ok(_) => println!("Rocket shut down gracefully."),
        Err(err) => println!("Rocket had an error: {}", err),
    };
}

#[openapi()]
#[post("/phe_help", format = "json", data = "<request>")]
fn phe_help(
    request: Json<RatelimiterRequest>,
    c_state: &State<CryptoserviceStatePointer>,
) -> Json<RatelimiterResponse> {
    let (key, pp) = {
        let c_state = c_state.read().unwrap();
        (c_state.key.clone(), c_state.pp.clone())
    };

    Json(core::phe_ratelimiter(&key, &request, &pp))
}

#[openapi()]
#[get("/get_public_parameters")]
fn get_public_parameters(
    c_state: &State<CryptoserviceStatePointer>,
) -> Json<core::PublicParameters> {
    let pp = c_state.read().unwrap().pp.clone();
    Json(pp)
}

#[openapi()]
#[post("/set_key", format = "json", data = "<request>")]
fn set_key(
    request: Json<SetKeyHelper>,
    c_state: &State<CryptoserviceStatePointer>,
) -> Json<core::PublicParameters> {
    let new_c = make_public_parameters(request.key);

    let mut c_state = c_state.write().unwrap();
    c_state.key = new_c.key;
    c_state.pp = new_c.pp;

    Json(new_c.pp)
}

#[openapi()]
#[post("/rotate_key", format = "json", data = "<request>")]
fn rotate_key(
    request: Json<RatelimiterRotateRequest>,
    c_state: &State<CryptoserviceStatePointer>,
) -> Json<RatelimiterRotateResponse> {
    let mut c_state = c_state.write().unwrap();
    let (response, k_new) = core::phe_roate_ratelimiter(&c_state.key, &request.s);
    c_state.key = k_new;

    Json(response)
}

#[catch(422)]
fn serialize_failed(_req: &Request) -> String {
    format!("Malformed Request")
}