/*
 * This file is part of the Aaru Data Preservation Suite.
 * Copyright (c) 2019-2026 Natalia Portillo.
 *
 * This library is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation; either version 2.1 of the
 * License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, see <http://www.gnu.org/licenses/>.
 */

/**
 * @file erasure_coding.cpp
 * @brief Integration tests for erasure coding: create EC image from real test data,
 *        corrupt compressed blocks, and verify recovery via RS decoding.
 */

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <vector>

#include <gtest/gtest.h>

extern "C"
{
#include <aaruformat.h>
#include <aaruformat/errors.h>
}

class ErasureCodingTest : public ::testing::Test
{
  protected:
    /**
     * @brief Copy a file byte-for-byte.
     */
    static bool copy_file(const char *src, const char *dst)
    {
        FILE *in  = fopen(src, "rb");
        FILE *out = fopen(dst, "wb");
        if(!in || !out) { if(in) fclose(in); if(out) fclose(out); return false; }

        uint8_t buf[65536];
        size_t n;
        while((n = fread(buf, 1, sizeof(buf), in)) > 0)
            fwrite(buf, 1, n, out);

        fclose(in);
        fclose(out);
        return true;
    }

    /**
     * @brief Corrupt bytes at a specific offset in a file.
     */
    static void corrupt_file(const char *path, long offset, size_t len)
    {
        FILE *f = fopen(path, "r+b");
        ASSERT_NE(f, nullptr) << "Cannot open file for corruption: " << path;
        fseek(f, offset, SEEK_SET);
        for(size_t i = 0; i < len; i++)
            fputc((int)(0xDE ^ (i & 0xFF)), f);
        fclose(f);
    }

    /**
     * @brief Create a new image with EC from an existing test image.
     *
     * Opens source image, reads all sectors + computes golden CRC64,
     * creates new image with EC enabled, writes all sectors, closes.
     *
     * @param src_path  Path to source .aif image.
     * @param dst_path  Path for new EC-enabled image.
     * @param K         Data blocks per stripe.
     * @param M         Parity blocks per stripe.
     * @param golden_crc Output: CRC64 of all sectors.
     * @return true on success.
     */
    bool create_ec_image(const char *src_path, const char *dst_path,
                         uint16_t K, uint16_t M, uint64_t *golden_crc)
    {
        /* Open source image */
        void *src_ctx = aaruf_open(src_path, false, nullptr);
        if(!src_ctx) return false;

        ImageInfo info;
        if(aaruf_get_image_info(src_ctx, &info) != AARUF_STATUS_OK)
        { aaruf_close(src_ctx); return false; }

        /* Read all sectors and compute golden CRC64 */
        crc64_ctx *crc_ctx = aaruf_crc64_init();
        std::vector<uint8_t> sector_buf(info.SectorSize);

        for(uint64_t i = 0; i < info.Sectors; i++)
        {
            uint32_t len = info.SectorSize;
            uint8_t  status = 0;
            int32_t  rc = aaruf_read_sector(src_ctx, i, false, sector_buf.data(), &len, &status);
            if(rc != AARUF_STATUS_OK) { aaruf_crc64_free(crc_ctx); aaruf_close(src_ctx); return false; }
            aaruf_crc64_update(crc_ctx, sector_buf.data(), len);
        }
        aaruf_crc64_final(crc_ctx, golden_crc);
        aaruf_crc64_free(crc_ctx);

        /* Create new image with EC */
        void *dst_ctx = aaruf_create(dst_path, info.MediaType, info.SectorSize,
                                     info.Sectors, 0, 0,
                                     "compress=true;deduplicate=false",
                                     (const uint8_t *)"gtest_ec", 8, 0, 1, false);
        if(!dst_ctx) { aaruf_close(src_ctx); return false; }

        int32_t ec_rc = aaruf_set_erasure_coding(dst_ctx, 1 /* RS-Vandermonde */, K, M);
        if(ec_rc != AARUF_STATUS_OK) { aaruf_close(dst_ctx); aaruf_close(src_ctx); return false; }

        /* Write all sectors from source to destination */
        for(uint64_t i = 0; i < info.Sectors; i++)
        {
            uint32_t len = info.SectorSize;
            uint8_t  status = 0;
            aaruf_read_sector(src_ctx, i, false, sector_buf.data(), &len, &status);
            int32_t wrc = aaruf_write_sector(dst_ctx, i, false, sector_buf.data(), status, len);
            if(wrc != AARUF_STATUS_OK) { aaruf_close(dst_ctx); aaruf_close(src_ctx); return false; }
        }

        aaruf_close(src_ctx);
        int close_rc = aaruf_close(dst_ctx);
        return close_rc == AARUF_STATUS_OK;
    }

