const { webcrypto } = require('crypto');
const sharp = require('sharp');

function pemToDer(pem) {
    const b64 = pem.replace(/-----BEGIN PRIVATE KEY-----/, '').replace(/-----END PRIVATE KEY-----/, '').replace(/\\n/g, '').replace(/\s/g, '');
    const binaryDer = atob(b64);
    const buffer = new ArrayBuffer(binaryDer.length);
    const bytes = new Uint8Array(buffer);
    for (let i = 0; i < binaryDer.length; i++) { bytes[i] = binaryDer.charCodeAt(i); }
    return buffer;
}

/**
 * Extracts the hidden data from an image buffer by perfectly mirroring the Python LSB embedding logic.
 * @param {Buffer} imageBuffer - The raw buffer of the PNG image.
 * @returns {Promise<Uint8Array>} A promise that resolves with the extracted crypto payload.
 */
async function extractDataFromImage(imageBuffer) {
    // 1. Ensure the image is processed as 3-channel RGB, just like the Python encryptor.
    // .removeAlpha() converts RGBA to RGB. .toColourspace('srgb') handles other formats.
    const { data: pixelData, info } = await sharp(imageBuffer)
        .removeAlpha()
        .toColourspace('srgb')
        .raw()
        .toBuffer({ resolveWithObject: true });

    if (info.channels !== 3) {
        throw new Error(`Image processing error: Expected 3 (RGB) channels, but got ${info.channels}.`);
    }

    // 2. Read the first 32 bits (4 bytes) from the LSB stream to get the data length.
    const HEADER_BITS = 32;
    if (pixelData.length < HEADER_BITS) {
        throw new Error("Image is too small to contain a valid LSB header.");
    }

    let headerBinaryString = '';
    for (let i = 0; i < HEADER_BITS; i++) {
        headerBinaryString += pixelData[i] & 1;
    }

    const headerBytes = new Uint8Array(HEADER_BITS / 8);
    for (let i = 0; i < HEADER_BITS / 8; i++) {
        headerBytes[i] = parseInt(headerBinaryString.substring(i * 8, (i + 1) * 8), 2);
    }
    const dataLengthInBytes = new DataView(headerBytes.buffer).getUint32(0, false); // Big-endian

    if (dataLengthInBytes === 0) return new Uint8Array(0);

    // 3. Read the main data payload from the LSB stream.
    const dataLengthInBits = dataLengthInBytes * 8;
    const startOffset = HEADER_BITS;
    const endOffset = startOffset + dataLengthInBits;

    if (pixelData.length < endOffset) {
        throw new Error("Image data is corrupt or truncated: Header specifies a length greater than the available pixels.");
    }

    let dataBinaryString = '';
    for (let i = startOffset; i < endOffset; i++) {
        dataBinaryString += pixelData[i] & 1;
    }

    const cryptoPayload = new Uint8Array(dataLengthInBytes);
    for (let i = 0; i < dataLengthInBytes; i++) {
        cryptoPayload[i] = parseInt(dataBinaryString.substring(i * 8, (i + 1) * 8), 2);
    }

    return cryptoPayload;
}


/**
 * Decrypts the hybrid RSA-AES payload. This function remains the same as it was correct.
 * @param {Uint8Array} cryptoPayload - The extracted encrypted payload.
 * @param {string} privateKeyPem - The server's private key.
 * @returns {Promise<object>} A promise that resolves with the decrypted data.
 */
async function decryptHybridPayload(cryptoPayload, privateKeyPem) {
    const privateKey = await webcrypto.subtle.importKey('pkcs8', pemToDer(privateKeyPem), { name: 'RSA-OAEP', hash: 'SHA-256' }, true, ['decrypt']);
    const encryptedAesKeyLen = new DataView(cryptoPayload.buffer, 0, 4).getUint32(0, false);
    
    let offset = 4;
    const encryptedAesKey = cryptoPayload.slice(offset, offset + encryptedAesKeyLen);
    offset += encryptedAesKeyLen;
    const nonce = cryptoPayload.slice(offset, offset + 12);
    offset += 12;
    const ciphertextWithTag = cryptoPayload.slice(offset);
    
    const decryptedAesKeyBytes = await webcrypto.subtle.decrypt({ name: 'RSA-OAEP' }, privateKey, encryptedAesKey);
    const aesKey = await webcrypto.subtle.importKey('raw', decryptedAesKeyBytes, { name: 'AES-GCM', length: 256 }, true, ['decrypt']);
    const decryptedDataBuffer = await webcrypto.subtle.decrypt({ name: 'AES-GCM', iv: nonce }, aesKey, ciphertextWithTag);
    
    return JSON.parse(new TextDecoder().decode(decryptedDataBuffer));
}


/**
 * Main public function. Decodes auth data from an image buffer.
 * @param {Buffer} imageBuffer - The raw buffer of the uploaded image.
 * @param {string} privateKeyPem - The PEM-formatted private key.
 * @returns {Promise<object>} A promise that resolves to the decoded credentials object.
 */
async function decodeFromImageBuffer(imageBuffer, privateKeyPem) {
    const cryptoPayload = await extractDataFromImage(imageBuffer);
    if (cryptoPayload.length === 0) {
        // This is a valid case if an empty message was embedded.
        return {}; 
    }
    return await decryptHybridPayload(cryptoPayload, privateKeyPem);
}

module.exports = { decodeFromImageBuffer };