#include "sample_decoder.hpp"
#include <algorithm>

namespace pdfengine {

std::vector<uint8_t> SampleDecoder::unpackSamples(const std::vector<uint8_t>& rawBytes,
                                                  int width,
                                                  int height,
                                                  int components,
                                                  int bitsPerComponent) {
    if (rawBytes.empty() || width <= 0 || height <= 0 || components <= 0) {
        return {};
    }

    if (bitsPerComponent == 8) {
        return rawBytes;
    }

    const size_t expectedSamples = static_cast<size_t>(width) * height * components;
    std::vector<uint8_t> unpacked;
    unpacked.reserve(expectedSamples);

    if (bitsPerComponent == 1) {
        for (int y = 0; y < height; ++y) {
            int samplesInRow = width * components;
            int bytesInRow = (samplesInRow + 7) / 8;
            size_t rowOffset = static_cast<size_t>(y) * bytesInRow;
            if (rowOffset >= rawBytes.size()) break;

            for (int s = 0; s < samplesInRow; ++s) {
                size_t byteIdx = rowOffset + (s / 8);
                if (byteIdx >= rawBytes.size()) break;
                uint8_t b = rawBytes[byteIdx];
                uint8_t bit = (b >> (7 - (s % 8))) & 1;
                unpacked.push_back(bit ? 255 : 0);
            }
        }
    } else if (bitsPerComponent == 2) {
        for (int y = 0; y < height; ++y) {
            int samplesInRow = width * components;
            int bytesInRow = (samplesInRow + 3) / 4;
            size_t rowOffset = static_cast<size_t>(y) * bytesInRow;
            if (rowOffset >= rawBytes.size()) break;

            for (int s = 0; s < samplesInRow; ++s) {
                size_t byteIdx = rowOffset + (s / 4);
                if (byteIdx >= rawBytes.size()) break;
                uint8_t b = rawBytes[byteIdx];
                uint8_t val = (b >> (6 - 2 * (s % 4))) & 0x03;
                unpacked.push_back(static_cast<uint8_t>(val * 85)); // 0x03 -> 255
            }
        }
    } else if (bitsPerComponent == 4) {
        for (int y = 0; y < height; ++y) {
            int samplesInRow = width * components;
            int bytesInRow = (samplesInRow + 1) / 2;
            size_t rowOffset = static_cast<size_t>(y) * bytesInRow;
            if (rowOffset >= rawBytes.size()) break;

            for (int s = 0; s < samplesInRow; ++s) {
                size_t byteIdx = rowOffset + (s / 2);
                if (byteIdx >= rawBytes.size()) break;
                uint8_t b = rawBytes[byteIdx];
                uint8_t val = (s % 2 == 0) ? ((b >> 4) & 0x0F) : (b & 0x0F);
                unpacked.push_back(static_cast<uint8_t>(val * 17)); // 0x0F -> 255
            }
        }
    } else if (bitsPerComponent == 16) {
        // Take upper 8 bits of each 16-bit sample
        for (size_t i = 0; i + 1 < rawBytes.size() && unpacked.size() < expectedSamples; i += 2) {
            unpacked.push_back(rawBytes[i]);
        }
    } else {
        // Fallback pass-through
        return rawBytes;
    }

    return unpacked;
}

} // namespace pdfengine