//! Integration tests for the proxy engine.
//! Uses wiremock to simulate upstream LLM providers.

use axum::http::StatusCode;
use std::sync::Arc;
use tokio::sync::mpsc;
use wiremock::{Mock, MockServer, ResponseTemplate};
use wiremock::matchers::{method, path, header};

use dd0c_route::{
    AppConfig, TelemetryEvent, RouterBrain,
    proxy::{create_router, ProxyState},
};

// --- Test Auth Provider (always succeeds) ---

struct TestAuthProvider {
    org_id: String,
}

#[async_trait::async_trait]
impl dd0c_route::AuthProvider for TestAuthProvider {
    async fn authenticate(
        &self,
        _headers: &axum::http::HeaderMap,
    ) -> Result<dd0c_route::AuthContext, dd0c_route::AuthError> {
        Ok(dd0c_route::AuthContext {
            org_id: self.org_id.clone(),
            user_id: None,
            role: dd0c_route::Role::Member,
        })
    }
}

// --- Test Helpers ---

async fn setup_proxy(mock_url: &str) -> (axum::Router, mpsc::Receiver<TelemetryEvent>) {
    let (tx, rx) = mpsc::channel::<TelemetryEvent>(100);

    let mut providers = std::collections::HashMap::new();
    providers.insert("openai".to_string(), dd0c_route::config::ProviderConfig {
        api_key: "test-key".to_string(),
        base_url: mock_url.to_string(),
    });

    let config = Arc::new(AppConfig {
        proxy_port: 0,
        api_port: 0,
        database_url: String::new(),
        redis_url: String::new(),
        timescale_url: String::new(),
        jwt_secret: "test".to_string(),
        auth_mode: dd0c_route::config::AuthMode::Local,
        governance_mode: dd0c_route::config::GovernanceMode::Audit,
        providers,
        telemetry_channel_size: 100,
    });

    let state = Arc::new(ProxyState {
        auth: Arc::new(TestAuthProvider { org_id: "test-org".to_string() }),
        router: Arc::new(RouterBrain::new()),
        telemetry_tx: tx,
        http_client: reqwest::Client::new(),
        config,
    });

    (create_router(state), rx)
}

fn chat_request(model: &str, message: &str) -> String {
    serde_json::json!({
        "model": model,
        "messages": [{"role": "user", "content": message}]
    }).to_string()
}

fn chat_response(content: &str, prompt_tokens: u32, completion_tokens: u32) -> String {
    serde_json::json!({
        "id": "chatcmpl-test",
        "object": "chat.completion",
        "model": "gpt-4o-mini",
        "choices": [{
            "index": 0,
            "message": {"role": "assistant", "content": content},
            "finish_reason": "stop"
        }],
        "usage": {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": prompt_tokens + completion_tokens
        }
    }).to_string()
}

// --- Tests ---

#[tokio::test]
async fn proxy_forwards_request_to_upstream() {
    let mock = MockServer::start().await;

    Mock::given(method("POST"))
        .and(path("/v1/chat/completions"))
        .respond_with(
            ResponseTemplate::new(200)
                .set_body_string(chat_response("Hello!", 10, 5))
                .insert_header("content-type", "application/json"),
        )
        .expect(1)
        .mount(&mock)
        .await;

    let (app, _rx) = setup_proxy(&mock.uri()).await;

    let resp = axum::http::Request::builder()
        .method("POST")
        .uri("/v1/chat/completions")
        .header("content-type", "application/json")
        .header("authorization", "Bearer test-key")
        .body(axum::body::Body::from(chat_request("gpt-4o", "Hello")))
        .unwrap();

    let response = tower::ServiceExt::oneshot(app, resp).await.unwrap();
    assert_eq!(response.status(), StatusCode::OK);

    let body = axum::body::to_bytes(response.into_body(), 1024 * 1024).await.unwrap();
    let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
    assert_eq!(json["choices"][0]["message"]["content"], "Hello!");
}

