use std::borrow::Cow;

use anyhow::anyhow;
use log::*;
use trillium::{Handler, KnownHeaderName, conn_try};

use crate::options::ServerOptions;

struct FixLocation {
    options: &'static ServerOptions,
}

pub fn fix_location_headers(options: &'static ServerOptions) -> impl Handler {
    FixLocation { options }
}

#[trillium::async_trait]
impl Handler for FixLocation {
    async fn run(&self, conn: trillium::Conn) -> trillium::Conn {
        let host = if let Some(h) = &self.options.host {
            Cow::Borrowed(h.as_str())
        } else if let Some(hh) = conn.request_headers().get(KnownHeaderName::Host) {
            hh.to_string().into()
        } else {
            warn!("no host header");
            return conn;
        };

        let scheme = if host.starts_with("http://") || host.starts_with("https://") {
            ""
        } else if conn.is_secure() {
            "https://"
        } else {
            "http://"
        };
        let mut conn = conn;
        let hdr = conn.response_headers_mut();
        if let Some(loc) = hdr.get(KnownHeaderName::Location) {
            let loc = conn_try!(
                loc.as_str()
                    .ok_or_else(|| anyhow!("non-UTF8 location header")),
                conn
            );
            if loc.starts_with("/") {
                debug!("fixing location header {}", loc);
                let mut url = scheme.to_string();
                url += &host;
                url += loc;
                hdr.insert(KnownHeaderName::Location, url);
            }
        }

        conn
    }
}
