use axum::{
    body::Body,
    http::{Request, StatusCode},
};
use mq::sqlite_queue::SqliteTaskQueue;
use server::{api::AppState, db::Database, scheduler::GridScheduler};
use std::sync::Arc;
use tower::ServiceExt; // for oneshot

#[tokio::test]
async fn test_server_api_flow() {
    let temp_dir = tempfile::tempdir().unwrap();
    let db_path = temp_dir.path().join("api_db.db");
    let queue_db_path = temp_dir.path().join("api_queue.db");
    let results_dir = temp_dir.path().join("results");
    std::fs::create_dir_all(&results_dir).unwrap();

    let db = Database::new(&db_path.to_string_lossy()).await.unwrap();
    let queue = Arc::new(
        SqliteTaskQueue::new(&queue_db_path.to_string_lossy())
            .await
            .unwrap(),
    );
    let scheduler = Arc::new(GridScheduler::new(
        db.clone(),
        queue.clone(),
        results_dir.to_string_lossy().to_string(),
    ));

    let state = AppState {
        db,
        queue,
        scheduler,
        results_dir: results_dir.to_string_lossy().to_string(),
        rate_limiter: server::api::rate_limit::RateLimiter::new(
            5,
            std::time::Duration::from_secs(300),
        ),
        auth_token: None,
        admin_token: None,
        auth_disabled: false,
        admin_sessions: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
    };

    let app = axum::Router::new()
        .route(
            "/api/node/register",
            axum::routing::post(server::api::node::register_node),
        )
        .route(
            "/api/node/heartbeat",
            axum::routing::post(server::api::node::heartbeat_node),
        )
        .route(
            "/api/task/claim",
            axum::routing::post(server::api::task::claim_task),
        )
        .route(
            "/api/status",
            axum::routing::get(server::api::status::get_status),
        )
        .route(
            "/api/workflows",
            axum::routing::get(server::api::workflow::list_workflows)
                .post(server::api::workflow::save_workflow),
        )
        .with_state(state);

    // 1. Check status API
    let response = app
        .clone()
        .oneshot(
            Request::builder()
                .uri("/api/status")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);

    // 2. Register node API
    let reg_body = serde_json::json!({
        "node_id": "test-node-api",
        "host_name": "api-host",
        "max_slots": 8
    });
    let response = app
        .clone()
        .oneshot(
            Request::builder()
                .method("POST")
                .uri("/api/node/register")
                .header("content-type", "application/json")
                .body(Body::from(serde_json::to_vec(&reg_body).unwrap()))
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);

    // 3. Save Workflow API
    let wf_body = serde_json::json!({
        "name": "test_api_wf",
        "description": "Test Workflow Description",
        "config_yaml": "grid:\n  teff: [35000]\n  logg: [5.5]\n  loghe: [-1]\n  logc: [-2]\n  logn: [-2]\n  logo: [-2]"
    });
    let response = app
        .clone()
        .oneshot(
            Request::builder()
                .method("POST")
                .uri("/api/workflows")
                .header("content-type", "application/json")
                .body(Body::from(serde_json::to_vec(&wf_body).unwrap()))
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);
}

#[tokio::test]
async fn test_auth_middleware_scope_and_running_status() {
    let temp_dir = tempfile::tempdir().unwrap();
    let db_path = temp_dir.path().join("auth_db.db");
    let queue_db_path = temp_dir.path().join("auth_queue.db");
    let results_dir = temp_dir.path().join("results");
    std::fs::create_dir_all(&results_dir).unwrap();

    let db = Database::new(&db_path.to_string_lossy()).await.unwrap();
    let queue = Arc::new(
        SqliteTaskQueue::new(&queue_db_path.to_string_lossy())
            .await
            .unwrap(),
    );
    let scheduler = Arc::new(GridScheduler::new(
        db.clone(),
        queue.clone(),
        results_dir.to_string_lossy().to_string(),
    ));

    let state = AppState {
        db: db.clone(),
        queue: queue.clone(),
        scheduler,
        results_dir: results_dir.to_string_lossy().to_string(),
        rate_limiter: server::api::rate_limit::RateLimiter::new(
            5,
            std::time::Duration::from_secs(300),
        ),
        auth_token: Some("secret_token_123".to_string()),
        admin_token: Some("secret_token_123".to_string()),
        auth_disabled: false,
        admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
            std::collections::HashMap::new(),
        )),
    };

    let api_router = axum::Router::new().route(
        "/status",
        axum::routing::get(server::api::status::get_status),
    );

    let auth_layer =
        axum::middleware::from_fn_with_state(state.clone(), server::api::auth_middleware);
    let api_router = api_router.layer(auth_layer);

    let app = axum::Router::new()
        .nest("/api", api_router)
        .with_state(state);

    // Unauthenticated API request -> 401 Unauthorized
    let res = app
        .clone()
        .oneshot(
            Request::builder()
                .uri("/api/status")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::UNAUTHORIZED);

    // Authenticated API request -> 200 OK
    let res = app
        .clone()
        .oneshot(
            Request::builder()
                .uri("/api/status")
                .header("authorization", "Bearer secret_token_123")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::OK);

    // Test mark_grid_point_running
    let params = common::models::GridPointParams {
        teff: 35000.0,
        logg: 5.5,
        loghe: -1.0,
        logc: -2.0,
        logn: -2.0,
        logo: -2.0,
    };
    db.upsert_grid_point(&params, 0, "test_wf").await.unwrap();
    db.mark_grid_point_running(&params.model_name(), "test_wf")
        .await
        .unwrap();

    let stats = db.get_grid_summary_stats(None).await.unwrap();
    assert_eq!(stats["running"], 1);
}

