use cookie::Cookie;
use fastly::http::{header, Method, StatusCode, Url};
use fastly::{mime, Error, Request, Response};
use rand::{thread_rng, Rng};

const BACKEND_NAME: &str = "origin_0";
const CAPTCHA_BACKEND_NAME: &str = "origin_1";

fn captcha_form(failed: bool, url: &mut Url) -> Result<Response, Error> {
    url.query_pairs_mut().append_pair("captcha", "1");
    let mut resp = Response::new()
        .with_content_type(mime::TEXT_HTML)
        .with_header(header::CACHE_CONTROL, "private, no-store")
        .with_body(format!(
            r#"
<html>
<head>
  <script src='https://www.google.com/recaptcha/api.js'></script>
  <script>function submit() {{ document.forms[0].submit(); }} </script>
</head>
  <body>
    <h1>A quick check...</h1>
    <p>Based on analysis of your activity, we've mistakenly classified you as a robot!  Sorry about that.  Please complete the question below to set us straight and we'll have you straight back to what you were doing in no time.</p>
    <form method="POST" action="{}">
      <div
        class="g-recaptcha"
        data-sitekey="6LdurkYUAAAAADsbJu5xi1r8irfESBX8F5XBAUgd"
        data-callback="submit"
      ></div>
    </form>
  </body>
</html>
"#, htmlescape::encode_minimal(url.as_str())
        ));
    if failed {
        resp.set_header("captchaFail", "1");
    }
    Ok(resp)
}

#[fastly::main]
fn main(mut req: Request) -> Result<Response, Error> {
    log_fastly::init_simple("my_log", log::LevelFilter::Info);

    // If the user is submitting a solution to a CAPTCHA challenge,
    // send it to the validation service to verify it.
    let mut qs_vec: Vec<(String, String)> = req.get_query()?;
    if req.get_method() == Method::POST
        && qs_vec
            .iter()
            .any(|(name, value)| name.trim().eq_ignore_ascii_case("captcha") && value.trim() == "1")
    {
        log::info!("Sending to CAPTCHA service to verify");

        // Remove captcha for the query string
        qs_vec.retain(|(name, _)| name != "captcha");
        req.set_query(&qs_vec)?;
        let mut url = req.get_url().clone();

        // Prepare to send the request to the captcha service
        req.set_header("origURL", &url);
        req.set_path("/validate/6LdurkYUAAAAADDdbQveEGdjTy6kIKXulAv8N7a8");
        req.remove_query();
        req.set_pass(true);
        let beresp = req.send(CAPTCHA_BACKEND_NAME)?;
        match beresp.get_header_str("Captcha-Result") {
            Some("OK") => {
                // It's a pass! Set a cookie, so that this user is not
                // challenged again within an hour. You would probably
                // want to make this cookie harder to fake, by using
                // a signature of some kind, incorporating the user's IP.
                Ok(Response::new()
                    .with_status(StatusCode::TEMPORARY_REDIRECT)
                    .with_header(header::CACHE_CONTROL, "private, no-store")
                    .with_header(header::SET_COOKIE, "captchaAuth=1; path=/; max-age=3600")
                    .with_header(header::LOCATION, url))
            }
            _ => {
                log::info!("Failed CAPTCHA check. Redisplay form");
                captcha_form(true, &mut url)
            }
        }
    } else {
  
        if is_risky_request(&req) {
            // TODO: Validate cookie signature
            if req
                .get_header_str(header::COOKIE)
                .unwrap_or("")
                .split(';')
                .filter_map(|entry| Cookie::parse(entry.trim()).ok())
                .all(|cookie| cookie.name() != "captchaAuth")
            {
                log::info!("Request flagged for CAPTCHA");
                return captcha_form(false, req.get_url_mut());
            }
        } else {
            log::info!("Request not flagged for CAPTCHA");
        }

        // Otherwise set the regular backend
        Ok(req.send(BACKEND_NAME)?)
    }
}

fn is_risky_request(_: &Request) -> bool {
  // Assess risk and decide whether to trigger CAPTCHA.
  // In practice you'd use an ACL or bot detection service.
  true
}