const crypto = require("crypto");

const MAGIC = Buffer.from("STEG");
const VERSION = 1;

const FLAG_ENCRYPTED = 1;

const HEADER_SIZE =
    4 + // magic
    1 + // version
    1 + // flags
    4;  // payload length

const IV_SIZE = 12;
const AUTH_TAG_SIZE = 16;

function getCapacity(width, height) {
    const availableBits =
        width * height * 3;

    const availableBytes =
        Math.floor(
            availableBits / 8
        );

    return Math.max(
        0,
        availableBytes - HEADER_SIZE
    );
}

function deriveKey(password) {
    return crypto
        .createHash("sha256")
        .update(password)
        .digest();
}

function encrypt(data, password) {
    const key =
        deriveKey(password);

    const iv =
        crypto.randomBytes(
            IV_SIZE
        );

    const cipher =
        crypto.createCipheriv(
            "aes-256-gcm",
            key,
            iv
        );

    const encrypted =
        Buffer.concat([
            cipher.update(data),
            cipher.final()
        ]);

    const authTag =
        cipher.getAuthTag();

    return Buffer.concat([
        iv,
        authTag,
        encrypted
    ]);
}

function decrypt(data, password) {
    if (
        data.length <
        IV_SIZE + AUTH_TAG_SIZE
    ) {
        throw new Error(
            "Invalid password or corrupted payload"
        );
    }

    const iv =
        data.subarray(
            0,
            IV_SIZE
        );

    const authTag =
        data.subarray(
            IV_SIZE,
            IV_SIZE + AUTH_TAG_SIZE
        );

    const encrypted =
        data.subarray(
            IV_SIZE + AUTH_TAG_SIZE
        );

    const key =
        deriveKey(password);

    try {
        const decipher =
            crypto.createDecipheriv(
                "aes-256-gcm",
                key,
                iv
            );

        decipher.setAuthTag(
            authTag
        );

        return Buffer.concat([
            decipher.update(encrypted),
            decipher.final()
        ]);
    } catch {
        throw new Error(
            "Invalid password or corrupted payload"
        );
    }
}

function setBit(
    buffer,
    bitIndex,
    value
) {
    const byteIndex =
        Math.floor(
            bitIndex / 8
        );

    const bit =
        7 - (bitIndex % 8);

    if (value) {
        buffer[byteIndex] |=
            1 << bit;
    } else {
        buffer[byteIndex] &=
            ~(1 << bit);
    }
}

function getBit(
    buffer,
    bitIndex
) {
    const byteIndex =
        Math.floor(
            bitIndex / 8
        );

    const bit =
        7 - (bitIndex % 8);

    return (
        (buffer[byteIndex] >> bit) &
        1
    );
}

function embed(
    pixels,
    payload
) {
    const output =
        Buffer.from(pixels);

    const totalBits =
        payload.length * 8;

    const availableBits =
        pixels.length -
        Math.floor(
            pixels.length / 4
        );

    if (
        totalBits >
        availableBits
    ) {
        throw new Error(
            "Payload too large"
        );
    }

    let bitIndex = 0;

    for (
        let pixelIndex = 0;
        pixelIndex < output.length;
        pixelIndex++
    ) {
        /*
         * Skip alpha channel.
         */
        if (
            pixelIndex % 4 === 3
        ) {
            continue;
        }

        if (
            bitIndex >= totalBits
        ) {
            break;
        }

        const value =
            getBit(
                payload,
                bitIndex
            );

        output[pixelIndex] =
            (output[pixelIndex] & 0xfe) |
            value;

        bitIndex++;
    }

    return output;
}

function extract(
    pixels,
    byteLength
) {
    const result =
        Buffer.alloc(
            byteLength
        );

    let bitIndex = 0;

    const totalBits =
        byteLength * 8;

    for (
        let pixelIndex = 0;
        pixelIndex < pixels.length;
        pixelIndex++
    ) {
        if (
            pixelIndex % 4 === 3
        ) {
            continue;
        }

        if (
            bitIndex >= totalBits
        ) {
            break;
        }

        const value =
            pixels[pixelIndex] & 1;

        setBit(
            result,
            bitIndex,
            value
        );

        bitIndex++;
    }

    return result;
}