/// L2 鉴权核心流程测试：
/// 注册节点 → 颁发专属 token → 用 token 调 heartbeat（200）→ 吊销 → 再调（401）。
#[tokio::test]
async fn test_l2_node_token_issue_revoke_flow() {
    let temp_dir = tempfile::tempdir().unwrap();
    let db_path = temp_dir.path().join("l2_db.db");

    let db = Database::new(&db_path.to_string_lossy()).await.unwrap();

    // 1. 注册节点
    let reg = common::models::NodeRegisterRequest {
        node_id: "node-l2-test".to_string(),
        host_name: "l2-host".to_string(),
        max_slots: 4,
    };
    db.register_node(&reg).await.unwrap();

    // 2. 颁发专属 token，返回明文
    let token = db.issue_node_token("node-l2-test").await.unwrap();
    assert!(!token.is_empty());

    // 2b. 一次性取走暂存明文（take_pending_node_token）：首次取到与颁发一致的明文，
    //     再次取为 None（取走即焚）。验证 #4 简化为单一 UPDATE...RETURNING 后行为一致。
    let pending = db.take_pending_node_token("node-l2-test").await.unwrap();
    assert_eq!(pending.as_deref(), Some(token.as_str()));
    let pending2 = db.take_pending_node_token("node-l2-test").await.unwrap();
    assert!(
        pending2.is_none(),
        "取走即焚：第二次 take_pending 必须返回 None"
    );
    // 不存在的 node take 也应返回 None（不报错）
    assert!(db
        .take_pending_node_token("node-not-exist")
        .await
        .unwrap()
        .is_none());

    // 3. token 可反查到 node_id（此调用会把 token 写入内存缓存）
    let found = db.find_node_by_token(&token).await;
    assert_eq!(found.as_deref(), Some("node-l2-test"));

    // 4. 错误 token 查不到
    assert!(db.find_node_by_token("wrong-token").await.is_none());

    // 5. 吊销后 token 立即失效
    db.revoke_node_token("node-l2-test").await.unwrap();
    assert!(db.find_node_by_token(&token).await.is_none());

    // 6. 重新颁发后恢复（覆盖式）
    let token2 = db.issue_node_token("node-l2-test").await.unwrap();
    assert_eq!(
        db.find_node_by_token(&token2).await.as_deref(),
        Some("node-l2-test")
    );
    // 旧 token 仍失效（已被覆盖）
    assert!(db.find_node_by_token(&token).await.is_none());
}

/// 验证中间件对 node token 的端到端鉴权：
/// 用真实 node token 调 /node/heartbeat 应 200；吊销后再调应 401。
#[tokio::test]
async fn test_middleware_node_role_with_token() {
    let temp_dir = tempfile::tempdir().unwrap();
    let db_path = temp_dir.path().join("l2_mw_db.db");
    let queue_db_path = temp_dir.path().join("l2_mw_queue.db");
    let results_dir = temp_dir.path().join("results");
    std::fs::create_dir_all(&results_dir).unwrap();

    let db = Database::new(&db_path.to_string_lossy()).await.unwrap();
    let queue = Arc::new(
        SqliteTaskQueue::new(&queue_db_path.to_string_lossy())
            .await
            .unwrap(),
    );
    let scheduler = Arc::new(GridScheduler::new(
        db.clone(),
        queue.clone(),
        results_dir.to_string_lossy().to_string(),
    ));

    // 注册并颁发 token
    let reg = common::models::NodeRegisterRequest {
        node_id: "node-mw-test".to_string(),
        host_name: "mw-host".to_string(),
        max_slots: 2,
    };
    db.register_node(&reg).await.unwrap();
    let token = db.issue_node_token("node-mw-test").await.unwrap();

    let state = AppState {
        db: db.clone(),
        queue,
        scheduler,
        results_dir: results_dir.to_string_lossy().to_string(),
        rate_limiter: server::api::rate_limit::RateLimiter::new(
            5,
            std::time::Duration::from_secs(300),
        ),
        auth_token: Some("placeholder".to_string()),
        admin_token: None,
        auth_disabled: false,
        admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
            std::collections::HashMap::new(),
        )),
    };

    let api_router = axum::Router::new().route(
        "/node/heartbeat",
        axum::routing::post(server::api::node::heartbeat_node),
    );
    let auth_layer =
        axum::middleware::from_fn_with_state(state.clone(), server::api::auth_middleware);
    let api_router = api_router.layer(auth_layer);
    let app = axum::Router::new()
        .nest("/api", api_router)
        .with_state(state);

    // 用有效 node token 调 heartbeat → 200
    let hb_body = serde_json::json!({
        "node_id": "node-mw-test", "active_slots": 1, "cpu_usage": 10.0, "memory_usage": 20.0
    });
    let res = app
        .clone()
        .oneshot(
            Request::builder()
                .method("POST")
                .uri("/api/node/heartbeat")
                .header("authorization", format!("Bearer {}", token))
                .header("content-type", "application/json")
                .body(Body::from(serde_json::to_vec(&hb_body).unwrap()))
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::OK);

    // 吊销后同样请求 → 401
    db.revoke_node_token("node-mw-test").await.unwrap();
    let res = app
        .oneshot(
            Request::builder()
                .method("POST")
                .uri("/api/node/heartbeat")
                .header("authorization", format!("Bearer {}", token))
                .header("content-type", "application/json")
                .body(Body::from(serde_json::to_vec(&hb_body).unwrap()))
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::UNAUTHORIZED);
}

