mod claim;
mod kvmap;
mod serviceauth;
mod sites;
mod xrpc;

use claim::SitesStore;
use kvmap::DomainMapping;
use worker::*;

fn r2_key(did: &str, rkey: &str, site_path: &str) -> String {
    let base = format!("{}/{}/", did, rkey);
    if site_path.is_empty() || site_path == "/" {
        format!("{}index.html", base)
    } else {
        let trimmed = site_path.trim_start_matches('/');
        if trimmed.is_empty() || trimmed.ends_with('/') {
            format!("{}{}index.html", base, trimmed)
        } else {
            format!("{}{}", base, trimmed)
        }
    }
}

// decodes per segment so an encoded slash can't start a new path
fn site_path(path: &str) -> Option<String> {
    let mut segments = Vec::new();
    for segment in path.split('/') {
        let decoded = percent_encoding::percent_decode_str(segment)
            .decode_utf8()
            .ok()?;
        if decoded == "." || decoded == ".." || decoded.contains(['/', '\\', '\0']) {
            return None;
        }
        segments.push(decoded);
    }
    Some(segments.join("/"))
}

fn is_service_host(audience: &str, host: &str) -> bool {
    audience.strip_prefix("did:web:") == Some(host)
}

fn needs_trailing_slash(path: &str) -> bool {
    if path == "/" || path.ends_with('/') {
        return false;
    }
    let last_segment = path.rsplit('/').next().unwrap_or(path);
    !last_segment.contains('.')
}

fn with_trailing_slash(url: &Url) -> String {
    let mut url = url.clone();
    url.set_path(&format!("{}/", url.path()));
    url.to_string()
}

async fn fetch_from_r2(bucket: &Bucket, key: &str) -> Result<Option<Object>> {
    if let Some(obj) = bucket.get(key).execute().await? {
        return Ok(Some(obj));
    }

    let last_segment = key.rsplit('/').next().unwrap_or(key);
    if !last_segment.contains('.') {
        let index_key = format!("{}/index.html", key.trim_end_matches('/'));
        if let Some(obj) = bucket.get(&index_key).execute().await? {
            return Ok(Some(obj));
        }
    }

    Ok(None)
}

fn response_from_object(obj: Object) -> Result<Response> {
    let content_type = obj
        .http_metadata()
        .content_type
        .unwrap_or_else(|| "application/octet-stream".to_string());

    let body = obj
        .body()
        .ok_or_else(|| Error::RustError("empty R2 body".into()))?;
    let mut resp = Response::from_body(body.response_body()?)?;
    resp.headers_mut().set("Content-Type", &content_type)?;
    resp.headers_mut()
        .set("Cache-Control", "public, max-age=60")?;
    Ok(resp)
}

#[event(fetch)]
async fn fetch(req: Request, env: Env, ctx: Context) -> Result<Response> {
    let origin = req.headers().get("Origin")?.filter(|o| !o.is_empty());

    if req.method() == Method::Options {
        return preflight_response(&origin);
    }

    let resp = fetch_inner(req, env, ctx).await?;
    let Some(o) = &origin else {
        return Ok(resp);
    };
    // a response passed through from a fetch has immutable headers
    let headers = resp.headers().clone();
    headers.set("Access-Control-Allow-Origin", o)?;
    headers.set("Vary", "Origin")?;
    Ok(resp.with_headers(headers))
}

fn preflight_response(origin: &Option<String>) -> Result<Response> {
    let mut resp = Response::empty()?.with_status(204);
    if let Some(o) = origin {
        resp.headers_mut().set("Access-Control-Allow-Origin", o)?;
        resp.headers_mut()
            .set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")?;
        resp.headers_mut().set(
            "Access-Control-Allow-Headers",
            "authorization, content-type, accept",
        )?;
        resp.headers_mut().set("Access-Control-Max-Age", "86400")?;
    }
    Ok(resp)
}

const HANDLE_TITLE_TTL: u64 = 300;

// shorter negative cache so fresh signups aren't delayed
const UNKNOWN_HANDLE_TTL: u64 = 60;

async fn handle_title(kv: &KvStore, host: &str) -> Result<Option<String>, String> {
    let key = format!("handle-title:{}", host);
    if let Some(memo) = kv.get(&key).text().await.map_err(|e| e.to_string())? {
        let memo = memo.trim();
        return Ok((!memo.is_empty()).then(|| memo.to_string()));
    }

    let resolved = claim::resolve_handle(host.to_string()).await?;
    kv.put(&key, resolved.as_deref().unwrap_or(""))
        .map_err(|e| e.to_string())?
        .expiration_ttl(if resolved.is_none() {
            UNKNOWN_HANDLE_TTL
        } else {
            HANDLE_TITLE_TTL
        })
        .execute()
        .await
        .map_err(|e| e.to_string())?;
    Ok(resolved)
}

