// trifid-api, an open source reimplementation of the Defined Networking nebula management server.
//     Copyright (C) 2023 c0repwn3r
//
//     This program is free software: you can redistribute it and/or modify
//     it under the terms of the GNU General Public License as published by
//     the Free Software Foundation, either version 3 of the License, or
//     (at your option) any later version.
//
//     This program is distributed in the hope that it will be useful,
//     but WITHOUT ANY WARRANTY; without even the implied warranty of
//     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//     GNU General Public License for more details.
//
//     You should have received a copy of the GNU General Public License
//     along with this program.  If not, see <https://www.gnu.org/licenses/>.

extern crate core;

use std::error::Error;
use std::fs;
use std::path::Path;
use dotenvy::dotenv;
use log::{error, info};
use rocket::{catchers, Request, Response, routes};
use rocket::fairing::{Fairing, Info, Kind};
use rocket::http::Header;
use sha2::Sha256;
use sqlx::migrate::Migrator;
use sqlx::postgres::PgPoolOptions;
use crate::config::TFConfig;
use crate::kv::{kv_get, kv_set};
use sha2::Digest;

pub mod format;
pub mod util;
pub mod db;
pub mod config;
pub mod tokens;
pub mod routes;
pub mod auth;
pub mod crypto;
pub mod org;
pub mod kv;
pub mod role;

static MIGRATOR: Migrator = sqlx::migrate!();

pub struct CORS;

#[rocket::async_trait]
impl Fairing for CORS {
    fn info(&self) -> Info {
        Info {
            name: "Add CORS headers to responses",
            kind: Kind::Response,
        }
    }

    async fn on_response<'r>(&self, _request: &'r Request<'_>, response: &mut Response<'r>) {
        response.set_header(Header::new("Access-Control-Allow-Origin", "*"));
        response.set_header(Header::new("Access-Control-Allow-Methods", "POST, GET, PATCH, OPTIONS"));
        response.set_header(Header::new("Access-Control-Allow-Headers", "*"));
        response.set_header(Header::new("Access-Control-Allow-Credentials", "true"));
    }
}

#[rocket::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let _ = rocket::build();

    info!("[tfapi] loading config");

    let _ = dotenv();

    if std::env::var("CONFIG_FILE").is_err() && !Path::new("config.toml").exists() {
        error!("[tfapi] fatal: the environment variable CONFIG_FILE is not set");
        error!("[tfapi]  help: try creating a .env file that sets it");
        error!("[tfapi]  help: or, create a file config.toml with your config, as it is loaded automatically");
        std::process::exit(1);
    }

    let config_file = if Path::new("config.toml").exists() {
        "config.toml".to_string()
    } else {
        std::env::var("CONFIG_FILE").unwrap()
    };

    let config_data = match fs::read_to_string(&config_file) {
        Ok(d) => d,
        Err(e) => {
            error!("[tfapi] fatal: unable to read config from {}", config_file);
            error!("[tfapi] fatal: {}", e);
            std::process::exit(1);
        }
    };

    let config: TFConfig = match toml::from_str(&config_data) {
        Ok(c) => c,
        Err(e) => {
            error!("[tfapi] fatal: unable to parse config from {}", config_file);
            error!("[tfapi] fatal: {}", e);
            std::process::exit(1);
        }
    };

    info!("[tfapi] connecting to database pool");

    let pool = match PgPoolOptions::new().max_connections(5).connect(&config.db_url).await {
        Ok(p) => p,
        Err(e) => {
            error!("[tfapi] fatal: unable to connect to database pool");
            error!("[tfapi] fatal: {}", e);
            std::process::exit(1);
        }
    };

    info!("[tfapi] running database migrations");

    MIGRATOR.run(&pool).await?;

    info!("[tfapi] verifying encryption key");

    let kv_hash = kv_get("pmk_hash", &pool).await.expect("Unable to get pmk hash from kv store");

    let mut hasher = Sha256::new();
    hasher.update(config.data_key.as_bytes());
    let config_hash = hex::encode(hasher.finalize());
    if let Some(k_hash) = kv_hash {
        if config_hash != k_hash {
            error!("[tfapi] fatal: instance master key does not match key used to encrypt data");
            error!("[tfapi] fatal: datastore was encrypted with keyid {k_hash}");
            error!("[tfapi] fatal: the key in your config has the keyid {config_hash}");
            error!("[tfapi] fatal: you probably changed data_key. please return it to it's original value");
            std::process::exit(1);
        } else {
            info!("[tfapi] data keyid is {config_hash}");
        }
    } else {
        info!("[tfapi] detected first run");
        info!("[tfapi] welcome to trifid!");
        info!("[tfapi] data keyid is {config_hash}");

        if let Err(e) = kv_set("pmk_hash", config_hash.as_str(), &pool).await {
            error!("[tfapi] fatal: unable to set pmk_hash in kv store");
            error!("[tfapi] fatal: the database returned the following error:");
            error!("[tfapi] fatal: {e}");
            std::process::exit(1);
        } else {
            info!("[tfapi] configured instance information in kv store");
        }
    }

    info!("[tfapi] building rocket config");

    let figment = rocket::Config::figment().merge(("port", config.listen_port));

    let _ = rocket::custom(figment)
        .mount("/", routes![
            crate::routes::v1::auth::magic_link::magiclink_request,
            crate::routes::v1::auth::magic_link::options,
            crate::routes::v1::signup::signup_request,
            crate::routes::v1::signup::options,
            crate::routes::v1::auth::verify_magic_link::verify_magic_link,
            crate::routes::v1::auth::verify_magic_link::options,
            crate::routes::v1::totp_authenticators::totp_authenticators_request,
            crate::routes::v1::totp_authenticators::options,
            crate::routes::v1::verify_totp_authenticator::verify_totp_authenticator_request,
            crate::routes::v1::verify_totp_authenticator::options,
            crate::routes::v1::auth::totp::totp_request,
            crate::routes::v1::auth::totp::options,
            crate::routes::v1::auth::check_session::check_session,
            crate::routes::v1::auth::check_session::check_session_auth,
            crate::routes::v1::auth::check_session::options,
            crate::routes::v1::auth::check_session::options_auth,
            crate::routes::v2::whoami::whoami_request,
            crate::routes::v2::whoami::options,

            crate::routes::v1::organization::options,
            crate::routes::v1::organization::orgidoptions,
            crate::routes::v1::organization::orginfo_req,
            crate::routes::v1::organization::orglist_req,
            crate::routes::v1::organization::create_org,
            crate::routes::v1::user::get_user,
            crate::routes::v1::user::options,
            crate::routes::v1::organization::createorgoptions,
            crate::routes::v1::ca::get_cas_for_org,
            crate::routes::v1::ca::options,
            crate::routes::v1::roles::get_roles,
            crate::routes::v1::roles::options,
            crate::routes::v1::roles::options_roleadd,
            crate::routes::v1::roles::role_add
        ])
        .register("/", catchers![
            crate::routes::handler_400,
            crate::routes::handler_401,
            crate::routes::handler_403,
            crate::routes::handler_404,
            crate::routes::handler_422,

            crate::routes::handler_500,
            crate::routes::handler_501,
            crate::routes::handler_502,
            crate::routes::handler_503,
            crate::routes::handler_504,
            crate::routes::handler_505,
        ])
        .attach(CORS)
        .manage(pool)
        .manage(config)
        .launch().await?;

    Ok(())
}