mod common;

use std::sync::Arc;
use std::time::Duration;

use bytes::Bytes;

use fafb_q::codec::Frame;
use fafb_q::codec::frame::{CAP_GOING_AWAY, ErrorCode, GoingAwayReason, SERVER_CAPABILITIES};
use fafb_q::config::TcpConfig;
use fafb_q::server::state::{ServerLifecycle, ServerState};
use fafb_q::types::{ApiKey, NamespaceId, Permission, TopicName};

use common::{
    TEST_KEY, TEST_NS, TestClient, start_test_tcp_server, start_test_tcp_server_with_perms,
};

// -- Capability negotiation tests --

#[tokio::test]
async fn v1_client_without_caps_connects_successfully() {
    let server = start_test_tcp_server().await;
    let mut client = TestClient::connect(&server.addr).await;

    // Send Connect without capabilities (legacy v1 client — capabilities = 0)
    client
        .send(Frame::Connect {
            version: 1,
            namespace: NamespaceId(TEST_NS.into()),
            api_key: ApiKey(TEST_KEY.into()),
            capabilities: 0,
        })
        .await;

    let frame = client.recv().await;
    match frame {
        Frame::ConnectAck {
            success,
            server_capabilities,
            ..
        } => {
            assert!(success);
            // Server returns 0 because the intersection of 0 & SERVER_CAPABILITIES = 0
            assert_eq!(server_capabilities, 0);
        }
        other => panic!("expected ConnectAck, got {other:?}"),
    }

    server.shutdown();
}

#[tokio::test]
async fn capability_aware_client_gets_agreed_caps() {
    let server = start_test_tcp_server().await;
    let mut client = TestClient::connect(&server.addr).await;

    // Send Connect with CAP_GOING_AWAY
    client
        .send(Frame::Connect {
            version: 1,
            namespace: NamespaceId(TEST_NS.into()),
            api_key: ApiKey(TEST_KEY.into()),
            capabilities: CAP_GOING_AWAY,
        })
        .await;

    let frame = client.recv().await;
    match frame {
        Frame::ConnectAck {
            success,
            server_capabilities,
            ..
        } => {
            assert!(success);
            assert_eq!(server_capabilities, CAP_GOING_AWAY);
        }
        other => panic!("expected ConnectAck, got {other:?}"),
    }

    server.shutdown();
}

#[tokio::test]
async fn capability_negotiation_full_server_caps() {
    let server = start_test_tcp_server().await;
    let mut client = TestClient::connect(&server.addr).await;

    // Client requests all possible capabilities
    client
        .send(Frame::Connect {
            version: 1,
            namespace: NamespaceId(TEST_NS.into()),
            api_key: ApiKey(TEST_KEY.into()),
            capabilities: 0xFFFFFFFF, // all bits set
        })
        .await;

    let frame = client.recv().await;
    match frame {
        Frame::ConnectAck {
            success,
            server_capabilities,
            ..
        } => {
            assert!(success);
            // Server only agrees to what it supports
            assert_eq!(server_capabilities, SERVER_CAPABILITIES);
        }
        other => panic!("expected ConnectAck, got {other:?}"),
    }

    server.shutdown();
}

// -- ServerLifecycle unit-level tests for drain wait early exit --

#[tokio::test]
async fn drain_wait_exits_early_when_connections_reach_zero() {
    let (lifecycle, _rx) = ServerLifecycle::new();
    let lifecycle = Arc::new(lifecycle);
    lifecycle.set_initialized();

    // Simulate 2 active connections
    lifecycle.inc_connections();
    lifecycle.inc_connections();

    lifecycle.transition_to(ServerState::Draining);
    assert_eq!(lifecycle.active_connections(), 2);

    // Simulate connections closing
    lifecycle.dec_connections();
    lifecycle.dec_connections();
    assert_eq!(lifecycle.active_connections(), 0);

    // The drain wait loop in main.rs checks active_connections() == 0 and exits early.
    // We verify the condition is met.
    assert_eq!(lifecycle.active_connections(), 0);
    assert!(lifecycle.is_draining());
}

// -- Drain-aware connection handler tests --

