use axum::Router;
use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
use narsil_engine::{
    BatchingConfig, BurnBackend, ObservabilityState, SharedInferenceService, http_router,
    http_router_with_observability, observability_router,
};
use reqwest::StatusCode;
use serde_json::{Value, json};

async fn spawn_router(app: Router) -> String {
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });
    format!("http://{addr}")
}

async fn spawn_gateway(service: SharedInferenceService<BurnBackend>) -> String {
    spawn_router(http_router(service)).await
}

async fn spawn_http_gateway() -> String {
    spawn_gateway(SharedInferenceService::new(BurnBackend::with_affine(
        2.0, 1.0,
    )))
    .await
}

fn f32_bytes(values: &[f32]) -> Vec<u8> {
    values
        .iter()
        .flat_map(|value| value.to_le_bytes())
        .collect()
}

fn f32_values(bytes: &[u8]) -> Vec<f32> {
    bytes
        .chunks_exact(4)
        .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap()))
        .collect()
}

fn json_request(values: &[f32]) -> Value {
    json!({
        "model_name": "burn-affine",
        "inputs": [{
            "name": "input",
            "dtype": "F32",
            "shape": [values.len()],
            "data": BASE64.encode(f32_bytes(values)),
        }],
        "parameters": {},
    })
}

#[tokio::test]
async fn http_gateway_serves_infer_and_infer_batch() {
    let endpoint = spawn_http_gateway().await;
    let client = reqwest::Client::new();

    let response = client
        .get(format!("{endpoint}/health/live"))
        .send()
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::NO_CONTENT);

    let response = client
        .get(format!("{endpoint}/health/ready"))
        .send()
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::NO_CONTENT);

    let response = client
        .post(format!("{endpoint}/v1/infer"))
        .json(&json_request(&[1.0, 2.0, 3.0]))
        .send()
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::OK);
    let response: Value = response.json().await.unwrap();
    let output = BASE64
        .decode(response["outputs"][0]["data"].as_str().unwrap())
        .unwrap();
    assert_eq!(f32_values(&output), vec![3.0, 5.0, 7.0]);
    assert_eq!(response["metadata"]["backend"], "burn");

    let response = client
        .post(format!("{endpoint}/v1/infer_batch"))
        .json(&json!({
            "requests": [
                json_request(&[4.0]),
                json_request(&[5.0, 6.0])
            ]
        }))
        .send()
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::OK);
    let response: Value = response.json().await.unwrap();
    assert_eq!(response["responses"][0]["status"], "ok");
    assert_eq!(response["responses"][1]["status"], "ok");

    let first = BASE64
        .decode(
            response["responses"][0]["response"]["outputs"][0]["data"]
                .as_str()
                .unwrap(),
        )
        .unwrap();
    let second = BASE64
        .decode(
            response["responses"][1]["response"]["outputs"][0]["data"]
                .as_str()
                .unwrap(),
        )
        .unwrap();
    assert_eq!(f32_values(&first), vec![9.0]);
    assert_eq!(f32_values(&second), vec![11.0, 13.0]);

    let response = client
        .get(format!("{endpoint}/metrics"))
        .send()
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::OK);
    let metrics = response.text().await.unwrap();
    assert!(metrics.contains("narsil_requests_total"));
    assert!(metrics.contains("narsil_request_latency_seconds"));
    assert!(metrics.contains("narsil_in_flight_requests"));
    assert!(metrics.contains("narsil_batch_size"));
    assert!(metrics.contains("transport=\"http\""));
}

#[tokio::test]
async fn http_gateway_rejects_oversized_body() {
    let endpoint = spawn_http_gateway().await;
    let client = reqwest::Client::new();

    // ~3 MiB body exceeds the 2 MiB DefaultBodyLimit.
    let huge = "A".repeat(3 * 1024 * 1024);
    let response = client
        .post(format!("{endpoint}/v1/infer"))
        .json(&json!({
            "model_name": "burn-affine",
            "inputs": [{
                "name": "input",
                "dtype": "F32",
                "shape": [1],
                "data": huge,
            }],
            "parameters": {},
        }))
        .send()
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
}

