// Copyright 2024 RustFS Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::common::RustFSTestClusterEnvironment;
use aws_sdk_s3::Client;
use aws_sdk_s3::error::SdkError;
use aws_sdk_s3::types::{CorsConfiguration, CorsRule};
use bytes::Bytes;
use std::sync::Arc;
use tokio::sync::Barrier;
use tracing::{info, warn};

const BUCKET: &str = "conditional-put-race-bucket";
const BUCKET_METADATA_RELOAD_BUCKET: &str = "bucket-metadata-reload-barrier";

async fn cleanup_object(client: &Client, key: &str) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    client.delete_object().bucket(BUCKET).key(key).send().await?;
    Ok(())
}

async fn assert_bucket_cors_missing(client: &Client) {
    let result = client.get_bucket_cors().bucket(BUCKET_METADATA_RELOAD_BUCKET).send().await;
    match result {
        Err(SdkError::ServiceError(error)) => {
            assert_eq!(error.err().meta().code(), Some("NoSuchCORSConfiguration"));
        }
        result => panic!("expected the peer to report a missing CORS configuration: {result:?}"),
    }
}

async fn conditional_put(
    client: &Client,
    key: &str,
    data: &[u8],
    client_id: usize,
) -> Result<bool, Box<dyn std::error::Error + Send + Sync>> {
    let result = client
        .put_object()
        .bucket(BUCKET)
        .key(key)
        .body(Bytes::copy_from_slice(data).into())
        .if_none_match("*")
        .send()
        .await;

    match result {
        Ok(resp) => {
            info!("  Client {} SUCCEEDED - ETag: {}", client_id, resp.e_tag().unwrap_or("none"));
            Ok(true)
        }
        Err(SdkError::ServiceError(e)) => {
            let code = e.err().meta().code().unwrap_or("");
            let message = e.err().meta().message().unwrap_or("");
            if code == "PreconditionFailed" {
                warn!("  Client {} got 412 PreconditionFailed", client_id);
                Ok(false)
            } else {
                warn!("  Client {} got ServiceError: code={}, message={}", client_id, code, message);
                Err(e.into_err().into())
            }
        }
        Err(e) => {
            warn!("  Client {} got non-ServiceError: {:?}", client_id, e);
            Err(e.into())
        }
    }
}

async fn run_race_iteration(
    clients: &[Client],
    test_key: &str,
    iteration: usize,
) -> Result<usize, Box<dyn std::error::Error + Send + Sync>> {
    cleanup_object(&clients[0], test_key).await?;
    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;

    match clients[0].head_object().bucket(BUCKET).key(test_key).send().await {
        Ok(_) => return Err(format!("object still exists after cleanup in iteration {iteration}").into()),
        Err(error) if error.as_service_error().is_some_and(|error| error.is_not_found()) => {}
        Err(error) => return Err(format!("failed to verify cleanup in iteration {iteration}: {error:?}").into()),
    }

    info!("\n=== Iteration {} ===", iteration);
    info!("Launching {} concurrent conditional PUTs to different nodes...", clients.len());

    let barrier = Arc::new(Barrier::new(clients.len()));
    let test_key = test_key.to_string();

    let mut handles = vec![];
    for (i, client) in clients.iter().enumerate() {
        let client = client.clone();
        let barrier = barrier.clone();
        let key = test_key.clone();
        let data = format!("data from client {}", i).into_bytes();

        let handle = tokio::spawn(async move {
            barrier.wait().await;
            conditional_put(&client, &key, &data, i).await
        });
        handles.push(handle);
    }

    let mut success_count = 0;
    let mut had_error = false;
    for handle in handles {
        match handle.await {
            Ok(Ok(true)) => success_count += 1,
            Ok(Ok(false)) => {}
            Ok(Err(e)) => {
                had_error = true;
                info!("  Error: {}", e);
            }
            Err(e) => {
                had_error = true;
                info!("  Task error: {}", e);
            }
        }
    }

    info!("Result: {} out of {} succeeded", success_count, clients.len());

    if had_error {
        return Err("one or more conditional PUTs failed unexpectedly".into());
    }

    if success_count > 1 {
        info!(">>> RACE CONDITION DETECTED!");
    } else if success_count == 1 {
        info!(">>> Correct behavior: exactly 1 writer succeeded.");
    } else {
        return Err("no conditional PUT succeeded".into());
    }

    Ok(success_count)
}