#[tokio::test]
async fn drain_sends_server_draining_error() {
    let server = start_test_tcp_server().await;
    let mut client = TestClient::connect(&server.addr).await;

    // Connect successfully
    client
        .send(Frame::Connect {
            version: 1,
            namespace: NamespaceId(TEST_NS.into()),
            api_key: ApiKey(TEST_KEY.into()),
            capabilities: 0,
        })
        .await;
    let frame = client.recv().await;
    assert!(matches!(frame, Frame::ConnectAck { success: true, .. }));

    // Trigger drain — the connection handler should send ServerDraining
    server.lifecycle.transition_to(ServerState::Draining);

    // The connection should receive a ServerDraining error and close
    let frame = client.recv().await;
    match frame {
        Frame::Error { code, .. } => {
            assert_eq!(code, ErrorCode::ServerDraining);
        }
        other => panic!("expected Error frame with ServerDraining, got {other:?}"),
    }
}

#[tokio::test]
async fn subscribe_rejected_during_drain() {
    let server = start_test_tcp_server_with_perms(&[
        Permission::Subscribe,
        Permission::Publish,
        Permission::Query,
    ])
    .await;

    let mut client = TestClient::connect(&server.addr).await;

    // Connect
    client
        .send(Frame::Connect {
            version: 1,
            namespace: NamespaceId(TEST_NS.into()),
            api_key: ApiKey(TEST_KEY.into()),
            capabilities: 0,
        })
        .await;
    let _ = client.recv().await; // ConnectAck

    // Publish a message first (publishes should work during drain)
    client
        .send(Frame::Publish {
            topic: TopicName("events".into()),
            payload: Bytes::from_static(b"hello"),
            headers: vec![],
            ttl_ms: 0,
        })
        .await;
    let _ = client.recv().await; // PublishAck

    // Trigger drain
    server.lifecycle.transition_to(ServerState::Draining);

    // Give the connection handler a moment to process the lifecycle change
    tokio::time::sleep(Duration::from_millis(50)).await;

    // Try to subscribe — should get either ServerDraining rejection or the
    // connection-level drain error. The drain notification races with our
    // subscribe frame.
    client
        .send(Frame::Subscribe {
            topic: TopicName("events".into()),
            group: None,
        })
        .await;

    // We should receive an error (either the per-frame rejection or the
    // connection-level drain signal).
    let frame = client.recv().await;
    match frame {
        Frame::Error { code, .. } => {
            assert!(
                code == ErrorCode::ServerDraining,
                "expected ServerDraining, got {code:?}"
            );
        }
        other => panic!("expected Error frame, got {other:?}"),
    }
}

// -- GoingAway frame tests --

#[tokio::test]
async fn capable_client_receives_going_away_on_drain() {
    let server = start_test_tcp_server().await;
    let mut client = TestClient::connect(&server.addr).await;

    // Connect with CAP_GOING_AWAY
    client
        .send(Frame::Connect {
            version: 1,
            namespace: NamespaceId(TEST_NS.into()),
            api_key: ApiKey(TEST_KEY.into()),
            capabilities: CAP_GOING_AWAY,
        })
        .await;
    let frame = client.recv().await;
    assert!(matches!(
        frame,
        Frame::ConnectAck {
            success: true,
            server_capabilities,
            ..
        } if server_capabilities & CAP_GOING_AWAY != 0
    ));

    // Trigger drain
    server.lifecycle.transition_to(ServerState::Draining);

    // Capable client should receive GoingAway (not Error)
    let frame = client.recv().await;
    match frame {
        Frame::GoingAway {
            reason,
            deadline_ms,
            message,
        } => {
            assert_eq!(reason, GoingAwayReason::UPGRADE);
            assert!(deadline_ms > 0, "deadline_ms should be positive");
            assert!(!message.is_empty());
        }
        other => panic!("expected GoingAway frame, got {other:?}"),
    }

    // Connection should still be open — client can continue working.
    // Transition to ShuttingDown to force-close.
    server.lifecycle.transition_to(ServerState::ShuttingDown);
}

#[tokio::test]
async fn no_cap_client_gets_error_not_going_away_on_drain() {
    let server = start_test_tcp_server().await;
    let mut client = TestClient::connect(&server.addr).await;

    // Connect WITHOUT CAP_GOING_AWAY
    client
        .send(Frame::Connect {
            version: 1,
            namespace: NamespaceId(TEST_NS.into()),
            api_key: ApiKey(TEST_KEY.into()),
            capabilities: 0,
        })
        .await;
    let _ = client.recv().await; // ConnectAck

    // Trigger drain
    server.lifecycle.transition_to(ServerState::Draining);

    // Non-capable client gets Error with ServerDraining and connection closes
    let frame = client.recv().await;
    match frame {
        Frame::Error { code, .. } => {
            assert_eq!(code, ErrorCode::ServerDraining);
        }
        other => panic!("expected Error frame with ServerDraining, got {other:?}"),
    }
}

