resuma 1.3.1

//! Staged per-request metadata for the next SSR page response.
//!
//! Storage is **task-local** when the request runs inside [`scope_page_staging`]
//! (installed around every request by `request_id_middleware`). Thread-local
//! storage is unsafe here: page renders use `block_in_place` + `block_on`,
//! which can interleave other tasks on the same worker thread and clobber the
//! staged CSRF token / CSP nonce mid-render. Code running outside a scoped
//! request task (tests, direct `render_to_string` callers) falls back to a
//! thread-local slot, preserving the previous synchronous behavior.

use std::cell::RefCell;
use std::future::Future;

#[derive(Default)]
struct PageStaging {
    cache_control: Option<String>,
    csrf: String,
    csp_nonce: String,
    status: Option<u16>,
    redirect: Option<String>,
    title: Option<String>,
    description: Option<String>,
    robots: Option<String>,
    canonical: Option<String>,
    json_ld: Option<String>,
    dir: Option<String>,
    theme: Option<String>,
}

tokio::task_local! {
    static PAGE_STAGING: RefCell<PageStaging>;
}

thread_local! {
    static FALLBACK_STAGING: RefCell<PageStaging> = RefCell::new(PageStaging::default());
}

/// Run `fut` with fresh, task-isolated page staging (one scope per request).
pub async fn scope_page_staging<F: Future>(fut: F) -> F::Output {
    PAGE_STAGING
        .scope(RefCell::new(PageStaging::default()), fut)
        .await
}

fn with_staging<R>(f: impl FnOnce(&mut PageStaging) -> R) -> R {
    let mut f = Some(f);
    match PAGE_STAGING.try_with(|cell| (f.take().expect("staging fn"))(&mut cell.borrow_mut())) {
        Ok(out) => out,
        Err(_) => {
            FALLBACK_STAGING.with(|cell| (f.take().expect("staging fn"))(&mut cell.borrow_mut()))
        }
    }
}

/// Clear all staged per-request page metadata (call at request boundaries).
pub fn clear_request_staging() {
    with_staging(|s| *s = PageStaging::default());
}

/// Stage the HTTP status for the page about to be returned (e.g. 404/500 error
/// pages). Defaults to 200 when unset. Non-200 statuses also stage `noindex`
/// unless [`set_page_robots`] already ran.
pub fn stage_response_status(status: u16) {
    with_staging(|s| {
        s.status = Some(status);
        if status != 200 && s.robots.is_none() {
            s.robots = Some("noindex".into());
        }
    });
}

/// Take the staged HTTP status (consumed once per response).
pub fn take_response_status() -> Option<u16> {
    with_staging(|s| s.status.take())
}

/// Stage a `Location` redirect for the page about to be returned — turns the
/// SSR response into a real `3xx` with no HTML body. Pair with
/// `stage_response_status` (e.g. `ResumaError::Redirect` already stages both
/// via `FlowError`/`error_page`).
pub fn stage_response_redirect(location: impl Into<String>) {
    let location = location.into();
    with_staging(|s| s.redirect = Some(location));
}

/// Take the staged redirect location (consumed once per response).
pub fn take_response_redirect() -> Option<String> {
    with_staging(|s| s.redirect.take())
}

/// Stage a `Cache-Control` header for the page about to be returned.
pub fn stage_response_cache_control(value: impl Into<String>) {
    let value = value.into();
    with_staging(|s| s.cache_control = Some(value));
}

/// When a response sets a session CSRF cookie, shared caches must not store it.
pub fn sanitize_cache_for_session(
    cache: Option<String>,
    sets_session_cookie: bool,
) -> Option<String> {
    if !sets_session_cookie {
        return cache;
    }
    if cache.as_deref().is_some_and(|c| {
        let lower = c.to_ascii_lowercase();
        lower.contains("public") || lower.contains("max-age")
    }) {
        tracing::warn!("overriding Cache-Control to private, no-store — page sets CSRF cookie");
    }
    Some("private, no-store".to_string())
}

/// Take a staged cache header (consumed once per response).
pub fn take_response_cache_control() -> Option<String> {
    with_staging(|s| s.cache_control.take())
}

/// Override `<title>` for this SSR response. Call during page render; the
/// document wrapper reads it after the view is written (non-streaming).
pub fn set_page_title(title: impl Into<String>) {
    let title = title.into();
    with_staging(|s| s.title = Some(title));
}

/// Staged document title, if the page set one.
pub fn page_title_override() -> Option<String> {
    with_staging(|s| s.title.clone())
}