function buildPayload({
    message,
    file,
    password
}) {
    let data;
    let type;

    if (
        file &&
        file.buffer
    ) {
        type = 1;

        const filename =
            Buffer.from(
                file.originalname || "",
                "utf8"
            );

        const mimetype =
            Buffer.from(
                file.mimetype || "application/octet-stream",
                "utf8"
            );

        const filenameLength =
            Buffer.alloc(2);

        filenameLength.writeUInt16BE(
            filename.length
        );

        const mimetypeLength =
            Buffer.alloc(2);

        mimetypeLength.writeUInt16BE(
            mimetype.length
        );

        data =
            Buffer.concat([
                Buffer.from([type]),
                filenameLength,
                filename,
                mimetypeLength,
                mimetype,
                file.buffer
            ]);
    } else if (
        typeof message === "string"
    ) {
        type = 0;

        const messageBuffer =
            Buffer.from(
                message,
                "utf8"
            );

        data =
            Buffer.concat([
                Buffer.from([type]),
                messageBuffer
            ]);
    } else {
        throw new Error(
            "Message or file is required"
        );
    }

    let flags = 0;

    if (password) {
        flags |= FLAG_ENCRYPTED;
        data = encrypt(
            data,
            password
        );
    }

    const header =
        Buffer.alloc(
            HEADER_SIZE
        );

    MAGIC.copy(
        header,
        0
    );

    header.writeUInt8(
        VERSION,
        4
    );

    header.writeUInt8(
        flags,
        5
    );

    header.writeUInt32BE(
        data.length,
        6
    );

    return Buffer.concat([
        header,
        data
    ]);
}

function parsePayload(
    payload,
    password
) {
    if (
        payload.length <
        HEADER_SIZE
    ) {
        throw new Error(
            "No valid steganography payload found"
        );
    }

    const magic =
        payload.subarray(
            0,
            4
        );

    if (
        !magic.equals(MAGIC)
    ) {
        throw new Error(
            "No valid steganography payload found"
        );
    }

    const version =
        payload.readUInt8(4);

    if (
        version !== VERSION
    ) {
        throw new Error(
            "Unsupported payload version"
        );
    }

    const flags =
        payload.readUInt8(5);

    const length =
        payload.readUInt32BE(6);

    if (
        payload.length <
        HEADER_SIZE + length
    ) {
        throw new Error(
            "Invalid password or corrupted payload"
        );
    }

    let data =
        payload.subarray(
            HEADER_SIZE,
            HEADER_SIZE + length
        );

    const encrypted =
        (flags & FLAG_ENCRYPTED) !== 0;

    if (encrypted) {
        if (!password) {
            throw new Error(
                "Password is required for this payload"
            );
        }

        data =
            decrypt(
                data,
                password
            );
    }

    if (
        data.length < 1
    ) {
        throw new Error(
            "Invalid password or corrupted payload"
        );
    }

    const type =
        data.readUInt8(0);

    if (type === 0) {
        return {
            type: "message",
            message:
                data
                    .subarray(1)
                    .toString("utf8")
        };
    }

    if (type === 1) {
        let offset = 1;

        if (
            data.length <
            offset + 2
        ) {
            throw new Error(
                "Invalid password or corrupted payload"
            );
        }

        const filenameLength =
            data.readUInt16BE(
                offset
            );

        offset += 2;

        if (
            data.length <
            offset + filenameLength
        ) {
            throw new Error(
                "Invalid password or corrupted payload"
            );
        }

        const filename =
            data
                .subarray(
                    offset,
                    offset + filenameLength
                )
                .toString("utf8");

        offset += filenameLength;

        if (
            data.length <
            offset + 2
        ) {
            throw new Error(
                "Invalid password or corrupted payload"
            );
        }

        const mimetypeLength =
            data.readUInt16BE(
                offset
            );

        offset += 2;

        if (
            data.length <
            offset + mimetypeLength
        ) {
            throw new Error(
                "Invalid password or corrupted payload"
            );
        }

        const mimetype =
            data
                .subarray(
                    offset,
                    offset + mimetypeLength
                )
                .toString("utf8");

        offset += mimetypeLength;

        return {
            type: "file",
            filename,
            mimetype,
            file:
                Buffer.from(
                    data.subarray(
                        offset
                    )
                )
        };
    }

    throw new Error(
        "Unknown payload type"
    );
}

async function encode({
    pixels,
    width,
    height,
    message,
    password,
    file
}) {
    if (!pixels) {
        throw new Error(
            "Image data is required"
        );
    }

    const payload =
        buildPayload({
            message,
            file,
            password
        });

    const capacity =
        getCapacity(
            width,
            height
        );

    if (
        payload.length >
        capacity
    ) {
        throw new Error(
            `Payload too large. ` +
            `Maximum size is ${capacity} bytes.`
        );
    }

    return {
        pixels:
            embed(
                pixels,
                payload
            )
    };
}

async function decode({
    pixels,
    password
}) {
    if (!pixels) {
        throw new Error(
            "Image data is required"
        );
    }

    /*
     * Read the header first.
     */
    const header =
        extract(
            pixels,
            HEADER_SIZE
        );

    if (
        !header
            .subarray(0, 4)
            .equals(MAGIC)
    ) {
        throw new Error(
            "No valid steganography payload found"
        );
    }

    const length =
        header.readUInt32BE(6);

    const totalLength =
        HEADER_SIZE + length;

    const payload =
        extract(
            pixels,
            totalLength
        );

    return parsePayload(
        payload,
        password
    );
}

module.exports = {
    encode,
    decode,
    getCapacity
};