// 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.

//! Reliability tests driven by in-process fault injection (see `chaos.rs`):
//! degraded reads/writes with an offline disk, bitrot read-through, and
//! fresh-disk replacement heal after a SIGKILL restart.
//!
//! All tests use a single-node 4-disk topology (default erasure coding for
//! 4 drives is 2 data + 2 parity) and verify object content via sha256
//! manifests recorded at write time.

#[cfg(test)]
mod tests {
    use crate::chaos::{DiskFaultHarness, VersionShardCensus, signed_admin_post};
    use crate::common::init_logging;
    use aws_sdk_s3::Client;
    use aws_sdk_s3::primitives::ByteStream;
    use aws_sdk_s3::types::{BucketVersioningStatus, CompletedMultipartUpload, CompletedPart, VersioningConfiguration};
    use serial_test::serial;
    use sha2::{Digest, Sha256};
    use std::collections::HashSet;
    use std::error::Error;
    use tokio::time::{Duration, Instant, interval, timeout};
    use tracing::info;

    const GET_TIMEOUT: Duration = Duration::from_secs(60);
    const PUT_TIMEOUT: Duration = Duration::from_secs(60);

    fn sha256_hex(data: &[u8]) -> String {
        let digest = Sha256::digest(data);
        digest.iter().map(|byte| format!("{byte:02x}")).collect()
    }

    /// Deterministic pseudo-random payload so tests stay reproducible.
    fn payload(len: usize, seed: u8) -> Vec<u8> {
        (0..len)
            .map(|i| (i as u64).wrapping_mul(2654435761).wrapping_add(seed as u64) as u8)
            .collect()
    }

    async fn put_and_record(
        client: &Client,
        bucket: &str,
        key: &str,
        body: Vec<u8>,
        manifest: &mut Vec<(String, String)>,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let digest = sha256_hex(&body);
        timeout(
            PUT_TIMEOUT,
            client
                .put_object()
                .bucket(bucket)
                .key(key)
                .body(ByteStream::from(body))
                .send(),
        )
        .await
        .map_err(|_| format!("PUT {key} timed out"))??;
        manifest.push((key.to_string(), digest));
        Ok(())
    }

    async fn multipart_put_and_record(
        client: &Client,
        bucket: &str,
        key: &str,
        parts: Vec<Vec<u8>>,
        manifest: &mut Vec<(String, String)>,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let full_body: Vec<u8> = parts.iter().flatten().copied().collect();
        let digest = sha256_hex(&full_body);

        let create = client.create_multipart_upload().bucket(bucket).key(key).send().await?;
        let upload_id = create.upload_id().ok_or("missing upload id")?.to_string();

        let mut completed_parts = Vec::with_capacity(parts.len());
        for (index, part_body) in parts.into_iter().enumerate() {
            let part_number = (index + 1) as i32;
            let uploaded = timeout(
                PUT_TIMEOUT,
                client
                    .upload_part()
                    .bucket(bucket)
                    .key(key)
                    .upload_id(&upload_id)
                    .part_number(part_number)
                    .body(ByteStream::from(part_body))
                    .send(),
            )
            .await
            .map_err(|_| format!("upload_part {part_number} for {key} timed out"))??;
            completed_parts.push(
                CompletedPart::builder()
                    .part_number(part_number)
                    .e_tag(uploaded.e_tag().ok_or("missing part etag")?)
                    .build(),
            );
        }

        client
            .complete_multipart_upload()
            .bucket(bucket)
            .key(key)
            .upload_id(&upload_id)
            .multipart_upload(CompletedMultipartUpload::builder().set_parts(Some(completed_parts)).build())
            .send()
            .await?;

        manifest.push((key.to_string(), digest));
        Ok(())
    }

    async fn verify_manifest(
        client: &Client,
        bucket: &str,
        manifest: &[(String, String)],
        phase: &str,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        for (key, expected_sha256) in manifest {
            let response = timeout(GET_TIMEOUT, client.get_object().bucket(bucket).key(key).send())
                .await
                .map_err(|_| format!("GET {key} timed out during {phase}"))?
                .map_err(|err| format!("GET {key} failed during {phase}: {err}"))?;
            let body = response
                .body
                .collect()
                .await
                .map_err(|err| format!("GET {key} body collect failed during {phase}: {err}"))?
                .into_bytes();
            let actual_sha256 = sha256_hex(&body);
            if actual_sha256 != *expected_sha256 {
                return Err(format!(
                    "sha256 mismatch for {key} during {phase}: expected {expected_sha256}, got {actual_sha256} ({} bytes)",
                    body.len()
                )
                .into());
            }
        }
        info!("Verified {} objects during {}", manifest.len(), phase);
        Ok(())
    }

