use anyhow::{anyhow, Result};
use fastly::http::StatusCode;
use fastly::{Error, Request, Response};
use regex::Regex;

/// The name of a backend server associated with this service.
const BACKEND_NAME: &str = "origin_0";

/// Secret hmac key value.
const SECRET: &[u8] = b"THESECRET";

#[fastly::main]
fn main(req: Request) -> Result<Response, Error> {
    // Log initial setup
    log_fastly::init_simple("my_log", log::LevelFilter::Info);
    fastly::log::set_panic_endpoint("my_log")?;

    match verify_token(req) {
        Ok(request) => {
            // Token signature verification passed, send out the request to backend
            Ok(request.send(BACKEND_NAME)?)
        }
        Err(e) => {
            // Token signature verification failed
            // Log the error message
            log::info!("error {}", e);

            // Generate error response
            Ok(Response::from_body(e.to_string()).with_status(StatusCode::FORBIDDEN))
        }
    }
}

/// Valid request, verify the token if needed.
fn verify_token(mut req: Request) -> Result<Request> {
    // Get basename from URL
    let req_path = req.get_path();
    let base_name = req_path.rsplit('/').next().unwrap();
    log::info!("{} {}", req_path, base_name);

    // If request is not playlist.m3u8 or manifest.mpd, no need token verification
    if !matches!(base_name, "playlist.m3u8" | "manifest.mpd") {
        return Ok(req);
    }

    // Get token information from URI
    let token_cap = Regex::new(r"/token=nva=(\d+)~dirs=(\d+)~hash=0([[:xdigit:]]{20})(/.*)")?
        .captures(req_path)
        .ok_or_else(|| anyhow!("Missing token in url path"))?;
    let nva = token_cap
        .get(1)
        .ok_or_else(|| anyhow!("Failed to get token nva"))?
        .as_str();
    let dirs = token_cap
        .get(2)
        .ok_or_else(|| anyhow!("Failed to get token dirs"))?
        .as_str();
    let hash = token_cap
        .get(3)
        .ok_or_else(|| anyhow!("Failed to get token hash"))?
        .as_str();
    let path = token_cap
        .get(4)
        .ok_or_else(|| anyhow!("Failed to get token path"))?
        .as_str()
        .to_owned();
    log::info!("nav={}; dirs={}; hash={}; path={}", nva, dirs, hash, path);

    // Validate the number of dirs
    let expected_dir_count = path.matches('/').count() - 1;
    log::info!("expected_dir_count={}", expected_dir_count);
    if dirs.parse::<usize>()? != expected_dir_count {
        return Err(anyhow!("Invalid number of dirs"));
    }

    // Validate hash
    let token = format!("{}?nva={}&dirs={}", path, nva, dirs);
    log::info!("token={}", token);

    let hash_bytes = hex::decode(hash)?;
    let calculated_hash_bytes = hmacsha1::hmac_sha1(SECRET, token.as_bytes());
    let hash_bytes_expected = &calculated_hash_bytes[0..10];

    if !secure_equal(hash_bytes_expected, &hash_bytes) {
        return Err(anyhow!("Token is incorrect"));
    }

    // Validate timestamp
    let expire_time: i64 = nva.parse()?;
    let now = chrono::Utc::now().timestamp();

    if expire_time < now {
        return Err(anyhow!(
            "Token has expired, now {}, timestamp {}",
            now,
            expire_time
        ));
    }

    // Set URI for backend access for production
    if BACKEND_NAME == "origin_0" {
        // For debugging
        req.set_path("/json");
    } else {
        // For production
        req.set_path(&format!("/resource{}", path));
    }

    Ok(req)
}

/// Compare two byte slices with constant execution time. This provides some protection
/// against timing side-channel attacks.
fn secure_equal(a: &[u8], b: &[u8]) -> bool {
    a.len() == b.len() && 0 == a.iter().zip(b).fold(0, |r, (a, b)| r | (a ^ b))
}