#[tokio::test]
async fn proxy_emits_telemetry_event() {
    let mock = MockServer::start().await;

    Mock::given(method("POST"))
        .and(path("/v1/chat/completions"))
        .respond_with(
            ResponseTemplate::new(200)
                .set_body_string(chat_response("Hi", 10, 5))
                .insert_header("content-type", "application/json"),
        )
        .mount(&mock)
        .await;

    let (app, mut rx) = setup_proxy(&mock.uri()).await;

    let resp = axum::http::Request::builder()
        .method("POST")
        .uri("/v1/chat/completions")
        .header("content-type", "application/json")
        .header("authorization", "Bearer test-key")
        .body(axum::body::Body::from(chat_request("gpt-4o", "Hi")))
        .unwrap();

    tower::ServiceExt::oneshot(app, resp).await.unwrap();

    // Telemetry event should be emitted
    let event = tokio::time::timeout(
        std::time::Duration::from_secs(1),
        rx.recv(),
    ).await.unwrap().unwrap();

    assert_eq!(event.org_id, "test-org");
    assert_eq!(event.original_model, "gpt-4o");
    assert_eq!(event.status_code, 200);
    assert!(event.latency_ms < 5000); // Should be fast against local mock
}

#[tokio::test]
async fn proxy_routes_low_complexity_to_cheaper_model() {
    let mock = MockServer::start().await;

    // Expect the request to arrive with model=gpt-4o-mini (routed)
    Mock::given(method("POST"))
        .and(path("/v1/chat/completions"))
        .respond_with(
            ResponseTemplate::new(200)
                .set_body_string(chat_response("4", 5, 1))
                .insert_header("content-type", "application/json"),
        )
        .mount(&mock)
        .await;

    let (app, mut rx) = setup_proxy(&mock.uri()).await;

    let resp = axum::http::Request::builder()
        .method("POST")
        .uri("/v1/chat/completions")
        .header("content-type", "application/json")
        .header("authorization", "Bearer test-key")
        .body(axum::body::Body::from(chat_request("gpt-4o", "What is 2+2?")))
        .unwrap();

    let response = tower::ServiceExt::oneshot(app, resp).await.unwrap();
    assert_eq!(response.status(), StatusCode::OK);

    let event = rx.recv().await.unwrap();
    assert_eq!(event.original_model, "gpt-4o");
    assert_eq!(event.routed_model, "gpt-4o-mini");
    assert_eq!(event.strategy, "cheapest");
}

#[tokio::test]
async fn proxy_passes_through_upstream_errors() {
    let mock = MockServer::start().await;

    Mock::given(method("POST"))
        .and(path("/v1/chat/completions"))
        .respond_with(
            ResponseTemplate::new(429)
                .set_body_string(r#"{"error": {"message": "Rate limit exceeded"}}"#)
                .insert_header("content-type", "application/json"),
        )
        .mount(&mock)
        .await;

    let (app, _rx) = setup_proxy(&mock.uri()).await;

    let resp = axum::http::Request::builder()
        .method("POST")
        .uri("/v1/chat/completions")
        .header("content-type", "application/json")
        .header("authorization", "Bearer test-key")
        .body(axum::body::Body::from(chat_request("gpt-4o", "Hello")))
        .unwrap();

    let response = tower::ServiceExt::oneshot(app, resp).await.unwrap();
    // Upstream 429 should be passed through transparently
    assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
}

#[tokio::test]
async fn proxy_returns_bad_request_for_invalid_json() {
    let mock = MockServer::start().await;
    let (app, _rx) = setup_proxy(&mock.uri()).await;

    let resp = axum::http::Request::builder()
        .method("POST")
        .uri("/v1/chat/completions")
        .header("content-type", "application/json")
        .header("authorization", "Bearer test-key")
        .body(axum::body::Body::from("not json"))
        .unwrap();

    let response = tower::ServiceExt::oneshot(app, resp).await.unwrap();
    assert_eq!(response.status(), StatusCode::BAD_REQUEST);
}

#[tokio::test]
async fn health_endpoint_returns_ok() {
    let mock = MockServer::start().await;
    let (app, _rx) = setup_proxy(&mock.uri()).await;

    let resp = axum::http::Request::builder()
        .method("GET")
        .uri("/health")
        .body(axum::body::Body::empty())
        .unwrap();

    let response = tower::ServiceExt::oneshot(app, resp).await.unwrap();
    assert_eq!(response.status(), StatusCode::OK);
}