/// 管理 API（Admin）端到端测试：
/// 列出节点 → 吊销 → 重发 token → 鉴权（admin 放行，node/匿名 401）。
#[tokio::test]
async fn test_admin_node_management_api() {
    let temp_dir = tempfile::tempdir().unwrap();
    let db_path = temp_dir.path().join("admin_db.db");
    let queue_db_path = temp_dir.path().join("admin_queue.db");
    let results_dir = temp_dir.path().join("results");
    std::fs::create_dir_all(&results_dir).unwrap();

    let db = Database::new(&db_path.to_string_lossy()).await.unwrap();
    let queue = Arc::new(
        SqliteTaskQueue::new(&queue_db_path.to_string_lossy())
            .await
            .unwrap(),
    );
    let scheduler = Arc::new(GridScheduler::new(
        db.clone(),
        queue.clone(),
        results_dir.to_string_lossy().to_string(),
    ));

    // 预置一个已注册并颁发 token 的节点
    let reg = common::models::NodeRegisterRequest {
        node_id: "node-admin-test".to_string(),
        host_name: "admin-host".to_string(),
        max_slots: 2,
    };
    db.register_node(&reg).await.unwrap();
    let original_token = db.issue_node_token("node-admin-test").await.unwrap();

    let state = AppState {
        db: db.clone(),
        queue,
        scheduler,
        results_dir: results_dir.to_string_lossy().to_string(),
        rate_limiter: server::api::rate_limit::RateLimiter::new(
            5,
            std::time::Duration::from_secs(300),
        ),
        auth_token: Some("admin-secret".to_string()),
        admin_token: Some("admin-secret".to_string()),
        auth_disabled: false,
        admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
            std::collections::HashMap::new(),
        )),
    };

    let api_router = axum::Router::new()
        .route(
            "/admin/nodes",
            axum::routing::get(server::api::admin::list_nodes),
        )
        .route(
            "/admin/nodes/:node_id/revoke",
            axum::routing::post(server::api::admin::revoke_node),
        )
        .route(
            "/admin/nodes/:node_id/reissue",
            axum::routing::post(server::api::admin::reissue_node),
        );
    let auth_layer =
        axum::middleware::from_fn_with_state(state.clone(), server::api::auth_middleware);
    let app = axum::Router::new()
        .nest("/api", api_router.layer(auth_layer))
        .with_state(state);

    // 1. 匿名访问 /admin/nodes → 401
    let res = app
        .clone()
        .oneshot(
            Request::builder()
                .uri("/api/admin/nodes")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::UNAUTHORIZED);

    // 2. admin 访问 /admin/nodes → 200，且能看到预置节点 token_status=active
    let res = app
        .clone()
        .oneshot(
            Request::builder()
                .uri("/api/admin/nodes")
                .header("authorization", "Bearer admin-secret")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::OK);
    let body: serde_json::Value = serde_json::from_slice(
        &axum::body::to_bytes(res.into_body(), usize::MAX)
            .await
            .unwrap(),
    )
    .unwrap();
    let nodes = body["data"].as_array().expect("data 应为数组");
    let target = nodes
        .iter()
        .find(|n| n["node_id"] == "node-admin-test")
        .expect("应包含预置节点");
    assert_eq!(target["token_status"], "active");

    // 3. admin 吊销节点 token → 200
    let res = app
        .clone()
        .oneshot(
            Request::builder()
                .method("POST")
                .uri("/api/admin/nodes/node-admin-test/revoke")
                .header("authorization", "Bearer admin-secret")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::OK);
    // 旧 token 立即失效
    assert!(db.find_node_by_token(&original_token).await.is_none());

    // 4. admin 重发 token → 200，返回新明文 token
    let res = app
        .clone()
        .oneshot(
            Request::builder()
                .method("POST")
                .uri("/api/admin/nodes/node-admin-test/reissue")
                .header("authorization", "Bearer admin-secret")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::OK);
    let body: serde_json::Value = serde_json::from_slice(
        &axum::body::to_bytes(res.into_body(), usize::MAX)
            .await
            .unwrap(),
    )
    .unwrap();
    let new_token = body["node_token"].as_str().expect("应返回 node_token");
    assert!(!new_token.is_empty());
    assert_ne!(new_token, original_token);
    // 新 token 可用
    assert_eq!(
        db.find_node_by_token(new_token).await.as_deref(),
        Some("node-admin-test")
    );

    // 5. node token 不能访问 admin API → 401（即便持有有效 node token）
    let res = app
        .oneshot(
            Request::builder()
                .uri("/api/admin/nodes")
                .header("authorization", format!("Bearer {}", new_token))
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::UNAUTHORIZED);
}

