use fastly::http::{header, Method, StatusCode};
use fastly::config_store::ConfigStore;
use fastly::{Error, Request, Response};
use hmac_sha256::HMAC;
use base64::{
    alphabet,
    engine::{self, general_purpose},
    Engine as _,
};
const CUSTOM_ENGINE: engine::GeneralPurpose =
    engine::GeneralPurpose::new(&alphabet::URL_SAFE, general_purpose::NO_PAD);

const MS_API_VERSION: &str = "2017-11-09";

/// Holds Azure Blob storage access information
struct AzureConfig {
    account_key: String,
    account_name: String,
    blob_container: String,
    host: String,
}

impl AzureConfig {
    // Load the Azure configuration from Config Store.
    fn load_config() -> Self {
        let dict = ConfigStore::open("azure_config");
        let account_name = dict.get("account_name").expect("Account name not found");
        Self {
            account_key: dict.get("account_key").expect("Account key not found"),
            blob_container: dict.get("blob_container_name").expect("Account blob not found"),
            host: format!("{}blob.core.windows.net", &account_name),
            account_name,
        }
    }
}

#[fastly::main]
fn main(mut req: Request) -> Result<Response, Error> {
    let azure_config = AzureConfig::load_config();

    // Only generate authorize header for GET and HEAD request
    // Pass PURGE to backend as it is
    // Block other kinds of method
    let resp = match req.get_method() {
        &Method::GET | &Method::HEAD => {
            req = authorize_azure_storage_request(
                req,
                &azure_config.account_key,
                &azure_config.account_name,
                &azure_config.blob_container,
            )?;

            req.send("azure_backend")?
        }
        m if m == "PURGE" => {
            // url of cached object has no query string, and has blob container added to path
            // so when do the purge, we need to make sure the url matches
            req.remove_query();
            req.set_path(&format!(
                "/{}{}",
                azure_config.blob_container,
                req.get_path()
            ));

            req.send("azure_backend")?
        }
        _ => Response::from_status(StatusCode::METHOD_NOT_ALLOWED),
    };

    Ok(resp)
}

/// Generate Authorization headers of the Azure request
fn authorize_azure_storage_request(
    mut req: Request,
    storage_access_key: &str,
    storage_account_name: &str,
    blob_container: &str,
) -> Result<Request, Error> {
    // The date format should be something like
    // Fri, 01 Jan 2021 00:14:22 GMT
    let ms_date = chrono::Utc::now().format("%a, %d %b %Y %T GMT").to_string();

    // Decode base64 format access key to binary key
    let access_key_decoded = CUSTOM_ENGINE.decode(storage_access_key)?;

    // Canonical headers must be sorted and concatenated with newlines.
    // If there are query params there is a separate spec for that, not supported here yet
    let canonical_headers = format!("x-ms-date:{}\nx-ms-version:{}\n", ms_date, MS_API_VERSION);

    // Canonical Resource is /ACCOUNT/CONTAINER/FILE
    let canonical_resource = format!(
        "/{}/{}{}",
        storage_account_name,
        blob_container,
        req.get_path()
    );

    // Construct everything properly before signing
    // We are adding 4 newlines here, however the spec says we
    // can override any of these 4 headers with values if we want.
    // For now we are just blanking them out
    let string_to_sign = format!("GET\n\n\n\n{}{}", canonical_headers, canonical_resource);
    // HMAC-sign with SHA256 and Base64-encode the result
    let signature_binary = HMAC::mac(string_to_sign.as_bytes(), &access_key_decoded);
    let signature = CUSTOM_ENGINE.encode(signature_binary);

    let authorization = format!("SharedKeyLite {}:{}", storage_account_name, signature);

    // request path for azure storage is /containder/object_id
    req.set_path(&format!("/{}{}", blob_container, req.get_path()));
    req.remove_query();

    // For the the backend request method as GET
    req.set_method(Method::GET);

    // Add authorization related headers to request
    req.set_header(header::AUTHORIZATION, &authorization);
    req.set_header("x-ms-date", &ms_date);
    req.set_header("x-ms-version", MS_API_VERSION);

    println!("Path: {}, Authorization: {}", req.get_path(), authorization);

    Ok(req)
}