axum_static_include 0.3.0

//! 用于 `axum` 包含静态文件的宏
//! 
//! 这个宏主要用于将静态文件以 `include_*!()` 的形式来包含进二进制文件中
//! 
//! 


mod util;

use proc_macro::TokenStream;
use syn::{parse_macro_input, Item};
use quote::quote;

// 生成的 handler 函数放在一个单独的私有 mod 里
// handler 函数命名以路径的 hash 做

/// 打包一个文件夹里的所有文件进二进制程序，并产生以此文件夹为根目录的路由
/// 
/// # Example
/// ```no_run
/// #[axum_static_include::static_serve("assert")]
/// fn assert_fold() -> axum::Router {}
/// 
/// #[tokio::main]
/// async fn main() {
///     let asserts = assert_fold();
///     let app = Router::new()
///         .nest("/static", asserts)
///     // ...
/// }
/// ```
#[proc_macro_attribute]
pub fn static_serve(attr: TokenStream, item: TokenStream) -> TokenStream {
    let dir = attr.to_string()
        .strip_prefix("\"")
        .and_then(|s| s.strip_suffix("\""))
        .unwrap()
        .to_string();

    let path = std::path::Path::new(&dir);
    if !path.exists() {
        panic!("path not exists: {dir}");
    }
    if !path.is_dir() {
        panic!("path is not a dir: {dir}");
    }
    let dir = direx::DirInfo::walk_dir(dir).unwrap();
    let all = dir.all_files();
    let fpaths = dir.all_files_map(|f| {
        f.path_str()
            .replace(r"\", "/")
    });
    let fpaths_abs = dir.all_files_abs();
    let (_dirs, fnames) = all.into_iter()
        .map(|f| (f.dir_path.as_ref(), &f.fname))
        .unzip::<_, _, Vec<_>, Vec<_>>();

    let ftypes = fnames.iter()
        .map(|fname| util::parse_file_type(fname) )
        .collect::<Vec<_>>();

    let hash_etags = fpaths.iter()
        .map(|s| sha256::digest(s))
        .collect::<Vec<_>>();

    let handler_names = fpaths.iter()
        .map(|s| format!("h{}", sha256::digest(s)).parse::<proc_macro2::TokenStream>() )
        .collect::<Result<Vec<_>, _>>()
        .unwrap();

    let router_names = fpaths.iter()
        .map(|fpath| fpath.strip_prefix(dir.path.as_ref()))
        .collect::<Option<Vec<_>>>()
        .unwrap();

    let mod_name: proc_macro2::TokenStream = format!("m{}", sha256::digest(fpaths.join(""))).parse().unwrap();

    let ast = parse_macro_input!(item as Item);
    let r = match &ast {
        Item::Fn(f) => {
            let vis = &f.vis;
            let sig = &f.sig;

            quote! {
                #vis #sig {
                    axum::Router::new()
                        #(.route(#router_names, axum::routing::get(#mod_name::#handler_names)))*
                }

                mod #mod_name {
                    use axum::http::{HeaderMap, HeaderValue, StatusCode, header};
                    #(pub async fn #handler_names(hm: HeaderMap) -> (StatusCode, HeaderMap, &'static [u8]) {
                        let rhm = HeaderMap::from_iter([
                            (header::CONTENT_TYPE, HeaderValue::from_static(#ftypes)),
                            (header::ETAG, HeaderValue::from_str(#hash_etags).unwrap()),
                        ]);
                        match hm.get(header::IF_NONE_MATCH) {
                            Some(he) if he.to_str().unwrap() == #hash_etags => (
                                StatusCode::NOT_MODIFIED,
                                rhm,
                                &[],
                            ),
                            _ => (
                                StatusCode::OK, 
                                rhm,
                                include_bytes!(#fpaths_abs)
                            )
                        }
                    })*
                }
            }
        },
        _ => panic!("Only support for function"),
    };
    r.into()
}