#[tokio::test]
async fn subscribe_rejected_during_drain_for_capable_client() {
    let server = start_test_tcp_server_with_perms(&[
        Permission::Subscribe,
        Permission::Publish,
        Permission::Query,
    ])
    .await;

    let mut client = TestClient::connect(&server.addr).await;

    // Connect with CAP_GOING_AWAY
    client
        .send(Frame::Connect {
            version: 1,
            namespace: NamespaceId(TEST_NS.into()),
            api_key: ApiKey(TEST_KEY.into()),
            capabilities: CAP_GOING_AWAY,
        })
        .await;
    let _ = client.recv().await; // ConnectAck

    // Trigger drain
    server.lifecycle.transition_to(ServerState::Draining);

    // Wait for GoingAway to arrive
    let frame = client.recv().await;
    assert!(
        matches!(frame, Frame::GoingAway { .. }),
        "expected GoingAway, got {frame:?}"
    );

    // Try to subscribe — should get ServerDraining rejection
    client
        .send(Frame::Subscribe {
            topic: TopicName("events".into()),
            group: None,
        })
        .await;

    let frame = client.recv().await;
    match frame {
        Frame::Error { code, .. } => {
            assert_eq!(code, ErrorCode::ServerDraining);
        }
        other => panic!("expected ServerDraining error, got {other:?}"),
    }

    server.lifecycle.transition_to(ServerState::ShuttingDown);
}

#[tokio::test]
async fn publish_accepted_during_drain_for_capable_client() {
    let server = start_test_tcp_server_with_perms(&[
        Permission::Subscribe,
        Permission::Publish,
        Permission::Query,
    ])
    .await;

    let mut client = TestClient::connect(&server.addr).await;

    // Connect with CAP_GOING_AWAY
    client
        .send(Frame::Connect {
            version: 1,
            namespace: NamespaceId(TEST_NS.into()),
            api_key: ApiKey(TEST_KEY.into()),
            capabilities: CAP_GOING_AWAY,
        })
        .await;
    let _ = client.recv().await; // ConnectAck

    // Trigger drain
    server.lifecycle.transition_to(ServerState::Draining);

    // Wait for GoingAway
    let frame = client.recv().await;
    assert!(
        matches!(frame, Frame::GoingAway { .. }),
        "expected GoingAway, got {frame:?}"
    );

    // Publish during drain — should succeed
    client
        .send(Frame::Publish {
            topic: TopicName("events".into()),
            payload: Bytes::from_static(b"drain-publish"),
            headers: vec![],
            ttl_ms: 0,
        })
        .await;

    let frame = client.recv().await;
    assert!(
        matches!(frame, Frame::PublishAck { .. }),
        "expected PublishAck during drain, got {frame:?}"
    );

    server.lifecycle.transition_to(ServerState::ShuttingDown);
}

// -- Health endpoint tests --

#[tokio::test]
async fn health_live_always_200() {
    let server = common::start_test_http_server().await;

    let resp = reqwest::get(format!("{}/health/live", server.base_url))
        .await
        .unwrap();
    assert_eq!(resp.status(), 200);

    let body: serde_json::Value = resp.json().await.unwrap();
    assert_eq!(body["status"], "alive");

    server.shutdown();
}

#[tokio::test]
async fn health_ready_200_when_initialized() {
    let server = common::start_test_http_server().await;

    let resp = reqwest::get(format!("{}/health/ready", server.base_url))
        .await
        .unwrap();
    assert_eq!(resp.status(), 200);

    let body: serde_json::Value = resp.json().await.unwrap();
    assert_eq!(body["status"], "ready");

    server.shutdown();
}

#[tokio::test]
async fn health_ready_503_when_draining() {
    let server = common::start_test_http_server().await;

    // Transition to draining
    server.lifecycle.transition_to(ServerState::Draining);

    let resp = reqwest::get(format!("{}/health/ready", server.base_url))
        .await
        .unwrap();
    assert_eq!(resp.status(), 503);

    let body: serde_json::Value = resp.json().await.unwrap();
    assert_eq!(body["status"], "draining");

    server.shutdown();
}