#[tokio::test]
async fn test_conditional_put_race_cluster() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    crate::common::init_logging();
    info!("Starting conditional PUT race test with auto cluster");

    let mut cluster = RustFSTestClusterEnvironment::new(4).await?;
    cluster.start().await?;

    cluster.create_test_bucket(BUCKET).await?;

    let clients = cluster.create_all_clients()?;

    let iterations = 5;
    let mut races_detected = 0;
    let mut correct_count = 0;
    let mut error_count = 0;

    for i in 1..=iterations {
        let test_key = format!("race-test-{}-{}", std::process::id(), i);

        match run_race_iteration(&clients, &test_key, i).await {
            Ok(success_count) => {
                if success_count > 1 {
                    races_detected += 1;
                } else if success_count == 1 {
                    correct_count += 1;
                }
            }
            Err(e) => {
                error_count += 1;
                warn!("Iteration {} failed (not a race; e.g. cluster/network): {}", i, e)
            }
        }

        cleanup_object(&clients[0], &test_key).await?;
        tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
    }

    info!("\n====================================");
    info!("SUMMARY");
    info!("====================================");
    info!("Total iterations:   {}", iterations);
    info!("Correct (1 winner): {}", correct_count);
    info!("Race conditions:    {}", races_detected);
    info!("Failed iterations:  {}", error_count);

    assert_eq!(races_detected, 0, "Race conditions detected: {}/{}", races_detected, iterations);
    assert_eq!(
        error_count, 0,
        "{} iteration(s) failed due to errors (e.g. cluster not ready)",
        error_count
    );
    assert_eq!(
        correct_count, iterations,
        "only {correct_count}/{iterations} iterations observed exactly one winner"
    );

    Ok(())
}

#[tokio::test]
async fn test_conditional_put_basic_cluster() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    crate::common::init_logging();
    info!("Starting basic conditional PUT test with auto cluster");

    let mut cluster = RustFSTestClusterEnvironment::new(4).await?;
    cluster.start().await?;

    cluster.create_test_bucket(BUCKET).await?;

    let client = cluster.create_s3_client(0)?;
    let test_key = "basic-conditional-put";
    cleanup_object(&client, test_key).await?;

    let result = client
        .put_object()
        .bucket(BUCKET)
        .key(test_key)
        .body(Bytes::from("first write").into())
        .if_none_match("*")
        .send()
        .await;
    assert!(result.is_ok(), "First PUT with If-None-Match:* should succeed");

    let result = client
        .put_object()
        .bucket(BUCKET)
        .key(test_key)
        .body(Bytes::from("second write").into())
        .if_none_match("*")
        .send()
        .await;
    assert!(result.is_err(), "Second PUT with If-None-Match:* should fail");

    assert!(
        matches!(result, Err(SdkError::ServiceError(_))),
        "Expected ServiceError but got different error type"
    );

    if let Err(SdkError::ServiceError(e)) = result {
        let code = e.err().meta().code().unwrap_or("");
        assert_eq!(code, "PreconditionFailed");
    }

    cleanup_object(&client, test_key).await?;
    Ok(())
}

#[tokio::test]
async fn test_bucket_cors_write_is_visible_on_peer_before_response() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    crate::common::init_logging();

    let mut cluster = RustFSTestClusterEnvironment::new(2).await?;
    cluster.start().await?;
    cluster.create_test_bucket(BUCKET_METADATA_RELOAD_BUCKET).await?;

    let writer = cluster.create_s3_client(0)?;
    let reader = cluster.create_s3_client(1)?;
    assert_bucket_cors_missing(&reader).await;

    let rule = CorsRule::builder()
        .allowed_methods("GET")
        .allowed_origins("https://example.com")
        .build()?;
    let configuration = CorsConfiguration::builder().cors_rules(rule).build()?;

    writer
        .put_bucket_cors()
        .bucket(BUCKET_METADATA_RELOAD_BUCKET)
        .cors_configuration(configuration)
        .send()
        .await?;

    let response = reader.get_bucket_cors().bucket(BUCKET_METADATA_RELOAD_BUCKET).send().await?;
    let rules = response.cors_rules();
    assert_eq!(
        rules.len(),
        1,
        "peer should observe the committed CORS rule before the write response returns"
    );
    assert_eq!(rules[0].allowed_methods(), ["GET"]);
    assert_eq!(rules[0].allowed_origins(), ["https://example.com"]);

    writer
        .delete_bucket_cors()
        .bucket(BUCKET_METADATA_RELOAD_BUCKET)
        .send()
        .await?;
    assert_bucket_cors_missing(&reader).await;
    writer.delete_bucket().bucket(BUCKET_METADATA_RELOAD_BUCKET).send().await?;
    Ok(())
}
