/**
 * E2E Encryption Service using TweetNaCl
 * https://tweetnacl.js.org/
 * 
 * Secret keys are encrypted at rest using AES-GCM with a key
 * derived from the user's UUID via PBKDF2.
 */

const STORAGE_KEY = 'dgray_keypair'
const SALT_KEY = 'dgray_keypair_salt'

class CryptoService {
    constructor() {
        this.nacl = null
        this.naclUtil = null
        this.keyPair = null
        this.ready = this.init()
    }

    async init() {
        if (!window.nacl) {
            await this.loadScript('/js/vendor/nacl-fast.min.js')
            await this.loadScript('/js/vendor/nacl-util.min.js')
        }

        this.nacl = window.nacl
        this.naclUtil = window.nacl.util
    }

    loadScript(src) {
        return new Promise((resolve, reject) => {
            if (document.querySelector(`script[src="${src}"]`)) {
                resolve()
                return
            }
            const script = document.createElement('script')
            script.src = src
            script.onload = resolve
            script.onerror = reject
            document.head.appendChild(script)
        })
    }

    async deriveKey(uuid) {
        const enc = new TextEncoder()
        let salt = localStorage.getItem(SALT_KEY)
        if (!salt) {
            const saltBytes = crypto.getRandomValues(new Uint8Array(16))
            salt = Array.from(saltBytes).map(b => b.toString(16).padStart(2, '0')).join('')
            localStorage.setItem(SALT_KEY, salt)
        }
        const saltBytes = new Uint8Array(salt.match(/.{2}/g).map(h => parseInt(h, 16)))

        const baseKey = await crypto.subtle.importKey(
            'raw', enc.encode(uuid), 'PBKDF2', false, ['deriveKey']
        )
        return crypto.subtle.deriveKey(
            { name: 'PBKDF2', salt: saltBytes, iterations: 100000, hash: 'SHA-256' },
            baseKey,
            { name: 'AES-GCM', length: 256 },
            false,
            ['encrypt', 'decrypt']
        )
    }

    async encryptForStorage(data, wrappingKey) {
        const iv = crypto.getRandomValues(new Uint8Array(12))
        const enc = new TextEncoder()
        const ciphertext = await crypto.subtle.encrypt(
            { name: 'AES-GCM', iv },
            wrappingKey,
            enc.encode(JSON.stringify(data))
        )
        return {
            iv: Array.from(iv).map(b => b.toString(16).padStart(2, '0')).join(''),
            ct: Array.from(new Uint8Array(ciphertext)).map(b => b.toString(16).padStart(2, '0')).join('')
        }
    }

    async decryptFromStorage(stored, wrappingKey) {
        const iv = new Uint8Array(stored.iv.match(/.{2}/g).map(h => parseInt(h, 16)))
        const ct = new Uint8Array(stored.ct.match(/.{2}/g).map(h => parseInt(h, 16)))
        const plainBuf = await crypto.subtle.decrypt(
            { name: 'AES-GCM', iv },
            wrappingKey,
            ct
        )
        return JSON.parse(new TextDecoder().decode(plainBuf))
    }

    async unlock(uuid) {
        await this.ready
        const wrappingKey = await this.deriveKey(uuid)

        const stored = localStorage.getItem(STORAGE_KEY)
        if (stored) {
            try {
                const parsed = JSON.parse(stored)

                if (parsed.ct && parsed.iv) {
                    const data = await this.decryptFromStorage(parsed, wrappingKey)
                    this.keyPair = {
                        publicKey: this.naclUtil.decodeBase64(data.publicKey),
                        secretKey: this.naclUtil.decodeBase64(data.secretKey)
                    }
                    return
                }

                if (parsed.publicKey && parsed.secretKey) {
                    this.keyPair = {
                        publicKey: this.naclUtil.decodeBase64(parsed.publicKey),
                        secretKey: this.naclUtil.decodeBase64(parsed.secretKey)
                    }
                    await this.saveKeyPair(wrappingKey)
                    return
                }
            } catch (e) {
                console.warn('Failed to load keypair, generating new one')
            }
        }

        await this.generateKeyPair(wrappingKey)
    }

    async generateKeyPair(wrappingKey) {
        this.keyPair = this.nacl.box.keyPair()
        await this.saveKeyPair(wrappingKey)
    }

    async saveKeyPair(wrappingKey) {
        const data = {
            publicKey: this.naclUtil.encodeBase64(this.keyPair.publicKey),
            secretKey: this.naclUtil.encodeBase64(this.keyPair.secretKey)
        }
        const encrypted = await this.encryptForStorage(data, wrappingKey)
        localStorage.setItem(STORAGE_KEY, JSON.stringify(encrypted))
    }

    lock() {
        this.keyPair = null
    }

    getPublicKey() {
        if (!this.keyPair) return null
        return this.naclUtil.encodeBase64(this.keyPair.publicKey)
    }

    /**
     * Encrypt a message for a recipient
     * @param {string} message - Plain text message
     * @param {string} recipientPublicKey - Base64 encoded public key
     * @returns {object} - { nonce, ciphertext } both base64 encoded
     */
    encrypt(message, recipientPublicKey) {
        const nonce = this.nacl.randomBytes(this.nacl.box.nonceLength)
        const messageUint8 = this.naclUtil.decodeUTF8(message)
        const recipientKey = this.naclUtil.decodeBase64(recipientPublicKey)
        const sharedKey = this.nacl.box.before(recipientKey, this.keyPair.secretKey)

        const encrypted = this.nacl.secretbox(messageUint8, nonce, sharedKey)

        return {
            nonce: this.naclUtil.encodeBase64(nonce),
            ciphertext: this.naclUtil.encodeBase64(encrypted)
        }
    }

    /**
     * Decrypt a message from a sender
     * @param {string} ciphertext - Base64 encoded ciphertext
     * @param {string} nonce - Base64 encoded nonce
     * @param {string} otherPublicKey - Base64 encoded public key of the other party
     * @returns {string|null} - Decrypted message or null if failed
     */
    decrypt(ciphertext, nonce, otherPublicKey) {
        try {
            const otherKey = this.naclUtil.decodeBase64(otherPublicKey)
            const sharedKey = this.nacl.box.before(otherKey, this.keyPair.secretKey)

            const decrypted = this.nacl.secretbox.open(
                this.naclUtil.decodeBase64(ciphertext),
                this.naclUtil.decodeBase64(nonce),
                sharedKey
            )

            if (!decrypted) return null
            return this.naclUtil.encodeUTF8(decrypted)
        } catch (e) {
            console.error('Decryption failed:', e)
            return null
        }
    }
}

export const cryptoService = new CryptoService()