    /**
     * @brief Read all sectors from an image and compute CRC64.
     */
    uint64_t compute_image_crc(const char *path, bool *success)
    {
        *success = false;
        void *ctx = aaruf_open(path, false, nullptr);
        if(!ctx) return 0;

        ImageInfo info;
        if(aaruf_get_image_info(ctx, &info) != AARUF_STATUS_OK) { aaruf_close(ctx); return 0; }

        crc64_ctx *crc_ctx = aaruf_crc64_init();
        std::vector<uint8_t> sector_buf(info.SectorSize);

        for(uint64_t i = 0; i < info.Sectors; i++)
        {
            uint32_t len = info.SectorSize;
            uint8_t  status = 0;
            int32_t  rc = aaruf_read_sector(ctx, i, false, sector_buf.data(), &len, &status);
            if(rc != AARUF_STATUS_OK)
            {
                aaruf_crc64_free(crc_ctx);
                aaruf_close(ctx);
                return 0;
            }
            aaruf_crc64_update(crc_ctx, sector_buf.data(), len);
        }

        uint64_t crc;
        aaruf_crc64_final(crc_ctx, &crc);
        aaruf_crc64_free(crc_ctx);
        aaruf_close(ctx);
        *success = true;
        return crc;
    }

    /**
     * @brief Find a data block offset in the file by scanning for DBLK identifiers.
     *
     * Skips past the header and finds the Nth data block (BlockHeader with identifier == DataBlock
     * and type == kDataTypeUserData).
     *
     * @param path File path.
     * @param skip_count Number of data blocks to skip (0 = first).
     * @return File offset of the block, or -1 if not found.
     */
    long find_data_block_offset(const char *path, int skip_count)
    {
        FILE *f = fopen(path, "rb");
        if(!f) return -1;

        /* Start after the header (128 bytes) */
        fseek(f, 128, SEEK_SET);

        uint8_t buf[4];
        int found = 0;
        long offset = -1;

        while(fread(buf, 1, 4, f) == 4)
        {
            /* Check for DataBlock identifier (0x4B4C4244 = "DBLK" LE) */
            uint32_t id = buf[0] | (buf[1] << 8) | (buf[2] << 16) | (buf[3] << 24);
            if(id == 0x4B4C4244)
            {
                /* Read type field (next 2 bytes) */
                uint16_t type;
                if(fread(&type, 2, 1, f) != 1) break;

                if(type == 1) /* kDataTypeUserData */
                {
                    if(found == skip_count)
                    {
                        offset = ftell(f) - 6; /* Back to start of BlockHeader */
                        break;
                    }
                    found++;
                }

                /* Skip past this block */
                fseek(f, -2, SEEK_CUR); /* back to after identifier */
                /* Read cmpLength to skip the rest */
                fseek(f, 4, SEEK_CUR); /* skip compression(2) + sectorSize(4) wait, no... */
            }

            /* Scan forward byte by byte (slow but reliable for testing) */
            fseek(f, -3, SEEK_CUR);
        }

        fclose(f);
        return offset;
    }

    void SetUp() override
    {
        /* Clean up any leftover temp files from previous runs */
        remove("ec_test_output.aif");
        remove("ec_test_corrupt.aif");
    }

    void TearDown() override
    {
        /* Clean up temp files */
        remove("ec_test_output.aif");
        remove("ec_test_corrupt.aif");
    }
};

/**
 * @test Create EC image from mf2hd.aif, verify uncorrupted reads match golden CRC.
 */
TEST_F(ErasureCodingTest, CreateAndVerifyUncorrupted)
{
    uint64_t golden_crc = 0;
    ASSERT_TRUE(create_ec_image("data/mf2hd.aif", "ec_test_output.aif", 4, 2, &golden_crc));

    bool success = false;
    uint64_t read_crc = compute_image_crc("ec_test_output.aif", &success);
    ASSERT_TRUE(success) << "Failed to read all sectors from EC image";
    EXPECT_EQ(read_crc, golden_crc) << "CRC mismatch on uncorrupted EC image";
}

/**
 * @test Create EC image, corrupt a data block, verify recovery produces correct data.
 */
TEST_F(ErasureCodingTest, RecoverCorruptedDataBlock)
{
    uint64_t golden_crc = 0;
    ASSERT_TRUE(create_ec_image("data/mf2hd.aif", "ec_test_output.aif", 4, 2, &golden_crc));

    /* Copy to corrupt version */
    ASSERT_TRUE(copy_file("ec_test_output.aif", "ec_test_corrupt.aif"));

    /* The first data block starts right after the 128-byte header, aligned.
     * With default blockAlignmentShift=9 (512 bytes), the first block starts at offset 512.
     * Corrupt 128 bytes at offset 550 (past BlockHeader, into compressed payload). */
    corrupt_file("ec_test_corrupt.aif", 550, 128);

    /* Read all sectors — should recover via EC */
    bool success = false;
    uint64_t read_crc = compute_image_crc("ec_test_corrupt.aif", &success);
    ASSERT_TRUE(success) << "EC recovery failed — could not read all sectors from corrupted image";
    EXPECT_EQ(read_crc, golden_crc) << "CRC mismatch after EC recovery";
}