#[tokio::test]
async fn health_alias_matches_ready() {
    let server = common::start_test_http_server().await;

    let resp = reqwest::get(format!("{}/health", server.base_url))
        .await
        .unwrap();
    assert_eq!(resp.status(), 200);

    let body: serde_json::Value = resp.json().await.unwrap();
    assert_eq!(body["status"], "ready");

    server.shutdown();
}

// -- Init-gate middleware test --

#[tokio::test]
async fn business_routes_503_before_init() {
    // Build an HTTP server manually without calling set_initialized()
    use fafb_q::auth::{MockAuthenticator, PermissionSet};
    use fafb_q::namespace::{DefaultNamespaceManager, ManagerConfig};
    use fafb_q::queue::{EngineConfig, InMemoryQueueEngine};
    use fafb_q::server::http::HttpServer;
    use fafb_q::storage::mock::MockStorage;

    let storage = Arc::new(MockStorage::new());
    let engine = Arc::new(InMemoryQueueEngine::new(
        Arc::clone(&storage),
        EngineConfig::default(),
    ));
    let manager = Arc::new(DefaultNamespaceManager::new(
        Arc::clone(&engine),
        vec![NamespaceId(TEST_NS.into())],
        common::test_shared_config(),
        ManagerConfig::default(),
    ));
    let auth: Arc<dyn fafb_q::auth::Authenticator> = Arc::new(MockAuthenticator {
        namespace: NamespaceId(TEST_NS.into()),
        key: ApiKey(TEST_KEY.into()),
        permissions: PermissionSet::from_slice(&[
            Permission::Subscribe,
            Permission::Publish,
            Permission::Query,
        ]),
    });
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let base_url = format!("http://{addr}");

    let dummy_reload_handle = {
        use tracing_subscriber::layer::SubscriberExt;
        let (_, handle) =
            tracing_subscriber::reload::Layer::new(tracing_subscriber::EnvFilter::new("info"));
        Arc::new(handle)
    };

    let (lifecycle, _rx) = ServerLifecycle::new();
    let lifecycle = Arc::new(lifecycle);
    // NOT calling lifecycle.set_initialized()

    let server = HttpServer::new(
        fafb_q::config::HttpConfig {
            bind_address: "127.0.0.1".into(),
            port: addr.port(),
            ..Default::default()
        },
        fafb_q::server::http::AppState {
            manager,
            auth,
            message_id_counter: Arc::new(std::sync::atomic::AtomicU64::new(1)),
            prom_handle: None,
            log_reload_handle: dummy_reload_handle,
            lifecycle: Arc::clone(&lifecycle),
            admin_api_key: None,
            drain_trigger: None,
            deploy_config: fafb_q::config::DeployConfig::default(),
            shared_config: fafb_q::config::shared_config(&fafb_q::config::ServerConfig::default()),
            config_path: None,
            reload_state: None,
            grace_pending: std::sync::Arc::new(tokio::sync::RwLock::new(
                std::collections::HashMap::new(),
            )),
        },
        None, // no TLS
    );

    tokio::spawn(async move {
        let _ = server.run_with_listener(listener).await;
    });
    tokio::time::sleep(Duration::from_millis(50)).await;

    // /health/live should still work (not gated)
    let resp = reqwest::get(format!("{base_url}/health/live"))
        .await
        .unwrap();
    assert_eq!(resp.status(), 200);

    // /health/ready should return 503 "initializing"
    let resp = reqwest::get(format!("{base_url}/health/ready"))
        .await
        .unwrap();
    assert_eq!(resp.status(), 503);
    let body: serde_json::Value = resp.json().await.unwrap();
    assert_eq!(body["status"], "initializing");

    // Business route should return 503
    let client = reqwest::Client::new();
    let resp = client
        .get(format!("{base_url}/v1/{TEST_NS}/topics"))
        .header("Authorization", format!("Bearer {TEST_KEY}"))
        .send()
        .await
        .unwrap();
    assert_eq!(resp.status(), 503);

    // Now initialize
    lifecycle.set_initialized();

    // Business route should now work
    let resp = client
        .get(format!("{base_url}/v1/{TEST_NS}/topics"))
        .header("Authorization", format!("Bearer {TEST_KEY}"))
        .send()
        .await
        .unwrap();
    assert_eq!(resp.status(), 200);

    lifecycle.transition_to(ServerState::ShuttingDown);
}
