use fastly::{Error, Request, Response};
use glob::Pattern;

/// Backend name configured in service
const BACKEND_NAME: &str = "origin_0";

/// Set `true` to use allow list, `false` to use block list
const CONFIG_USE_ALLOW_LIST: bool = true;

/// Allow List of header pattern (Unix shell style)
const ALLOW_HEADER_PATTERNS: &[&str] = &[
    "access-control-allow-*",
    "content-type",
    "x-cache-*",
    "x-my-info-*",
];

/// Block List of header pattern (Unix shell style)
const BLOCK_HEADER_PATTERNS: &[&str] = &[
    "fastly-debug-*",
    "x-goog-*",
    "x-github-*",
    "x-amz-*",
    "x-generator",
    "via",
];

/// Policy used to clean up headers
enum CleanHeaderPolicy {
    /// Only keep headers matches patterns in pattern list
    Keep,
    /// Remove header that matches any pattern in the pattern list
    Remove,
}

#[fastly::main]
fn main(req: Request) -> Result<Response, Error> {
    // Send the request to backend
    let mut resp = req.send(BACKEND_NAME)?;

    // Clean up the response headers
    if CONFIG_USE_ALLOW_LIST {
        clean_up_headers(&mut resp, ALLOW_HEADER_PATTERNS, CleanHeaderPolicy::Keep);
    } else {
        clean_up_headers(&mut resp, BLOCK_HEADER_PATTERNS, CleanHeaderPolicy::Remove);
    }

    Ok(resp)
}

/// Clean up response headers
fn clean_up_headers(resp: &mut Response, patterns_list: &[&str], policy: CleanHeaderPolicy) {
    let globs: Vec<_> = patterns_list
        .iter()
        .map(|s| Pattern::new(s).unwrap())
        .collect();

    let remove_header_vec: Vec<_> = resp
        .get_header_names()
        .filter(|h| {
            let matches = globs.iter().any(|g| g.matches(h.as_str()));

            match policy {
                CleanHeaderPolicy::Keep => !matches,
                CleanHeaderPolicy::Remove => matches,
            }
        })
        .cloned()
        .collect();

    for header in remove_header_vec {
        resp.remove_header(header);
    }
}