// keylock.js

class KeyLock {

    constructor() {
        this.PREFERRED_FONTS = "bold 30px Arial, 'DejaVu Sans', Helvetica, sans-serif";
    }

    // --- Core Cryptography Methods ---

    /**
     * Converts a PEM-formatted key string to an ArrayBuffer.
     * @param {string} pem - The PEM string (e.g., -----BEGIN...-----).
     * @returns {ArrayBuffer} The raw binary data of the key.
     */
    _pemToBinary(pem) {
        const lines = pem.split('\n');
        const base64 = lines
            .filter(line => !line.startsWith('-----'))
            .join('');
        const binaryDer = window.atob(base64);
        const uint8Array = new Uint8Array(binaryDer.length);
        for (let i = 0; i < binaryDer.length; i++) {
            uint8Array[i] = binaryDer.charCodeAt(i);
        }
        return uint8Array.buffer;
    }
    
    /**
     * Converts an ArrayBuffer key to a PEM-formatted string.
     * @param {ArrayBuffer} buffer - The raw binary data of the key.
     * @param {string} label - The label for the PEM file (e.g., "PRIVATE KEY").
     * @returns {string} The PEM-formatted key string.
     */
    _binaryToPem(buffer, label) {
        const binaryStr = String.fromCharCode.apply(null, new Uint8Array(buffer));
        const base64 = window.btoa(binaryStr);
        let pem = `-----BEGIN ${label}-----\n`;
        for (let i = 0; i < base64.length; i += 64) {
            pem += base64.slice(i, i + 64) + '\n';
        }
        pem += `-----END ${label}-----`;
        return pem;
    }

    /**
     * Imports an RSA public key from PEM format into a CryptoKey object.
     * @param {string} pem - The public key in PEM format (SPKI).
     * @returns {Promise<CryptoKey>}
     */
    async importRsaPublicKey(pem) {
        return await window.crypto.subtle.importKey(
            'spki',
            this._pemToBinary(pem), {
                name: 'RSA-OAEP',
                hash: 'SHA-256'
            },
            true,
            ['encrypt']
        );
    }

    /**
     * Imports an RSA private key from PEM format into a CryptoKey object.
     * @param {string} pem - The private key in PEM format (PKCS#8).
     * @returns {Promise<CryptoKey>}
     */
    async importRsaPrivateKey(pem) {
        return await window.crypto.subtle.importKey(
            'pkcs8',
            this._pemToBinary(pem), {
                name: 'RSA-OAEP',
                hash: 'SHA-256'
            },
            true,
            ['decrypt']
        );
    }

    /**
     * Generates a new RSA-2048 key pair.
     * @returns {Promise<{privateKeyPem: string, publicKeyPem: string}>}
     */
    async generatePemKeys() {
        const keyPair = await window.crypto.subtle.generateKey({
                name: 'RSA-OAEP',
                modulusLength: 2048,
                publicExponent: new Uint8Array([1, 0, 1]), // 65537
                hash: 'SHA-256',
            },
            true,
            ['encrypt', 'decrypt']
        );

        const privateKeyDer = await window.crypto.subtle.exportKey('pkcs8', keyPair.privateKey);
        const publicKeyDer = await window.crypto.subtle.exportKey('spki', keyPair.publicKey);

        return {
            privateKeyPem: this._binaryToPem(privateKeyDer, 'PRIVATE KEY'),
            publicKeyPem: this._binaryToPem(publicKeyDer, 'PUBLIC KEY'),
        };
    }

    // --- Image Generation and Steganography ---