async fn fetch_inner(req: Request, env: Env, ctx: Context) -> Result<Response> {
    let url = req.url()?;
    let path = url.path();

    if path.starts_with("/xrpc/") {
        return xrpc::route(req, env, ctx).await;
    }

    let kv = env.kv("SITES")?;
    let bucket = env.bucket("SITES_BUCKET")?;

    let host = req.headers().get("host")?.unwrap_or_default();
    let host = host.split(':').next().unwrap_or("").to_string();

    if host.is_empty() {
        return Response::error("Bad Request: missing host", 400);
    }
    // the service's own host is its did:web, so no site is served there
    let audience = env
        .var("SITES_AUDIENCE_DID")
        .map(|v| v.to_string())
        .unwrap_or_default();
    if is_service_host(&audience, &host) {
        return Response::error("Not Found", 404);
    }

    if path == "/.well-known/atproto-did" {
        // from the domain claim, never the site
        if req.method() != Method::Get {
            let mut resp = Response::error("Method Not Allowed", 405)?;
            resp.headers_mut().set("Allow", "GET")?;
            return Ok(resp);
        }
        // handles are the pdses' to answer, and a fetch on our own zone skips
        // the worker and lands on the origin in front of them
        if claim::is_handle_domain(&host) {
            return Fetch::Request(req).send().await;
        }
        let store = SitesStore::d1(env.d1("SITES_DB")?);
        return match store.select_claim_by_domain(&host).await {
            Ok(Some(claim)) if claim.deleted.is_none() => {
                let mut resp = Response::ok(claim.did)?;
                resp.headers_mut()
                    .set("Cache-Control", "public, max-age=300")?;
                Ok(resp)
            }
            Ok(_) => Response::error("Not Found", 404),
            Err(e) => {
                console_warn!("atproto-did claim lookup for {} failed: {}", host, e);
                Response::error("Internal Error", 500)
            }
        };
    }

    if needs_trailing_slash(path) {
        let redirect_url = with_trailing_slash(&url);
        return Response::redirect_with_status(redirect_url.parse()?, 308);
    }

    let mapping = match kv.get(&host).text().await? {
        Some(raw) => {
            let trimmed = raw.trim();
            if trimmed.is_empty() || trimmed == "{}" {
                return Response::error("site not found!", 404);
            }
            match serde_json::from_str::<DomainMapping>(&raw) {
                Ok(m) => m,
                Err(_) => return Response::error("Internal Error: bad mapping", 500),
            }
        }
        None => return Response::error("site not found!", 404),
    };

    if claim::is_handle_domain(&host) {
        match handle_title(&kv, &host).await {
            Ok(resolved) if resolved.as_deref() == Some(mapping.did.as_str()) => {}
            Ok(_) => {
                console_warn!(
                    "refusing to serve {}: the zone does not resolve it to its mapping's owner",
                    host
                );
                return Response::error("Not Found", 404);
            }
            Err(e) => {
                console_warn!("handle authority lookup for {} failed: {}", host, e);
                return Response::error("Service Unavailable", 503);
            }
        }
    }

    let Some(path) = site_path(url.path()) else {
        return Response::error("Not Found", 404);
    };
    let path = path.as_str();

    let first_segment = path
        .trim_start_matches('/')
        .split('/')
        .next()
        .unwrap_or("")
        .to_string();

    if !first_segment.is_empty() {
        if let Some(entry) = mapping.repos.get(&first_segment) {
            if !entry.is_index() {
                let site_path = path
                    .trim_start_matches('/')
                    .trim_start_matches(&first_segment)
                    .to_string();

                let key = r2_key(&mapping.did, entry.rkey(&first_segment), &site_path);
                return match fetch_from_r2(&bucket, &key).await? {
                    Some(obj) => response_from_object(obj),
                    None => Response::error("Not Found", 404),
                };
            }
        }
    }

    if let Some((name, entry)) = mapping.index_repo() {
        let key = r2_key(&mapping.did, entry.rkey(name), path);
        return match fetch_from_r2(&bucket, &key).await? {
            Some(obj) => response_from_object(obj),
            None => Response::error("Not Found", 404),
        };
    }

    Response::error("Not Found", 404)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn a_served_path_is_the_uploaded_file_name() {
        for (requested, key) in [
            ("/sp%20ace.txt", "did/site/sp ace.txt"),
            ("/%C3%BC.txt", "did/site/\u{fc}.txt"),
            ("/a%2520b.txt", "did/site/a%20b.txt"),
            (
                "/docs/%E6%97%A5%E6%9C%AC/",
                "did/site/docs/\u{65e5}\u{672c}/index.html",
            ),
            ("/plain.txt", "did/site/plain.txt"),
        ] {
            let path = site_path(requested).unwrap();
            assert_eq!(r2_key("did", "site", &path), key, "{}", requested);
        }
    }

    #[test]
    fn the_service_host_is_not_a_site() {
        let audience = "did:web:sites.next-sites.tngl.sh";
        assert!(is_service_host(audience, "sites.next-sites.tngl.sh"));
        assert!(!is_service_host(audience, "alpha.next-sites.tngl.sh"));
        assert!(!is_service_host(audience, "next-sites.tngl.sh"));
        assert!(!is_service_host(
            "did:plc:sites",
            "sites.next-sites.tngl.sh"
        ));
    }

    #[test]
    fn a_decoded_path_cannot_leave_its_segment() {
        for requested in [
            "/%2e%2e/index.html",
            "/a/%2E/b",
            "/x%2f..%2findex.html",
            "/x%2Fy",
            "/a%5c..%5cb",
            "/nul%00.txt",
            "/bad%ff.txt",
        ] {
            assert_eq!(site_path(requested), None, "{}", requested);
        }
    }
}