#[tokio::test]
async fn http_gateway_reports_per_item_decode_errors() {
    let endpoint = spawn_http_gateway().await;
    let client = reqwest::Client::new();

    // Mix a valid request with one carrying non-base64 tensor data.
    let mut bad = json_request(&[4.0]);
    bad["inputs"][0]["data"] = json!("not valid base64!!");

    let response = client
        .post(format!("{endpoint}/v1/infer_batch"))
        .json(&json!({ "requests": [json_request(&[4.0]), bad] }))
        .send()
        .await
        .unwrap();
    // The whole batch is NOT rejected; the bad item becomes a per-item error.
    assert_eq!(response.status(), StatusCode::OK);
    let response: Value = response.json().await.unwrap();
    assert_eq!(response["responses"][0]["status"], "ok");
    assert_eq!(response["responses"][1]["status"], "error");
    assert_eq!(response["responses"][1]["error"]["code"], "schema_error");
}

#[tokio::test]
async fn http_gateway_rejects_batch_exceeding_max_batch_size() {
    let service = SharedInferenceService::with_batching(
        BurnBackend::with_affine(2.0, 1.0),
        BatchingConfig::new(2, 5, 1, 8).unwrap(),
    )
    .unwrap();
    let endpoint = spawn_gateway(service).await;
    let client = reqwest::Client::new();

    // 3 requests > max_batch_size of 2 -> 400 up front.
    let response = client
        .post(format!("{endpoint}/v1/infer_batch"))
        .json(&json!({
            "requests": [json_request(&[1.0]), json_request(&[2.0]), json_request(&[3.0])]
        }))
        .send()
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::BAD_REQUEST);

    // A batch within the limit still succeeds through the collector path.
    let response = client
        .post(format!("{endpoint}/v1/infer_batch"))
        .json(&json!({ "requests": [json_request(&[1.0]), json_request(&[2.0])] }))
        .send()
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::OK);
    let response: Value = response.json().await.unwrap();
    assert_eq!(response["responses"][0]["status"], "ok");
    assert_eq!(response["responses"][1]["status"], "ok");
}

#[tokio::test]
async fn metrics_off_inference_gateway_but_on_observability_listener() {
    let service = SharedInferenceService::new(BurnBackend::with_affine(2.0, 1.0));
    let observability = ObservabilityState::new(&service).unwrap();

    // Default inference gateway must NOT expose /metrics (private-bind recommendation).
    let gateway = http_router_with_observability(service, observability.clone());
    let gateway_endpoint = spawn_router(gateway).await;

    // The dedicated observability listener is the canonical /metrics home.
    let obs_endpoint = spawn_router(observability_router(observability)).await;
    let client = reqwest::Client::new();

    let on_gateway = client
        .get(format!("{gateway_endpoint}/metrics"))
        .send()
        .await
        .unwrap();
    assert_eq!(on_gateway.status(), StatusCode::NOT_FOUND);

    // Health probes remain available on the inference gateway.
    let live = client
        .get(format!("{gateway_endpoint}/health/live"))
        .send()
        .await
        .unwrap();
    assert_eq!(live.status(), StatusCode::NO_CONTENT);

    // Drive one inference through the gateway so a request series exists, then confirm
    // it is visible on the dedicated observability listener (shared global recorder).
    let infer = client
        .post(format!("{gateway_endpoint}/v1/infer"))
        .json(&json_request(&[1.0]))
        .send()
        .await
        .unwrap();
    assert_eq!(infer.status(), StatusCode::OK);

    let on_obs = client
        .get(format!("{obs_endpoint}/metrics"))
        .send()
        .await
        .unwrap();
    assert_eq!(on_obs.status(), StatusCode::OK);
    assert!(
        on_obs
            .text()
            .await
            .unwrap()
            .contains("narsil_requests_total")
    );
}

#[tokio::test]
async fn readiness_reflects_backend_and_draining() {
    let service = SharedInferenceService::new(BurnBackend::with_affine(2.0, 1.0));
    let observability = ObservabilityState::new(&service).unwrap();
    let endpoint = spawn_router(observability_router(observability.clone())).await;
    let client = reqwest::Client::new();

    // Backend has valid metadata -> ready.
    let ready = client
        .get(format!("{endpoint}/health/ready"))
        .send()
        .await
        .unwrap();
    assert_eq!(ready.status(), StatusCode::NO_CONTENT);

    // Draining flips readiness to 503 but liveness stays 204.
    observability.set_draining(true);
    let ready = client
        .get(format!("{endpoint}/health/ready"))
        .send()
        .await
        .unwrap();
    assert_eq!(ready.status(), StatusCode::SERVICE_UNAVAILABLE);
    let live = client
        .get(format!("{endpoint}/health/live"))
        .send()
        .await
        .unwrap();
    assert_eq!(live.status(), StatusCode::NO_CONTENT);
}