/// S1 身份绑定测试：节点 A 用自己 token 冒充节点 B 发心跳 → 403；一致时 → 200。
#[tokio::test]
async fn test_node_heartbeat_node_id_binding() {
    let temp_dir = tempfile::tempdir().unwrap();
    let db_path = temp_dir.path().join("bind_db.db");
    let queue_db_path = temp_dir.path().join("bind_queue.db");
    let results_dir = temp_dir.path().join("results");
    std::fs::create_dir_all(&results_dir).unwrap();

    let db = Database::new(&db_path.to_string_lossy()).await.unwrap();
    let queue = Arc::new(
        SqliteTaskQueue::new(&queue_db_path.to_string_lossy())
            .await
            .unwrap(),
    );
    let scheduler = Arc::new(GridScheduler::new(
        db.clone(),
        queue,
        results_dir.to_string_lossy().to_string(),
    ));

    // 预置节点 A，颁发 token
    let reg = common::models::NodeRegisterRequest {
        node_id: "node-A".to_string(),
        host_name: "h".to_string(),
        max_slots: 2,
    };
    db.register_node(&reg).await.unwrap();
    let token_a = db.issue_node_token("node-A").await.unwrap();

    let state = AppState {
        db: db.clone(),
        queue: Arc::new(SqliteTaskQueue::new(":memory:").await.unwrap()),
        scheduler,
        results_dir: results_dir.to_string_lossy().to_string(),
        rate_limiter: server::api::rate_limit::RateLimiter::new(
            5,
            std::time::Duration::from_secs(300),
        ),
        auth_token: Some("x".into()),
        admin_token: None,
        auth_disabled: false,
        admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
            std::collections::HashMap::new(),
        )),
    };
    let api_router = axum::Router::new().route(
        "/node/heartbeat",
        axum::routing::post(server::api::node::heartbeat_node),
    );
    let auth_layer =
        axum::middleware::from_fn_with_state(state.clone(), server::api::auth_middleware);
    let app = axum::Router::new()
        .nest("/api", api_router.layer(auth_layer))
        .with_state(state);

    // A 用自己 token，但 body 声称 node_id=node-B（冒充）→ 403
    let hb =
        serde_json::json!({"node_id":"node-B","active_slots":1,"cpu_usage":0.0,"memory_usage":0.0});
    let res = app
        .clone()
        .oneshot(
            Request::builder()
                .method("POST")
                .uri("/api/node/heartbeat")
                .header("authorization", format!("Bearer {}", token_a))
                .header("content-type", "application/json")
                .body(Body::from(serde_json::to_vec(&hb).unwrap()))
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::FORBIDDEN);

    // A 用自己 token，body node_id=node-A（一致）→ 200
    let hb_ok =
        serde_json::json!({"node_id":"node-A","active_slots":1,"cpu_usage":0.0,"memory_usage":0.0});
    let res = app
        .oneshot(
            Request::builder()
                .method("POST")
                .uri("/api/node/heartbeat")
                .header("authorization", format!("Bearer {}", token_a))
                .header("content-type", "application/json")
                .body(Body::from(serde_json::to_vec(&hb_ok).unwrap()))
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::OK);
}