/// Override the meta description for this SSR response.
pub fn set_page_description(description: impl Into<String>) {
    let description = description.into();
    with_staging(|s| s.description = Some(description));
}

/// Staged meta description, if the page set one.
pub fn page_description_override() -> Option<String> {
    with_staging(|s| s.description.clone())
}

/// Override `meta name="robots"` for this SSR response.
pub fn set_page_robots(robots: impl Into<String>) {
    let robots = robots.into();
    with_staging(|s| s.robots = Some(robots));
}

/// Staged robots directive, if the page or a non-200 status set one.
pub fn page_robots_override() -> Option<String> {
    with_staging(|s| s.robots.clone())
}

/// Override the HTML canonical URL for this SSR response (absolute).
pub fn set_page_canonical(canonical: impl Into<String>) {
    let canonical = canonical.into();
    with_staging(|s| s.canonical = Some(canonical));
}

/// Staged canonical URL, if the page set one.
pub fn page_canonical_override() -> Option<String> {
    with_staging(|s| s.canonical.clone())
}

/// Override JSON-LD for this SSR response (raw JSON, not HTML-wrapped).
pub fn set_page_json_ld(json_ld: impl Into<String>) {
    let json_ld = json_ld.into();
    with_staging(|s| s.json_ld = Some(json_ld));
}

/// Staged JSON-LD, if the page set one.
pub fn page_json_ld_override() -> Option<String> {
    with_staging(|s| s.json_ld.clone())
}

/// Override `<html dir>` for this SSR response (`ltr`, `rtl`, or `auto`).
pub fn set_page_dir(dir: impl Into<String>) {
    let dir = dir.into();
    with_staging(|s| s.dir = Some(dir));
}

/// Staged document `dir`, if the page set one.
pub fn page_dir_override() -> Option<String> {
    with_staging(|s| s.dir.clone())
}

/// Override `<html data-theme>` for this SSR response.
pub fn set_page_theme(theme: impl Into<String>) {
    let theme = theme.into();
    with_staging(|s| s.theme = Some(theme));
}

/// Staged document theme id, if the page set one.
pub fn page_theme_override() -> Option<String> {
    with_staging(|s| s.theme.clone())
}

/// Stage the CSRF token for forms rendered during this page pass.
pub fn stage_page_csrf(token: impl Into<String>) {
    let token = token.into();
    with_staging(|s| s.csrf = token);
}

/// CSRF token for the current page render (forms).
pub fn page_csrf() -> String {
    with_staging(|s| s.csrf.clone())
}

/// Stage the CSP nonce for inline/module scripts rendered during this page pass.
pub fn stage_page_csp_nonce(nonce: impl Into<String>) {
    let nonce = nonce.into();
    with_staging(|s| s.csp_nonce = nonce);
}

/// CSP nonce for the current page render (client components, inline scripts).
pub fn page_csp_nonce() -> String {
    with_staging(|s| s.csp_nonce.clone())
}

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

    #[test]
    fn sanitize_cache_for_session_overrides_public() {
        let out = sanitize_cache_for_session(Some("public, max-age=3600".into()), true);
        assert_eq!(out.as_deref(), Some("private, no-store"));
    }

    #[test]
    fn redirect_staging_round_trips_and_is_consumed_once() {
        clear_request_staging();
        assert_eq!(take_response_redirect(), None);
        stage_response_redirect("/login");
        assert_eq!(take_response_redirect(), Some("/login".to_string()));
        // Consumed — a second take (e.g. a following request reusing the
        // same fallback thread-local slot) must not see a stale redirect.
        assert_eq!(take_response_redirect(), None);
    }

    #[test]
    fn sanitize_cache_for_session_keeps_when_no_cookie() {
        let cache = "public, max-age=60".to_string();
        let out = sanitize_cache_for_session(Some(cache.clone()), false);
        assert_eq!(out, Some(cache));
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn page_staging_isolated_per_scoped_task() {
        let mut handles = Vec::new();
        for i in 0..32u32 {
            handles.push(tokio::spawn(async move {
                scope_page_staging(async {
                    let token = format!("csrf-{i:04}");
                    let nonce = format!("nonce-{i:04}");
                    stage_page_csrf(token.clone());
                    stage_page_csp_nonce(nonce.clone());
                    tokio::task::yield_now().await;
                    assert_eq!(page_csrf(), token);
                    assert_eq!(page_csp_nonce(), nonce);
                })
                .await
            }));
        }
        for h in handles {
            h.await.unwrap();
        }
    }
}