    /// One disk goes offline at runtime: all previously written objects (from
    /// inline-small to multipart-sized) must stay readable with intact
    /// content, degraded writes must succeed, and everything must still
    /// verify after the disk returns.
    #[tokio::test]
    #[serial]
    async fn test_degraded_read_write_with_one_disk_offline() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        init_logging();
        info!("Reliability: degraded read/write with one of four disks offline");

        let mut harness = DiskFaultHarness::new(4).await?;
        harness.start_server().await?;
        let client = harness.env.create_s3_client();

        let bucket = "reliability-degraded-rw";
        client.create_bucket().bucket(bucket).send().await?;

        let mut manifest: Vec<(String, String)> = Vec::new();
        put_and_record(&client, bucket, "degraded/small.bin", payload(4 * 1024, 1), &mut manifest).await?;
        put_and_record(&client, bucket, "degraded/medium.bin", payload(1024 * 1024, 2), &mut manifest).await?;
        put_and_record(&client, bucket, "degraded/big.bin", payload(3 * 1024 * 1024, 3), &mut manifest).await?;
        multipart_put_and_record(
            &client,
            bucket,
            "degraded/multipart.bin",
            vec![payload(5 * 1024 * 1024, 4), payload(1024 * 1024, 5)],
            &mut manifest,
        )
        .await?;

        verify_manifest(&client, bucket, &manifest, "baseline with all disks online").await?;

        harness.take_disk_offline(0)?;

        verify_manifest(&client, bucket, &manifest, "degraded read with disk0 offline").await?;

        // Degraded writes: 3 of 4 disks still satisfy the write quorum for
        // an EC 2+2 set.
        put_and_record(
            &client,
            bucket,
            "degraded/written-while-offline.bin",
            payload(1024 * 1024, 9),
            &mut manifest,
        )
        .await?;
        verify_manifest(&client, bucket, &manifest, "read-back of degraded write").await?;

        harness.bring_disk_online(0)?;