#[tokio::test]
async fn test_admin_login_flow() {
    let temp_dir = tempfile::tempdir().unwrap();
    let db_path = temp_dir.path().join("login_db.db");
    let queue_db_path = temp_dir.path().join("login_queue.db");
    let results_dir = temp_dir.path().join("results");
    std::fs::create_dir_all(&results_dir).unwrap();

    let db = Database::new(&db_path.to_string_lossy()).await.unwrap();
    let queue = Arc::new(
        SqliteTaskQueue::new(&queue_db_path.to_string_lossy())
            .await
            .unwrap(),
    );
    let scheduler = Arc::new(GridScheduler::new(
        db.clone(),
        queue.clone(),
        results_dir.to_string_lossy().to_string(),
    ));

    let admin_pass = "my_short_admin_password_123";
    let state = AppState {
        db,
        queue,
        scheduler,
        results_dir: results_dir.to_string_lossy().to_string(),
        rate_limiter: server::api::rate_limit::RateLimiter::new(
            5,
            std::time::Duration::from_secs(300),
        ),
        auth_token: Some(admin_pass.to_string()),
        admin_token: Some(admin_pass.to_string()),
        auth_disabled: false,
        admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
            std::collections::HashMap::new(),
        )),
    };

    let api_router = axum::Router::new()
        .route("/login", axum::routing::post(server::api::auth::login))
        .route(
            "/auth/check",
            axum::routing::get(server::api::auth::check_auth),
        )
        .layer(axum::middleware::from_fn_with_state(
            state.clone(),
            server::api::auth_middleware,
        ));

    let app = axum::Router::new()
        .nest("/api", api_router)
        .with_state(state);

    // 1. 密码错误 ➔ 401
    let wrong_body = serde_json::json!({ "password": "wrong_password" });
    let res = app
        .clone()
        .oneshot(
            Request::builder()
                .method("POST")
                .uri("/api/login")
                .header("content-type", "application/json")
                .extension(axum::extract::ConnectInfo(std::net::SocketAddr::from((
                    [127, 0, 0, 1],
                    12345,
                ))))
                .body(Body::from(serde_json::to_vec(&wrong_body).unwrap()))
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::UNAUTHORIZED);

    // 2. 正确密码 ➔ 200 + token
    let right_body = serde_json::json!({ "password": admin_pass });
    let res = app
        .clone()
        .oneshot(
            Request::builder()
                .method("POST")
                .uri("/api/login")
                .header("content-type", "application/json")
                .extension(axum::extract::ConnectInfo(std::net::SocketAddr::from((
                    [127, 0, 0, 1],
                    12345,
                ))))
                .body(Body::from(serde_json::to_vec(&right_body).unwrap()))
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::OK);
    let bytes = axum::body::to_bytes(res.into_body(), 1024 * 1024)
        .await
        .unwrap();
    let json: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
    assert_eq!(json["success"], true);
    let token = json["token"].as_str().unwrap();
    assert!(token.len() == 64);

    // 3. 携带拿到的 Token 访问受保护的 /api/auth/check ➔ 200
    let res = app
        .oneshot(
            Request::builder()
                .method("GET")
                .uri("/api/auth/check")
                .header("authorization", format!("Bearer {}", token))
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::OK);
}

