#include "AudioConverter.h"

#include <factories/naudio/v1/BookFactory.h>
#include <factories/naudio/v1/LoopFactory.h>
#include <factories/naudio/v1/AudioContext.h>
#include <factories/naudio/v1/SampleFactory.h>
#include <factories/BaseFactory.h>
#include <Companion.h>
#include <cassert>
#include <cstring>
#include "hj/pyutils.h"

void AIFCWriter::End(std::string chunk, LUS::BinaryWriter& writer) {
    auto buffer = writer.ToVector();
    this->Chunks.push_back({ chunk, buffer });

    this->totalSize += ALIGN(buffer.size(), 2) + 8;
}

void AIFCWriter::Close(LUS::BinaryWriter& out) {
    out.SetEndianness(Torch::Endianness::Big);

    out.Write(AIFCMagicValues::FORM);
    out.Write((uint32_t)(this->totalSize + 4));
    out.Write(AIFCMagicValues::AIFC);

    for (auto& chunk : this->Chunks) {
        out.Write((char*)chunk.id.data(), chunk.id.size());
        out.Write((uint32_t)chunk.data.size());
        out.Write((char*)chunk.data.data(), chunk.data.size());

        if (chunk.data.size() % 2 == 1) {
            out.Write((uint8_t)0);
        }
    }
}

// Function to serialize double to 80-bit extended-precision
void SerializeF80(double num, LUS::BinaryWriter& writer) {
    // Convert the input double to a uint64_t representation
    uint64_t f64;
    memcpy((void*)&f64, (void*)&num, sizeof(double));

    // Extract the sign bit
    uint64_t f64_sign_bit = f64 & (1ULL << 63);

    // Handle the special case: zero
    if (num == 0.0) {
        if (f64_sign_bit) {
            writer.Write(static_cast<uint16_t>(0x8000)); // Sign bit set
        } else {
            writer.Write(static_cast<uint16_t>(0x0000)); // No sign bit
        }
        writer.Write(static_cast<uint64_t>(0x0000000000000000)); // Zero mantissa
        return;
    }

    // Extract the exponent and mantissa
    uint64_t exponent = (f64 >> 52) & 0x7FF; // Exponent bits
    assert(exponent != 0);                   // Ensure not denormal
    assert(exponent != 0x7FF);               // Ensure not infinity/NaN

    exponent -= 1023; // Adjust bias for 64-bit

    uint64_t f64_mantissa_bits = f64 & ((1ULL << 52) - 1); // Mantissa bits

    // Construct the 80-bit extended-precision fields
    uint64_t f80_sign_bit = f64_sign_bit << (80 - 64);                            // Shift sign
    uint64_t f80_exponent = (exponent + 0x3FFF) & 0x7FFF;                         // Adjust bias
    uint64_t f80_mantissa_bits = (1ULL << 63) | (f64_mantissa_bits << (63 - 52)); // Add implicit bit

    // Combine components into the 80-bit representation
    uint16_t high = static_cast<uint16_t>(f80_sign_bit >> 48) | static_cast<uint16_t>(f80_exponent);
    uint64_t low = f80_mantissa_bits;

    // Write the result in big-endian order
    writer.Write(high);
    writer.Write(low);
}

void AudioConverter::SampleV0ToAIFC(AudioBankSample* sample, LUS::BinaryWriter& out) {
    auto aifc = AIFCWriter();
    auto data = sample->data;

    uint32_t num_frames = data.size() * 16 / 9;
    uint32_t sample_rate = -1;

    if (sample->tunings.size() == 1) {
        sample_rate = 32000 * sample->tunings[0];
    } else {
        float tmin = PyUtils::min(sample->tunings);
        float tmax = PyUtils::max(sample->tunings);

        if (tmin <= 0.5f && 0.5f <= tmax) {
            sample_rate = 16000;
        } else if (tmin <= 1.0f && 1.0f <= tmax) {
            sample_rate = 32000;
        } else if (tmin <= 1.5f && 1.5f <= tmax) {
            sample_rate = 48000;
        } else if (tmin <= 2.5f && 2.5f <= tmax) {
            sample_rate = 80000;
        } else {
            sample_rate = 16000 * (tmin + tmax);
        }
    }

    int16_t num_channels = 1;
    int16_t sample_size = 16;

    // COMM Chunk
    auto comm = aifc.Start();
    comm.Write(num_channels);
    comm.Write(num_frames);
    comm.Write(sample_size);
    SerializeF80(sample_rate, comm);
    comm.Write(AIFCMagicValues::VAPC);
    comm.Write((char*)"\x0bVADPCM ~4-1", 12);
    aifc.End("COMM", comm);

    // INST Chunk
    auto inst = aifc.Start();
    for (size_t i = 0; i < 5; i++) {
        inst.Write((int32_t)0);
    }
    aifc.End("INST", inst);

    // VADPCMCODES Chunk
    auto vcodes = aifc.Start();
    vcodes.Write((char*)"stoc\x0bVADPCMCODES", 16);
    vcodes.Write((int16_t)1);
    vcodes.Write((int16_t)sample->book.order);
    vcodes.Write((int16_t)sample->book.npredictors);

    for (auto page : sample->book.table) {
        vcodes.Write(page);
    }
    aifc.End("APPL", vcodes);

    // SSND Chunk
    auto ssnd = aifc.Start();
    ssnd.Write((uint64_t)0);
    ssnd.Write((char*)data.data(), data.size());
    aifc.End("SSND", ssnd);

    // VADPCMLOOPS
    if (sample->loop.count != 0) {
        auto vloops = aifc.Start();
        vloops.Write((char*)"stoc\x0bVADPCMLOOPS", 16);
        vloops.Write((uint16_t)1);
        vloops.Write((uint16_t)1);
        vloops.Write(sample->loop.start);
        vloops.Write(sample->loop.end);
        vloops.Write(sample->loop.count);

        // The decoder reads back a fixed-size ALADPCMloop whose state is a
        // 16-entry array; always emit all 16 (zero-filled when absent) so the
        // chunk is exactly 44 bytes. Note: this must target vloops, not vcodes.
        const std::vector<int16_t>* state =
            sample->loop.state.has_value() ? &sample->loop.state.value() : nullptr;
        for (size_t i = 0; i < 16; i++) {
            vloops.Write((int16_t)(state != nullptr && i < state->size() ? (*state)[i] : 0));
        }
        aifc.End("APPL", vloops);
    }

    aifc.Close(out);
}

