use fastly::http::header;
use fastly::{Body, Error, Request, Response};
use quick_xml::{
    events::{BytesText, Event},
    Reader, Writer,
};
use std::time::{Duration, UNIX_EPOCH};

fn parse_since(since: &str) -> String {
    match since.parse::<u64>() {
        Ok(secs) => {
            match (UNIX_EPOCH + Duration::from_secs(secs))
                .elapsed()
                .map(|d| d.as_secs())
                .unwrap_or(0)
            {
                n if n < 60 => format!("{}s ago", n),
                n if n < 3600 => format!("{}m ago", n / 60),
                n if n < 96 * 3600 => format!("{}h ago", n / 3600),
                n => format!("{}d ago", n / 86400),
            }
        }
        _ => "".to_string(),
    }
}

#[fastly::main]
fn main(mut req: Request) -> Result<Response, Error> {
    // Make sure the origin server does not gzip the response. We
    // can't process ESIs in gzipped content.
    req.remove_header(header::ACCEPT_ENCODING);
    let mut beresp = req.send("origin_0")?;

    let mut reader = Reader::from_reader(beresp.get_body_mut());
    reader.config_mut().trim_markup_names_in_closing_tags = false;
    reader.config_mut().check_end_names = false;
    let mut writer = Writer::new(Body::new());
    let mut buf = Vec::new();
    let mut remove = false;

    loop {
        buf.clear();
        match reader.read_event_into(&mut buf) {
            // Handle <esi:remove> tags.
            Ok(quick_xml::events::Event::Start(elem)) if elem.starts_with(b"esi:remove") => {
              remove = true;
            },
            Ok(quick_xml::events::Event::End(elem)) if elem.starts_with(b"esi:remove") => {
              remove = false;
            }
            _ if remove => continue,
            // Handle <esi:include> tags.
            // For a complete ESI solution for Compute, see docs.rs/esi
            Ok(Event::Empty(elem)) => match elem.name().as_ref() {
              b"esi:include" => {
                if let Some(path) = elem.attributes().find_map(|attr| {
                  let attribute = attr.unwrap();
                  if attribute.key.as_ref() == b"src" {
                    attribute.decode_and_unescape_value(&reader).ok()
                  } else {
                    None
                  }
                }) {
                    // We assume here that the source page's ESI includes are
                    // in the form /__since/<unix-timestamp>, and we need to
                    // convert that into a relative time.
                    if let Some(since) = path.strip_prefix("/__since/") {
                        let ts = parse_since(since);
                        writer.write_event(Event::Text(BytesText::from_escaped(ts)))?;
                    }
                }
              },
              _ => {}
            }
            Ok(Event::Eof) => break,
            Ok(e) => writer.write_event(&e)?,
            _ => {}
        }
    }
    Ok(beresp.with_body(writer.into_inner()))
}