#[tokio::test]
async fn test_node_approval_workflow() {
    let temp_dir = tempfile::tempdir().unwrap();
    let db_path = temp_dir.path().join("appr_db.db");
    let queue_db_path = temp_dir.path().join("appr_queue.db");
    let results_dir = temp_dir.path().join("results");
    std::fs::create_dir_all(&results_dir).unwrap();

    let db = Database::new(&db_path.to_string_lossy()).await.unwrap();
    let queue = Arc::new(
        SqliteTaskQueue::new(&queue_db_path.to_string_lossy())
            .await
            .unwrap(),
    );
    let scheduler = Arc::new(GridScheduler::new(
        db.clone(),
        queue.clone(),
        results_dir.to_string_lossy().to_string(),
    ));

    let admin_pass = "admin_approval_secret";
    let state = AppState {
        db,
        queue,
        scheduler,
        results_dir: results_dir.to_string_lossy().to_string(),
        rate_limiter: server::api::rate_limit::RateLimiter::new(
            5,
            std::time::Duration::from_secs(300),
        ),
        auth_token: Some(admin_pass.to_string()),
        admin_token: Some(admin_pass.to_string()),
        auth_disabled: false,
        admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
            std::collections::HashMap::new(),
        )),
    };

    let api_router = axum::Router::new()
        .route(
            "/node/register",
            axum::routing::post(server::api::node::register_node),
        )
        .route(
            "/node/check_status",
            axum::routing::post(server::api::node::check_node_status),
        )
        .route(
            "/node/heartbeat",
            axum::routing::post(server::api::node::heartbeat_node),
        )
        .route(
            "/admin/nodes/:node_id/approve",
            axum::routing::post(server::api::admin::approve_node),
        )
        .layer(axum::middleware::from_fn_with_state(
            state.clone(),
            server::api::auth_middleware,
        ));

    let app = axum::Router::new()
        .nest("/api", api_router)
        .with_state(state);

    // 1. 新节点免凭据申请注册 ➔ 200 + status: pending_approval
    let reg_body = serde_json::json!({
        "node_id": "node-pending-01",
        "host_name": "worker-host",
        "max_slots": 4
    });
    let res = app
        .clone()
        .oneshot(
            Request::builder()
                .method("POST")
                .uri("/api/node/register")
                .header("content-type", "application/json")
                .body(Body::from(serde_json::to_vec(&reg_body).unwrap()))
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::OK);
    let bytes = axum::body::to_bytes(res.into_body(), 1024 * 1024)
        .await
        .unwrap();
    let json: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
    assert_eq!(json["status"], "pending_approval");

    // 2. Node 端轮询查状态 ➔ status: pending_approval
    let check_body = serde_json::json!({ "node_id": "node-pending-01" });
    let res = app
        .clone()
        .oneshot(
            Request::builder()
                .method("POST")
                .uri("/api/node/check_status")
                .header("content-type", "application/json")
                .body(Body::from(serde_json::to_vec(&check_body).unwrap()))
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::OK);
    let bytes = axum::body::to_bytes(res.into_body(), 1024 * 1024)
        .await
        .unwrap();
    let json: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
    assert_eq!(json["status"], "pending_approval");

    // 3. 管理员在 Dashboard 点击同意 ➔ 200
    let res = app
        .clone()
        .oneshot(
            Request::builder()
                .method("POST")
                .uri("/api/admin/nodes/node-pending-01/approve")
                .header("authorization", format!("Bearer {}", admin_pass))
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::OK);
    let bytes = axum::body::to_bytes(res.into_body(), 1024 * 1024)
        .await
        .unwrap();
    println!("Step 3 body: {}", String::from_utf8_lossy(&bytes));

    // 4. Node 端再次轮询查状态 ➔ status: approved + 获取 node_token
    let res = app
        .clone()
        .oneshot(
            Request::builder()
                .method("POST")
                .uri("/api/node/check_status")
                .header("content-type", "application/json")
                .body(Body::from(serde_json::to_vec(&check_body).unwrap()))
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::OK);
    let bytes = axum::body::to_bytes(res.into_body(), 1024 * 1024)
        .await
        .unwrap();
    let json: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
    println!("Step 4 status returned: {:?}", json);

    assert_eq!(json["status"], "approved");
    let node_token = json["node_token"].as_str().unwrap();

    // 5. Node 携带拿到到的专属 Token 发送心跳 ➔ 200
    let hb_body = serde_json::json!({
        "node_id": "node-pending-01",
        "active_slots": 1,
        "cpu_usage": 10.0,
        "memory_usage": 20.0
    });
    let res = app
        .oneshot(
            Request::builder()
                .method("POST")
                .uri("/api/node/heartbeat")
                .header("authorization", format!("Bearer {}", node_token))
                .header("content-type", "application/json")
                .body(Body::from(serde_json::to_vec(&hb_body).unwrap()))
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::OK);
}

#[tokio::test]
async fn test_admin_sessions_capacity_limit() {
    let temp_dir = tempfile::tempdir().unwrap();
    let db_path = temp_dir.path().join("sess_db.db");
    let queue_db_path = temp_dir.path().join("sess_queue.db");
    let results_dir = temp_dir.path().join("results");
    std::fs::create_dir_all(&results_dir).unwrap();

    let db = Database::new(&db_path.to_string_lossy()).await.unwrap();
    let queue = Arc::new(
        SqliteTaskQueue::new(&queue_db_path.to_string_lossy())
            .await
            .unwrap(),
    );
    let scheduler = Arc::new(GridScheduler::new(
        db.clone(),
        queue.clone(),
        results_dir.to_string_lossy().to_string(),
    ));

    let admin_pass = "admin_capacity_secret";
    let state = AppState {
        db,
        queue,
        scheduler,
        results_dir: results_dir.to_string_lossy().to_string(),
        rate_limiter: server::api::rate_limit::RateLimiter::new(
            1000,
            std::time::Duration::from_secs(300),
        ),
        auth_token: Some(admin_pass.to_string()),
        admin_token: Some(admin_pass.to_string()),
        auth_disabled: false,
        admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
            std::collections::HashMap::new(),
        )),
    };

    let api_router =
        axum::Router::new().route("/login", axum::routing::post(server::api::auth::login));
    let app = axum::Router::new()
        .nest("/api", api_router)
        .with_state(state.clone());

    // 连续登录 105 次，超过 100 容量上限
    for _ in 0..105 {
        let right_body = serde_json::json!({ "password": admin_pass });
        let res = app
            .clone()
            .oneshot(
                Request::builder()
                    .method("POST")
                    .uri("/api/login")
                    .header("content-type", "application/json")
                    .extension(axum::extract::ConnectInfo(std::net::SocketAddr::from((
                        [127, 0, 0, 1],
                        12345,
                    ))))
                    .body(Body::from(serde_json::to_vec(&right_body).unwrap()))
                    .unwrap(),
            )
            .await
            .unwrap();
        assert_eq!(res.status(), StatusCode::OK);
    }

    // 验证 sessions 集合保存的数量不超过 MAX_ADMIN_SESSIONS (100)
    let sessions = state.admin_sessions.read().await;
    assert!(sessions.len() <= server::api::MAX_ADMIN_SESSIONS);
    assert_eq!(sessions.len(), 100);
}