    /**
     * Generates a procedural starfield image on a canvas.
     * @param {number} w - Width of the canvas.
     * @param {number} h - Height of the canvas.
     * @returns {HTMLCanvasElement}
     */
    _generateStarfieldImage(w = 800, h = 800) {
        const canvas = document.createElement('canvas');
        canvas.width = w;
        canvas.height = h;
        const ctx = canvas.getContext('2d');

        // Background gradient
        const centerX = w / 2;
        const centerY = h / 2;
        const gradient = ctx.createRadialGradient(centerX, centerY, 0, centerX, centerY, w / 2);
        gradient.addColorStop(0, 'rgb(20, 25, 40)');
        gradient.addColorStop(1, 'rgb(0, 0, 5)');
        ctx.fillStyle = gradient;
        ctx.fillRect(0, 0, w, h);

        // Dim stars
        for (let i = 0; i < (w * h) / 200; i++) {
            const x = Math.random() * w;
            const y = Math.random() * h;
            const brightness = 30 + Math.random() * 60;
            ctx.fillStyle = `rgb(${Math.floor(brightness * 0.9)}, ${Math.floor(brightness * 0.9)}, ${brightness})`;
            ctx.fillRect(x, y, 1, 1);
        }

        // Bright stars with glow
        const starColors = ['rgb(255, 255, 255)', 'rgb(220, 230, 255)', 'rgb(255, 240, 220)'];
        for (let i = 0; i < (w * h) / 1000; i++) {
            const x = Math.random() * w;
            const y = Math.random() * h;
            const size = 0.5 + 2.5 * (Math.random() ** 2);
            const brightness = 120 + 135 * (Math.random() ** 1.5);
            const color = starColors[Math.floor(Math.random() * starColors.length)];

            ctx.beginPath();
            const glowGradient = ctx.createRadialGradient(x, y, 0, x, y, size * 3);
            glowGradient.addColorStop(0, color.replace(')', `, ${brightness / 255})`).replace('rgb', 'rgba'));
            glowGradient.addColorStop(1, color.replace(')', ', 0)').replace('rgb', 'rgba'));
            ctx.fillStyle = glowGradient;
            ctx.arc(x, y, size * 3, 0, 2 * Math.PI);
            ctx.fill();

            ctx.beginPath();
            ctx.fillStyle = color.replace(')', `, ${brightness / 255})`).replace('rgb', 'rgba');
            ctx.arc(x, y, size, 0, 2 * Math.PI);
            ctx.fill();
        }
        return canvas;
    }
    
    /**
     * Draws the text overlay on the image.
     * @param {HTMLCanvasElement} canvas - The canvas to draw on.
     * @returns {HTMLCanvasElement} The same canvas, now with an overlay.
     */
    _drawOverlay(canvas) {
        const ctx = canvas.getContext('2d');
        const { width, height } = canvas;
        
        ctx.fillStyle = 'rgba(10, 15, 30, 0.78)'; // 200/255 alpha
        ctx.fillRect(0, 20, width, 60);

        ctx.fillStyle = 'rgb(200, 220, 255)';
        ctx.font = this.PREFERRED_FONTS;
        ctx.textAlign = 'center';
        ctx.textBaseline = 'middle';
        ctx.fillText("KeyLock Secure Data", width / 2, 50);

        return canvas;
    }