/**
 * @test Create EC image, corrupt the BlockHeader of a data block.
 */
TEST_F(ErasureCodingTest, RecoverCorruptedBlockHeader)
{
    uint64_t golden_crc = 0;
    ASSERT_TRUE(create_ec_image("data/mf2hd.aif", "ec_test_output.aif", 4, 2, &golden_crc));

    ASSERT_TRUE(copy_file("ec_test_output.aif", "ec_test_corrupt.aif"));

    /* Corrupt the first block's header at offset 512 (after alignment).
     * Corrupt 32 bytes = entire BlockHeader, making identifier/compression/sectorSize garbage. */
    corrupt_file("ec_test_corrupt.aif", 512, 32);

    bool success = false;
    uint64_t read_crc = compute_image_crc("ec_test_corrupt.aif", &success);
    ASSERT_TRUE(success) << "EC recovery failed on corrupted BlockHeader";
    EXPECT_EQ(read_crc, golden_crc);
}

/**
 * @test Verify that corrupting a parity block doesn't affect normal reads.
 */
TEST_F(ErasureCodingTest, ParityCorruptionDoesNotAffectReads)
{
    uint64_t golden_crc = 0;
    ASSERT_TRUE(create_ec_image("data/mf2hd.aif", "ec_test_output.aif", 4, 1, &golden_crc));

    ASSERT_TRUE(copy_file("ec_test_output.aif", "ec_test_corrupt.aif"));

    /* Find a parity block (DataType == kDataTypeErasureParity = 104) by scanning */
    FILE *f = fopen("ec_test_corrupt.aif", "r+b");
    ASSERT_NE(f, nullptr);
    fseek(f, 0, SEEK_END);
    long file_size = ftell(f);

    /* Scan from offset 128 for parity block identifier */
    bool found_parity = false;
    for(long pos = 128; pos < file_size - 6; pos++)
    {
        fseek(f, pos, SEEK_SET);
        uint32_t id;
        uint16_t type;
        if(fread(&id, 4, 1, f) != 1) break;
        if(fread(&type, 2, 1, f) != 1) break;

        if(id == 0x4B4C4244 && type == 104) /* DataBlock + kDataTypeErasureParity */
        {
            /* Corrupt 32 bytes of parity payload */
            fseek(f, pos + 40, SEEK_SET);
            for(int i = 0; i < 32; i++) fputc(0xFF, f);
            found_parity = true;
            break;
        }
    }
    fclose(f);

    if(!found_parity) { GTEST_SKIP() << "No parity block found in image"; }

    /* Data should still read fine — parity is only needed during recovery */
    bool success = false;
    uint64_t read_crc = compute_image_crc("ec_test_corrupt.aif", &success);
    ASSERT_TRUE(success);
    EXPECT_EQ(read_crc, golden_crc);
}

/**
 * @test Create EC image with XOR (M=1), corrupt one block, verify recovery.
 */
TEST_F(ErasureCodingTest, RecoverWithXOR_M1)
{
    uint64_t golden_crc = 0;
    ASSERT_TRUE(create_ec_image("data/mf2hd.aif", "ec_test_output.aif", 4, 1, &golden_crc));

    ASSERT_TRUE(copy_file("ec_test_output.aif", "ec_test_corrupt.aif"));

    long block_offset = find_data_block_offset("ec_test_corrupt.aif", 0);
    ASSERT_GT(block_offset, 0);

    corrupt_file("ec_test_corrupt.aif", block_offset + 40, 64);

    bool success = false;
    uint64_t read_crc = compute_image_crc("ec_test_corrupt.aif", &success);
    ASSERT_TRUE(success) << "XOR recovery failed";
    EXPECT_EQ(read_crc, golden_crc);
}

/**
 * @test Corrupt the primary header (first 128 bytes), verify backup header from recovery footer is used.
 */
TEST_F(ErasureCodingTest, BackupHeaderRecovery)
{
    uint64_t golden_crc = 0;
    ASSERT_TRUE(create_ec_image("data/mf2hd.aif", "ec_test_output.aif", 4, 2, &golden_crc));

    ASSERT_TRUE(copy_file("ec_test_output.aif", "ec_test_corrupt.aif"));

    /* Corrupt the first 64 bytes of the file (part of AaruHeaderV2, including the magic) */
    corrupt_file("ec_test_corrupt.aif", 0, 64);

    /* Open should succeed using the backup header from the recovery footer */
    bool success = false;
    uint64_t read_crc = compute_image_crc("ec_test_corrupt.aif", &success);
    ASSERT_TRUE(success) << "Backup header recovery failed — could not open/read corrupted image";
    EXPECT_EQ(read_crc, golden_crc) << "CRC mismatch after backup header recovery";
}