#[tokio::test]
async fn test_node_register_rate_limit() {
    let temp_dir = tempfile::tempdir().unwrap();
    let db_path = temp_dir.path().join("reg_limit_db.db");
    let queue_db_path = temp_dir.path().join("reg_limit_queue.db");
    let results_dir = temp_dir.path().join("results");
    std::fs::create_dir_all(&results_dir).unwrap();

    let db = Database::new(&db_path.to_string_lossy()).await.unwrap();
    let queue = Arc::new(
        SqliteTaskQueue::new(&queue_db_path.to_string_lossy())
            .await
            .unwrap(),
    );
    let scheduler = Arc::new(GridScheduler::new(
        db.clone(),
        queue.clone(),
        results_dir.to_string_lossy().to_string(),
    ));

    let state = AppState {
        db,
        queue,
        scheduler,
        results_dir: results_dir.to_string_lossy().to_string(),
        rate_limiter: server::api::rate_limit::RateLimiter::new(
            5,
            std::time::Duration::from_secs(300),
        ),
        auth_token: None,
        admin_token: None,
        auth_disabled: false,
        admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
            std::collections::HashMap::new(),
        )),
    };

    // 注册端点专用限流器（count_all=true）：对成功请求也计数，模拟生产 main.rs 的注册节流配置。
    let register_limiter =
        server::api::rate_limit::RateLimiter::new_count_all(5, std::time::Duration::from_secs(60));
    let register_rate_limit_layer = axum::middleware::from_fn_with_state(
        register_limiter,
        server::api::rate_limit::rate_limit_middleware,
    );

    let app = axum::Router::new()
        .route(
            "/api/node/register",
            axum::routing::post(server::api::node::register_node).layer(register_rate_limit_layer),
        )
        .with_state(state);

    let reg_body = serde_json::json!({
        "node_id": "test-limit-node",
        "host_name": "limit-host",
        "max_slots": 4
    });

    // 5 次以内的注册尝试 ➔ 200
    for _ in 0..5 {
        let res = app
            .clone()
            .oneshot(
                Request::builder()
                    .method("POST")
                    .uri("/api/node/register")
                    .header("content-type", "application/json")
                    .extension(axum::extract::ConnectInfo(std::net::SocketAddr::from((
                        [127, 0, 0, 1],
                        12345,
                    ))))
                    .body(Body::from(serde_json::to_vec(&reg_body).unwrap()))
                    .unwrap(),
            )
            .await
            .unwrap();
        assert_eq!(res.status(), StatusCode::OK);
    }

    // 第 6 次触发限流 ➔ 429 TOO_MANY_REQUESTS
    let res = app
        .oneshot(
            Request::builder()
                .method("POST")
                .uri("/api/node/register")
                .header("content-type", "application/json")
                .extension(axum::extract::ConnectInfo(std::net::SocketAddr::from((
                    [127, 0, 0, 1],
                    12345,
                ))))
                .body(Body::from(serde_json::to_vec(&reg_body).unwrap()))
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::TOO_MANY_REQUESTS);
}