    /**
     * Parses a key-value string into a JavaScript object.
     * @param {string} kvString - The input string (e.g., 'USER="test"\nPASS:123').
     * @returns {object}
     */
    _parseKvString(kvString) {
        const payload = {};
        if (!kvString) return payload;
        
        const lines = kvString.trim().split('\n');
        for (const line of lines) {
            const trimmedLine = line.trim();
            if (!trimmedLine || trimmedLine.startsWith('#')) continue;

            const parts = trimmedLine.split(/[:=]/, 2);
            if (parts.length === 2) {
                let [key, value] = parts;
                key = key.trim().replace(/^['"]|['"]$/g, '');
                value = value.trim().replace(/^['"]|['"]$/g, '');
                if (key) {
                    payload[key] = value;
                }
            }
        }
        return payload;
    }
    
    // --- Main Public Methods: Encoding and Decoding ---

    /**
     * Generates an encrypted image containing the payload.
     * @param {string} payloadKvString - The key-value data to encrypt.
     * @param {string} publicKeyPem - The RSA public key in PEM format.
     * @returns {Promise<string>} A data URL of the generated PNG image.
     */
    async generateEncryptedImage(payloadKvString, publicKeyPem) {
        const payloadDict = this._parseKvString(payloadKvString);
        if (Object.keys(payloadDict).length === 0) {
            throw new Error("Payload is empty or could not be parsed.");
        }
        
        // 1. Prepare crypto primitives
        const publicKey = await this.importRsaPublicKey(publicKeyPem);
        const aesKey = await window.crypto.subtle.generateKey({ name: 'AES-GCM', length: 256 }, true, ['encrypt', 'decrypt']);
        const nonce = window.crypto.getRandomValues(new Uint8Array(12));

        // 2. Encrypt payload with AES
        const jsonBytes = new TextEncoder().encode(JSON.stringify(payloadDict));
        const ciphertext = await window.crypto.subtle.encrypt({ name: 'AES-GCM', iv: nonce }, aesKey, jsonBytes);

        // 3. Encrypt AES key with RSA
        const exportedAesKey = await window.crypto.subtle.exportKey('raw', aesKey);
        const encryptedAesKey = await window.crypto.subtle.encrypt({ name: 'RSA-OAEP' }, publicKey, exportedAesKey);

        // 4. Construct the final binary payload
        // Format: [4-byte len of RSA key][RSA key][12-byte nonce][AES ciphertext]
        const totalLength = 4 + encryptedAesKey.byteLength + nonce.byteLength + ciphertext.byteLength;
        const finalPayload = new Uint8Array(totalLength);
        const view = new DataView(finalPayload.buffer);
        
        let offset = 0;
        view.setUint32(offset, encryptedAesKey.byteLength, false); // Big-endian
        offset += 4;
        finalPayload.set(new Uint8Array(encryptedAesKey), offset);
        offset += encryptedAesKey.byteLength;
        finalPayload.set(nonce, offset);
        offset += nonce.byteLength;
        finalPayload.set(new Uint8Array(ciphertext), offset);

        // 5. Generate base image
        const canvas = this._generateStarfieldImage();
        this._drawOverlay(canvas);
        const ctx = canvas.getContext('2d');
        const imageData = ctx.getImageData(0, 0, canvas.width, canvas.height);
        const pixelData = imageData.data;

        // 6. Embed payload into image via LSB steganography
        // Format: [32-bit length header][payload bits]
        const payloadWithHeader = new Uint8Array(4 + finalPayload.length);
        const headerView = new DataView(payloadWithHeader.buffer);
        headerView.setUint32(0, finalPayload.length, false); // Big-endian
        payloadWithHeader.set(finalPayload, 4);

        let binaryPayload = '';
        payloadWithHeader.forEach(byte => {
            binaryPayload += byte.toString(2).padStart(8, '0');
        });

        if (binaryPayload.length > pixelData.length) {
             throw new Error("Payload is too large for the image.");
        }

        for (let i = 0; i < binaryPayload.length; i++) {
            pixelData[i] = (pixelData[i] & 0xFE) | parseInt(binaryPayload[i], 2);
        }

        // 7. Finalize and return image
        ctx.putImageData(imageData, 0, 0);
        return canvas.toDataURL('image/png');
    }

    /**
     * Decodes a payload hidden inside an image.
     * @param {HTMLImageElement} imageElement - The image containing the hidden data.
     * @param {string} privateKeyPem - The RSA private key in PEM format.
     * @returns {Promise<{status: string, payload?: object, message?: string}>}
     */
    async decodePayload(imageElement, privateKeyPem) {
        try {
            const privateKey = await this.importRsaPrivateKey(privateKeyPem);

            // 1. Extract pixel data from image
            const canvas = document.createElement('canvas');
            canvas.width = imageElement.naturalWidth;
            canvas.height = imageElement.naturalHeight;
            const ctx = canvas.getContext('2d');
            ctx.drawImage(imageElement, 0, 0);
            const pixelData = ctx.getImageData(0, 0, canvas.width, canvas.height).data;

            // 2. Extract binary data from LSBs
            let headerBinary = '';
            for (let i = 0; i < 32; i++) {
                headerBinary += (pixelData[i] & 1).toString();
            }
            const dataLength = parseInt(headerBinary, 2); // Length in bytes
            const requiredPixels = 32 + dataLength * 8;
            if (requiredPixels > pixelData.length) {
                throw new Error("Incomplete payload in image.");
            }
            
            let dataBinary = '';
            for (let i = 32; i < requiredPixels; i++) {
                dataBinary += (pixelData[i] & 1).toString();
            }

            // 3. Convert binary string to ArrayBuffer
            const cryptoPayload = new Uint8Array(dataLength);
            for (let i = 0; i < dataLength; i++) {
                cryptoPayload[i] = parseInt(dataBinary.substring(i * 8, (i + 1) * 8), 2);
            }

            // 4. Parse the crypto payload
            const view = new DataView(cryptoPayload.buffer);
            let offset = 0;
            const encryptedAesKeyLen = view.getUint32(offset, false); // Big-endian
            offset += 4;

            const encryptedAesKey = cryptoPayload.slice(offset, offset + encryptedAesKeyLen);
            offset += encryptedAesKeyLen;
            const nonce = cryptoPayload.slice(offset, offset + 12);
            offset += 12;
            const ciphertext = cryptoPayload.slice(offset);

            // 5. Decrypt AES key with RSA private key
            const recoveredAesKeyBytes = await window.crypto.subtle.decrypt({ name: 'RSA-OAEP' }, privateKey, encryptedAesKey);
            const recoveredAesKey = await window.crypto.subtle.importKey('raw', recoveredAesKeyBytes, { name: 'AES-GCM' }, true, ['decrypt']);

            // 6. Decrypt payload with AES key
            const decryptedPayloadBytes = await window.crypto.subtle.decrypt({ name: 'AES-GCM', iv: nonce }, recoveredAesKey, ciphertext);
            const payloadJson = new TextDecoder().decode(decryptedPayloadBytes);
            const payload = JSON.parse(payloadJson);

            return { status: "Success", payload: payload };

        } catch (e) {
            console.error("Decryption Failed:", e);
            return { status: "Error", message: `Decryption Failed: ${e.message}` };
        }
    }
}