use fastly::geo::geo_lookup;
use fastly::{Error, Request, Response};
use fastly::http::StatusCode;


#[fastly::main]
fn main(mut req: Request) -> Result<Response, Error> {
     log_fastly::init_simple("mylogs", log::LevelFilter::Info);
     
    let client_geo = req.get_client_ip_addr().and_then(geo_lookup).unwrap();
    let country_code = client_geo.country_code();

    match country_code {
        "US" | "PR" | "AS" | "GU" | "VI" => {
                req.set_header("X-Country-Code", country_code);
                log::info!( "Country Code is {}", country_code);

                Ok(Response::new()
                    .with_status(StatusCode::FORBIDDEN)
                    .with_body("Restricted Content"))
        }
        _ => {
            let beresp = req.send("origin_0")?;
            Ok(beresp)
        }
    }
}