/// 回归保护：通用 API 限流器（count_all=false）不应因 /node/register 的成功响应计数。
/// 历史缺陷：此前中间件对 is_register 无条件计数，导致通用限流器复用时成功注册会把 IP
/// 锁出整个 /api/*（跨端点连锁）。此测试断言多次成功注册后通用限流器仍放行。
#[tokio::test]
async fn test_general_limiter_does_not_count_successful_register() {
    let temp_dir = tempfile::tempdir().unwrap();
    let db_path = temp_dir.path().join("gen_limit_db.db");
    let queue_db_path = temp_dir.path().join("gen_limit_queue.db");
    let results_dir = temp_dir.path().join("results");
    std::fs::create_dir_all(&results_dir).unwrap();

    let db = Database::new(&db_path.to_string_lossy()).await.unwrap();
    let queue = Arc::new(
        SqliteTaskQueue::new(&queue_db_path.to_string_lossy())
            .await
            .unwrap(),
    );
    let scheduler = Arc::new(GridScheduler::new(
        db.clone(),
        queue.clone(),
        results_dir.to_string_lossy().to_string(),
    ));

    let state = AppState {
        db,
        queue,
        scheduler,
        results_dir: results_dir.to_string_lossy().to_string(),
        rate_limiter: server::api::rate_limit::RateLimiter::new(
            5,
            std::time::Duration::from_secs(300),
        ),
        auth_token: None,
        admin_token: None,
        auth_disabled: false,
        admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
            std::collections::HashMap::new(),
        )),
    };

    // 通用限流器（new = count_all=false），阈值仅 3，远低于下面的请求数。
    // 若仍对成功注册计数，第 4 次就会 429。
    let general_limiter =
        server::api::rate_limit::RateLimiter::new(3, std::time::Duration::from_secs(60));
    let rate_limit_layer = axum::middleware::from_fn_with_state(
        general_limiter,
        server::api::rate_limit::rate_limit_middleware,
    );

    let app = axum::Router::new()
        .route(
            "/api/node/register",
            axum::routing::post(server::api::node::register_node),
        )
        .layer(rate_limit_layer)
        .with_state(state);

    let reg_body = serde_json::json!({
        "node_id": "gen-limit-node",
        "host_name": "gen-host",
        "max_slots": 4
    });

    // 连续 6 次成功注册（远超阈值 3）：通用限流器不应计数成功响应，全部应为 200。
    for i in 0..6 {
        // 每次用不同 node_id 避免重复注册逻辑干扰
        let body = serde_json::json!({
            "node_id": format!("gen-limit-node-{}", i),
            "host_name": "gen-host",
            "max_slots": 4
        });
        let res = app
            .clone()
            .oneshot(
                Request::builder()
                    .method("POST")
                    .uri("/api/node/register")
                    .header("content-type", "application/json")
                    .extension(axum::extract::ConnectInfo(std::net::SocketAddr::from((
                        [127, 0, 0, 1],
                        12345,
                    ))))
                    .body(Body::from(serde_json::to_vec(&body).unwrap()))
                    .unwrap(),
            )
            .await
            .unwrap();
        assert_eq!(
            res.status(),
            StatusCode::OK,
            "第 {} 次注册应成功，通用限流器不应计数成功响应",
            i + 1
        );
    }

    // 静默 reg_body 未使用的警告（保留以对齐其它测试结构）
    let _ = &reg_body;
}

#[tokio::test]
async fn test_cors_same_origin_and_local_policy() {
    let temp_dir = tempfile::tempdir().unwrap();
    let db_path = temp_dir.path().join("cors_db.db");
    let queue_db_path = temp_dir.path().join("cors_queue.db");
    let results_dir = temp_dir.path().join("results");
    std::fs::create_dir_all(&results_dir).unwrap();

    let db = Database::new(&db_path.to_string_lossy()).await.unwrap();
    let queue = Arc::new(
        SqliteTaskQueue::new(&queue_db_path.to_string_lossy())
            .await
            .unwrap(),
    );
    let scheduler = Arc::new(GridScheduler::new(
        db.clone(),
        queue.clone(),
        results_dir.to_string_lossy().to_string(),
    ));

    let state = AppState {
        db,
        queue,
        scheduler,
        results_dir: results_dir.to_string_lossy().to_string(),
        rate_limiter: server::api::rate_limit::RateLimiter::new(
            100,
            std::time::Duration::from_secs(300),
        ),
        auth_token: None,
        admin_token: None,
        auth_disabled: false,
        admin_sessions: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
    };

    let app = axum::Router::new()
        .route(
            "/api/status",
            axum::routing::get(server::api::status::get_status),
        )
        .layer(server::cors::build_cors_layer())
        .with_state(state);

    // 1. 本地 localhost 请求 -> 允许 CORS
    let res = app
        .clone()
        .oneshot(
            Request::builder()
                .uri("/api/status")
                .header("origin", "http://localhost:5173")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::OK);
    assert_eq!(
        res.headers()
            .get("access-control-allow-origin")
            .unwrap()
            .to_str()
            .unwrap(),
        "http://localhost:5173"
    );

    // 2. 本地 127.0.0.1 请求 -> 允许 CORS
    let res = app
        .clone()
        .oneshot(
            Request::builder()
                .uri("/api/status")
                .header("origin", "http://127.0.0.1:3000")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::OK);
    assert_eq!(
        res.headers()
            .get("access-control-allow-origin")
            .unwrap()
            .to_str()
            .unwrap(),
        "http://127.0.0.1:3000"
    );

    // 3. 同源请求 (Origin 匹配 Host 标头) -> 允许 CORS
    let res = app
        .clone()
        .oneshot(
            Request::builder()
                .uri("/api/status")
                .header("host", "192.168.1.100:8090")
                .header("origin", "http://192.168.1.100:8090")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::OK);
    assert_eq!(
        res.headers()
            .get("access-control-allow-origin")
            .unwrap()
            .to_str()
            .unwrap(),
        "http://192.168.1.100:8090"
    );

    // 4. 外部非法跨源请求 -> 拒绝 CORS
    let res = app
        .oneshot(
            Request::builder()
                .uri("/api/status")
                .header("host", "192.168.1.100:8090")
                .header("origin", "https://attacker.example.com")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(res.status(), StatusCode::OK);
    assert!(res.headers().get("access-control-allow-origin").is_none());
}