use cookie::{Cookie, CookieJar};
use fastly::http::{header, HeaderName, HeaderValue, Method, StatusCode, Url};
use fastly::{mime, Error, Request, Response};

const PAYWALL_BACKEND: &str = "origin_1";
const CONTENT_BACKEND: &str = "origin_0";

#[fastly::main]
fn main(mut req: Request) -> Result<Response, Error> {
    // Ingest cookies
    if let Some(req_cookie_jar) = req
        .get_header(header::COOKIE)
        .and_then(|h| parse_cookies_to_jar(h).ok())
    {
        req.remove_header(header::COOKIE);
        // Getting a session ID from a cookie is rather simplistic in this
        // demo.  For a more advanced solution, combine the paywall pattern
        // with a JSON web token verification at the edge
        req.set_header(
            "auth-sessionid",
            req_cookie_jar
                .get("sessionid")
                .map(|c| c.value())
                .unwrap_or_default(),
        );
    }

    // Only GETs and HEADs are considered for paywall restriction.
    if !matches!(req.get_method(), &Method::GET | &Method::HEAD) {
        return Ok(req.send(CONTENT_BACKEND)?);
    }

    let _ = req.take_body(); // request bodies will not be forwarded.

    // run the content request
    let mut content_resp = req.send(CONTENT_BACKEND)?;

    // if 403 or 5xx - don't continue
    if content_resp.get_status() == 403 || content_resp.get_status().is_server_error() {
        return Ok(content_resp);
    }

    // If the response contains `paywall` header we run a check against the URL
    let paywall_hdr = match content_resp.remove_header("paywall") {
        None => {
            // No paywall header received, deliver the content
            return Ok(content_resp);
        }
        Some(paywall_hdr) => paywall_hdr,
    };

    let paywall_url = Url::parse(paywall_hdr.to_str()?)?;
    if paywall_url.host_str() != Some("fiddle-paywall.glitch.me") {
        // Not an expected paywall, deliver the content as is
        return Ok(content_resp);
    }

    // First, re-create a request with original headers we received from client
    let mut paywall_req = content_resp.take_backend_request().unwrap();
    paywall_req.set_url(paywall_url);

    // run the request against the paywall and get a response
    let paywall_resp = paywall_req.send(PAYWALL_BACKEND)?;

    // Check the paywall response
    if let Some(result_value) = paywall_resp.get_header("paywall-result") {
        if result_value == "BLOCK" {
            // if answer is `BLOCK` - deliver an error
            return Ok(Response::from_body(
                "<p>This content is only available to premium subscribers.</p>\n",
            )
            .with_content_type(mime::TEXT_HTML)
            .with_status(StatusCode::FORBIDDEN));
        }
        let meta_header = &HeaderName::from_static("paywall-meta");
        if let Some(paywall_meta) = paywall_resp.get_header(meta_header) {
            content_resp.set_header(meta_header, paywall_meta.clone());
        }
    }
    Ok(content_resp)
}

fn parse_cookies_to_jar(value: &HeaderValue) -> Result<CookieJar, Error> {
    let mut jar = CookieJar::new();
    for cookie in value.to_str()?.split(';').map(Cookie::parse) {
        jar.add_original(cookie?.into_owned());
    }
    Ok(jar)
}