void AudioConverter::SampleV1ToAIFC(NSampleData* sample, LUS::BinaryWriter& out) {
    auto loop =
        std::static_pointer_cast<ADPCMLoopData>(Companion::Instance->GetParseDataByAddr(sample->loop)->data.value());
    auto book =
        std::static_pointer_cast<ADPCMBookData>(Companion::Instance->GetParseDataByAddr(sample->book)->data.value());
    SampleV1ToAIFC(sample, loop.get(), book.get(), out);
}

void AudioConverter::SampleV1ToAIFC(NSampleData* sample, const ADPCMLoopData* loop, const ADPCMBookData* book,
                                    LUS::BinaryWriter& out) {
    auto entry = AudioContext::tables[AudioTableType::SAMPLE_TABLE];
    auto sampleData = entry.buffer.data() + entry.info->entries[sample->sampleBankId].addr + sample->sampleAddr;
    auto aifc = AIFCWriter();
    std::vector<uint8_t> data(sampleData, sampleData + sample->size);

    uint32_t num_frames = data.size() * 16 / 9;
    uint32_t sample_rate = sample->sampleRate;

    if (sample_rate == 0) {
        sample_rate = 32000 * sample->tuning;
    }
    // Explicit yaml entries carry no tuning; a zero rate makes the AIFF
    // invalid (decoders reject it), so fall back to the mixer reference.
    if (sample_rate == 0) {
        sample_rate = 32000;
    }

    int16_t num_channels = 1;
    int16_t sample_size = 16;

    // COMM Chunk
    auto comm = aifc.Start();
    comm.Write(num_channels);
    comm.Write(num_frames);
    comm.Write(sample_size);
    SerializeF80(sample_rate, comm);
    comm.Write(AIFCMagicValues::VAPC);
    comm.Write((char*)"\x0bVADPCM ~4-1", 12);
    aifc.End("COMM", comm);

    // INST Chunk
    auto inst = aifc.Start();
    for (size_t i = 0; i < 5; i++) {
        inst.Write((int32_t)0);
    }
    aifc.End("INST", inst);

    // VADPCMCODES Chunk
    auto vcodes = aifc.Start();
    vcodes.Write((char*)"stoc\x0bVADPCMCODES", 16);
    vcodes.Write((int16_t)1);
    vcodes.Write((int16_t)book->order);
    vcodes.Write((int16_t)book->numPredictors);

    for (auto page : book->book) {
        vcodes.Write(page);
    }
    aifc.End("APPL", vcodes);

    // SSND Chunk
    auto ssnd = aifc.Start();
    ssnd.Write((uint64_t)0);
    ssnd.Write((char*)data.data(), data.size());
    aifc.End("SSND", ssnd);

    // VADPCMLOOPS
    if (loop->count != 0) {
        auto vloops = aifc.Start();
        vloops.Write((char*)"stoc\x0bVADPCMLOOPS", 16);
        vloops.Write((uint16_t)1);
        vloops.Write((uint16_t)1);
        vloops.Write(loop->start);
        vloops.Write(loop->end);
        vloops.Write(loop->count);

        for (size_t i = 0; i < 16; i++) {
            vloops.Write(loop->predictorState[i]);
        }
        aifc.End("APPL", vloops);
    }

    aifc.Close(out);
}