        verify_manifest(&client, bucket, &manifest, "after disk0 came back online").await?;
        Ok(())
    }

    /// Silent bitrot in a single erasure shard must never surface corrupted
    /// bytes to a reader: per-shard bitrot checksums reject the bad shard and
    /// the object is reconstructed from the remaining shards.
    #[tokio::test]
    #[serial]
    async fn test_bitrot_corrupted_shard_read_returns_correct_data() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        init_logging();
        info!("Reliability: GET must read through a bitrot-corrupted shard");

        let mut harness = DiskFaultHarness::new(4).await?;
        harness.start_server().await?;
        let client = harness.env.create_s3_client();

        let bucket = "reliability-bitrot";
        client.create_bucket().bucket(bucket).send().await?;

        let key_a = "bitrot/object-a.bin";
        let key_b = "bitrot/object-b.bin";
        let mut manifest: Vec<(String, String)> = Vec::new();
        // 2 MiB objects are far above the 128 KiB inline threshold, so every
        // disk holds a real part.1 shard file to corrupt.
        put_and_record(&client, bucket, key_a, payload(2 * 1024 * 1024, 21), &mut manifest).await?;
        put_and_record(&client, bucket, key_b, payload(2 * 1024 * 1024, 22), &mut manifest).await?;

        verify_manifest(&client, bucket, &manifest, "baseline before corruption").await?;

        // Corrupt two shards per object (the maximum an EC 2+2 set can lose)
        // on different disk pairs. Which disks hold data vs parity depends on
        // the per-object distribution, so corrupting a pair makes it very
        // likely that at least one data shard is hit; either way the read
        // must return intact content reconstructed from the clean shards.
        harness.corrupt_object_shard(0, bucket, key_a)?;
        harness.corrupt_object_shard(1, bucket, key_a)?;
        harness.corrupt_object_shard(2, bucket, key_b)?;
        harness.corrupt_object_shard(3, bucket, key_b)?;

        verify_manifest(&client, bucket, &manifest, "first read after shard corruption").await?;
        // A second pass ensures repeated reads stay correct as well.
        verify_manifest(&client, bucket, &manifest, "second read after shard corruption").await?;
        Ok(())
    }

    /// Fresh-disk replacement: SIGKILL the server, swap one disk for an empty
    /// directory, restart with the same volumes/port, trigger an admin deep
    /// heal, and require the replaced disk to be rebuilt and all content to
    /// verify against the sha256 manifest.
    #[tokio::test]
    #[serial]
    async fn test_fresh_disk_replacement_heals_after_sigkill_restart() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        init_logging();
        info!("Reliability: fresh-disk replacement heals after SIGKILL restart");

        let mut harness = DiskFaultHarness::new(4).await?;
        harness.start_server().await?;
        let client = harness.env.create_s3_client();

        let bucket = "reliability-fresh-disk";
        client.create_bucket().bucket(bucket).send().await?;

        let mut manifest: Vec<(String, String)> = Vec::new();
        put_and_record(&client, bucket, "heal/tiny.bin", payload(4 * 1024, 31), &mut manifest).await?;
        put_and_record(&client, bucket, "heal/small.bin", payload(256 * 1024, 32), &mut manifest).await?;
        put_and_record(&client, bucket, "heal/medium.bin", payload(1024 * 1024, 33), &mut manifest).await?;
        put_and_record(&client, bucket, "heal/nested/large.bin", payload(2 * 1024 * 1024, 34), &mut manifest).await?;

        verify_manifest(&client, bucket, &manifest, "baseline before disk replacement").await?;
        let manifest_keys = manifest.iter().map(|(key, _)| key.clone()).collect::<Vec<_>>();
        let target_manifest: Vec<(String, VersionShardCensus)> = manifest_keys
            .iter()
            .map(|key| {
                let census = harness.census_object_version(0, bucket, key, None)?;
                if !census.is_complete() {
                    return Err(format!("disk 0 has incomplete physical census for {key}: {census:?}").into());
                }
                Ok((key.clone(), census))
            })
            .collect::<Result<_, Box<dyn Error + Send + Sync>>>()?;

        harness.kill_server();
        harness.replace_disk_with_empty(0)?;
        harness.restart_server().await?;

        let heal_body = r#"{"recursive":true,"dryRun":false,"remove":false,"recreate":true,"scanMode":2,"updateParity":false,"nolock":false}"#;
        let heal_url = format!("{}/rustfs/admin/v3/heal/{}?forceStart=true", harness.env.url, bucket);
        signed_admin_post(&heal_url, Some(heal_body), &harness.env.access_key, &harness.env.secret_key).await?;

        let client = harness.env.create_s3_client();
        let mut remaining: HashSet<String> = manifest_keys.iter().cloned().collect();
        let heal_timeout_secs = std::env::var("RUSTFS_RELIABILITY_HEAL_TIMEOUT_SECS")
            .ok()
            .and_then(|value| value.parse::<u64>().ok())
            .unwrap_or(120);
        let deadline = Instant::now() + Duration::from_secs(heal_timeout_secs);
        let mut retry = interval(Duration::from_secs(1));

        loop {
            remaining.retain(|key| {
                let expected = target_manifest
                    .iter()
                    .find(|(manifest_key, _)| manifest_key == key)
                    .map(|(_, manifest)| manifest)
                    .expect("every key has a physical manifest");
                harness
                    .census_object_version(0, bucket, key, None)
                    .map(|census| !census.matches_manifest(expected))
                    .unwrap_or(true)
            });
            if remaining.is_empty() {
                verify_manifest(&client, bucket, &manifest, "after fresh-disk heal completed").await?;
                return Ok(());
            }
            if Instant::now() >= deadline {
                break;
            }
            retry.tick().await;
        }

        Err(format!("fresh-disk heal did not rebuild {remaining:?} on the replaced disk within {heal_timeout_secs}s").into())
    }

    #[tokio::test]
    #[serial]
    async fn test_versioned_shard_census_selects_each_version_data_dir() -> Result<(), Box<dyn Error + Send + Sync>> {
        init_logging();
        info!("Reliability: physical shard census selects the requested object version");

        let mut harness = DiskFaultHarness::new(4).await?;
        harness.start_server().await?;
        let client = harness.env.create_s3_client();
        let bucket = "reliability-versioned-census";
        let key = "versions/large.bin";
        client.create_bucket().bucket(bucket).send().await?;
        client
            .put_bucket_versioning()
            .bucket(bucket)
            .versioning_configuration(
                VersioningConfiguration::builder()
                    .status(BucketVersioningStatus::Enabled)
                    .build(),
            )
            .send()
            .await?;

        let first_inline = client
            .put_object()
            .bucket(bucket)
            .key("versions/inline.bin")
            .body(ByteStream::from(payload(8 * 1024, 40)))
            .send()
            .await?;
        let first_inline_version = first_inline
            .version_id()
            .ok_or("first inline PUT did not return a version ID")?;
        let second_inline = client
            .put_object()
            .bucket(bucket)
            .key("versions/inline.bin")
            .body(ByteStream::from(payload(8 * 1024, 41)))
            .send()
            .await?;
        let second_inline_version = second_inline
            .version_id()
            .ok_or("second inline PUT did not return a version ID")?;

        let first = client
            .put_object()
            .bucket(bucket)
            .key(key)
            .body(ByteStream::from(payload(128 * 1024, 41)))
            .send()
            .await?;
        let first_version = first.version_id().ok_or("first PUT did not return a version ID")?;
        let second = client
            .put_object()
            .bucket(bucket)
            .key(key)
            .body(ByteStream::from(payload(3 * 1024 * 1024, 42)))
            .send()
            .await?;
        let second_version = second.version_id().ok_or("second PUT did not return a version ID")?;
        let delete = client.delete_object().bucket(bucket).key(key).send().await?;
        let delete_version = delete.version_id().ok_or("delete marker did not return a version ID")?;

        let first_inline_census = harness.census_object_version(0, bucket, "versions/inline.bin", Some(first_inline_version))?;
        let second_inline_census =
            harness.census_object_version(0, bucket, "versions/inline.bin", Some(second_inline_version))?;
        let first_census = harness.census_object_version(0, bucket, key, Some(first_version))?;
        let first_other_disk_census = harness.census_object_version(1, bucket, key, Some(first_version))?;
        let second_census = harness.census_object_version(0, bucket, key, Some(second_version))?;
        let delete_census = harness.census_object_version(0, bucket, key, Some(delete_version))?;
        assert!(
            first_inline_census.is_complete() && second_inline_census.is_complete(),
            "inline version physical census is incomplete: first={first_inline_census:?} second={second_inline_census:?}"
        );
        assert!(
            first_inline_census.present_part_fingerprints.is_empty() && second_inline_census.present_part_fingerprints.is_empty(),
            "inline versions must not select external shard files: first={first_inline_census:?} second={second_inline_census:?}"
        );
        assert!(
            first_inline_census.inline_data_fingerprint.is_some() && second_inline_census.inline_data_fingerprint.is_some(),
            "inline versions must fingerprint payload bytes stored in xl.meta"
        );
        assert_ne!(
            first_inline_census.inline_data_fingerprint, second_inline_census.inline_data_fingerprint,
            "same-size inline versions with different payloads must retain distinct xl.meta fingerprints"
        );
        assert!(
            first_census.is_complete(),
            "first version physical census is incomplete: {first_census:?}"
        );
        assert!(
            second_census.is_complete(),
            "second version physical census is incomplete: {second_census:?}"
        );
        assert!(
            first_other_disk_census.is_complete(),
            "first version physical census on the second disk is incomplete: {first_other_disk_census:?}"
        );
        assert_ne!(
            first_census.erasure_index, first_other_disk_census.erasure_index,
            "physical census must preserve each disk's erasure index"
        );
        assert_ne!(
            first_census.data_dir, second_census.data_dir,
            "distinct object versions must select distinct physical data directories"
        );
        assert_eq!(
            first_census.expected_part_numbers, second_census.expected_part_numbers,
            "same single-part shape should expose the same part numbers"
        );
        let first_part = first_census
            .present_part_fingerprints
            .values()
            .next()
            .ok_or("first version did not expose a physical part fingerprint")?;
        let second_part = second_census
            .present_part_fingerprints
            .values()
            .next()
            .ok_or("second version did not expose a physical part fingerprint")?;
        assert_ne!(
            first_part.size, second_part.size,
            "different shard lengths must retain their physical sizes"
        );
        assert_ne!(
            first_part.sha256, second_part.sha256,
            "different shard contents must retain their physical hashes"
        );
        assert!(
            delete_census.is_complete(),
            "delete marker physical census is incomplete: {delete_census:?}"
        );
        assert!(
            delete_census.expected_part_numbers.is_empty(),
            "delete marker must not declare object shards: {delete_census:?}"
        );
        assert!(
            delete_census.present_part_fingerprints.is_empty(),
            "delete marker must not select stale object shards: {delete_census:?}"
        );
        Ok(())
    }
}