use async_trait::async_trait;
use headers::{ContentLength, ContentType};
use pingora_core::server::Server;
use pingora_proxy::{ProxyHttp, Session};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use mime_guess::from_path;

struct StaticFileServer {
    root_dir: PathBuf,
}

impl StaticFileServer {
    fn new(root: impl AsRef<Path>) -> Self {
        StaticFileServer {
            root_dir: root.as_ref().to_path_buf(),
        }
    }
}

#[async_trait]
impl ProxyHttp for StaticFileServer {
    async fn response_filter(&self, session: &mut Session, response: &mut pingora_http::Response) {
        let path = session.req_header().uri.path();
        let file_path = self.root_dir.join(&path[1..]); // 去除前导/

        if let Ok(metadata) = tokio::fs::metadata(&file_path).await {
            if metadata.is_file() {
                if let Ok(file) = tokio::fs::read(&file_path).await {
                    // 设置正确的Content-Type
                    let mime = from_path(&file_path).first_or_octet_stream();
                    response.headers_mut().insert(
                        headers::CONTENT_TYPE,
                        ContentType::from(mime).to_string().parse().unwrap(),
                    );

                    // 设置Content-Length
                    response.headers_mut().insert(
                        headers::CONTENT_LENGTH,
                        ContentLength(metadata.len()).to_string().parse().unwrap(),
                    );

                    // 设置缓存头（CDN关键配置）
                    response.headers_mut().insert(
                        "cache-control",
                        "public, max-age=31536000".parse().unwrap(), // 1年缓存
                    );

                    *response.body_mut() = file.into();
                    response.set_status(200);
                    return;
                }
            }
        }

        // 文件不存在
        *response.status_mut() = 404;
        *response.body_mut() = b"Not Found".to_vec().into();
    }
}

#[tokio::main]
async fn main() {
    // 1. 初始化服务器
    let mut server = Server::new(Some("Static File Server")).unwrap();
    server.bootstrap();

    // 2. 创建文件服务（假设静态文件在./static目录）
    let static_service = StaticFileServer::new("./static");
    let mut service = pingora_proxy::http_proxy_service(&server.configuration, static_service);

    // 3. 监听端口
    service.add_tcp("0.0.0.0:8080");

    // 4. 启动服务
    server.add_service(service);
    println!("Static file server running on :8080");
    server.run_forever();
}

