use anthropic_async::AnthropicConfig;
use anthropic_async::Client;
use anthropic_async::types::content::*;
use anthropic_async::types::messages::*;
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;

#[tokio::test]
async fn test_retry_on_429_then_success() {
    let server = MockServer::start().await;
    let count = Arc::new(AtomicUsize::new(0));

    let count_clone = Arc::clone(&count);
    Mock::given(method("POST"))
        .and(path("/v1/messages"))
        .respond_with(move |_req: &wiremock::Request| {
            let i = count_clone.fetch_add(1, Ordering::SeqCst);
            if i == 0 {
                ResponseTemplate::new(429)
                    .insert_header("retry-after-ms", "100")
                    .set_body_json(serde_json::json!({
                        "error": {
                            "message": "Rate limit exceeded",
                            "type": "rate_limit_error"
                        }
                    }))
            } else {
                ResponseTemplate::new(200).set_body_json(serde_json::json!({
                    "id": "msg",
                    "type": "message",
                    "role": "assistant",
                    "content": [{"type": "text", "text": "Success"}],
                    "model": "claude"
                }))
            }
        })
        .mount(&server)
        .await;

    let cfg = AnthropicConfig::new()
        .with_api_base(server.uri())
        .with_api_key("test");
    let client = Client::with_config(cfg);

    let request = MessagesCreateRequest {
        model: "claude".into(),
        max_tokens: 10,
        messages: vec![MessageParam {
            role: MessageRole::User,
            content: "test".into(),
        }],
        ..Default::default()
    };

    let response = client.messages().create(request).await.unwrap();
    assert_eq!(response.kind, "message");
    assert!(count.load(Ordering::SeqCst) >= 2);
}

#[tokio::test]
async fn test_529_overloaded_retry() {
    let server = MockServer::start().await;
    let count = Arc::new(AtomicUsize::new(0));

    let count_clone = Arc::clone(&count);
    Mock::given(method("POST"))
        .and(path("/v1/messages"))
        .respond_with(move |_req: &wiremock::Request| {
            let i = count_clone.fetch_add(1, Ordering::SeqCst);
            if i == 0 {
                ResponseTemplate::new(529).set_body_json(serde_json::json!({
                    "error": {
                        "message": "Overloaded",
                        "type": "overloaded_error"
                    }
                }))
            } else {
                ResponseTemplate::new(200).set_body_json(serde_json::json!({
                    "id": "msg",
                    "type": "message",
                    "role": "assistant",
                    "content": [{"type": "text", "text": "Success"}],
                    "model": "claude"
                }))
            }
        })
        .mount(&server)
        .await;

    let cfg = AnthropicConfig::new()
        .with_api_base(server.uri())
        .with_api_key("test");
    let client = Client::with_config(cfg);

    let request = MessagesCreateRequest {
        model: "claude".into(),
        max_tokens: 10,
        messages: vec![MessageParam {
            role: MessageRole::User,
            content: "test".into(),
        }],
        ..Default::default()
    };

    let response = client.messages().create(request).await.unwrap();
    assert_eq!(response.kind, "message");
}

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

    Mock::given(method("POST"))
        .and(path("/v1/messages"))
        .respond_with(ResponseTemplate::new(400).set_body_json(serde_json::json!({
            "error": {
                "message": "Invalid request",
                "type": "invalid_request_error"
            }
        })))
        .mount(&server)
        .await;

    let cfg = AnthropicConfig::new()
        .with_api_base(server.uri())
        .with_api_key("test");
    let client = Client::with_config(cfg);

    let request = MessagesCreateRequest {
        model: "invalid-model".into(),
        max_tokens: 10,
        messages: vec![],
        ..Default::default()
    };

    let err = client.messages().create(request).await.unwrap_err();
    match err {
        anthropic_async::AnthropicError::Api(obj) => {
            assert_eq!(obj.message, "Invalid request");
        }
        _ => panic!("Expected Api error"),
    }
}
