package com.vgmlr.wedge

import android.util.Base64
import java.security.SecureRandom
import javax.crypto.Cipher
import javax.crypto.SecretKeyFactory
import javax.crypto.spec.GCMParameterSpec
import javax.crypto.spec.PBEKeySpec
import javax.crypto.spec.SecretKeySpec
import kotlin.collections.sliceArray
import android.content.Context

object WedgeSecurity {
    private const val ITERATIONS = 35690
    private const val KEY_LENGTH = 256
    private const val SALT_LENGTH = 16
    private const val IV_LENGTH = 12
    private const val PREFIX = "wgd:"
    private const val DIGITS = "0123456789"
    private const val WORDLIST = "wordlist.txt"
    private val WS = Regex("\\s+")
    @Volatile private var words: List<String>? = null
    
    fun isEncrypted(text: String): Boolean = text.startsWith(PREFIX)

    fun encrypt(plainText: String, passPhrase: String): String {
        val random = SecureRandom()
        val salt = ByteArray(SALT_LENGTH).apply { random.nextBytes(this) }
        val iv = ByteArray(IV_LENGTH).apply { random.nextBytes(this) }

        val factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA256")
        val spec = PBEKeySpec(passPhrase.toCharArray(), salt, ITERATIONS, KEY_LENGTH)
        val tmp = factory.generateSecret(spec)
        val secretKey = SecretKeySpec(tmp.encoded, "AES")

        val cipher = Cipher.getInstance("AES/GCM/NoPadding")
        cipher.init(Cipher.ENCRYPT_MODE, secretKey, GCMParameterSpec(128, iv))
        val cipherText = cipher.doFinal(plainText.toByteArray(Charsets.UTF_8))

        val combined = salt + iv + cipherText
        return PREFIX + Base64.encodeToString(combined, Base64.NO_WRAP)
    }

    fun decrypt(cipherTextBase64: String, passPhrase: String): String? {
        if (!isEncrypted(cipherTextBase64)) return null
        return try {
            val cleanInput = cipherTextBase64.removePrefix(PREFIX)
            val combined = Base64.decode(cleanInput, Base64.DEFAULT)

            val salt = combined.sliceArray(0 until SALT_LENGTH)
            val iv = combined.sliceArray(SALT_LENGTH until SALT_LENGTH + IV_LENGTH)
            val encrypted = combined.sliceArray(SALT_LENGTH + IV_LENGTH until combined.size)

            val factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA256")
            val spec = PBEKeySpec(passPhrase.toCharArray(), salt, ITERATIONS, KEY_LENGTH)
            val secretKey = SecretKeySpec(factory.generateSecret(spec).encoded, "AES")

            val cipher = Cipher.getInstance("AES/GCM/NoPadding")
            cipher.init(Cipher.DECRYPT_MODE, secretKey, GCMParameterSpec(128, iv))

            String(cipher.doFinal(encrypted), Charsets.UTF_8)
        } catch (_: Exception) { null }
    }

    private fun wordList(context: Context): List<String> = words ?: synchronized(this) {
        words ?: context.assets.open(WORDLIST).bufferedReader().useLines { lines ->
            lines.mapNotNull { it.trim().split(WS).lastOrNull() }
                .filter { w -> w.isNotEmpty() && w.all { it in 'a'..'z' } }
                .toList()
        }.also { words = it }
    }

    fun generatePhrase(context: Context, spec: PhraseSpec): String {
        val rnd = SecureRandom()
        val useWords = spec.words || !spec.numbers
        val parts = mutableListOf<String>()
        if (useWords) {
            val all = wordList(context)
            val max = spec.letters.coerceIn(3, 9)
            val pool = all.filter { it.length <= max }.ifEmpty { all }
            repeat(spec.wordCount.coerceIn(1, 5)) { parts += pool[rnd.nextInt(pool.size)] }
        }
        if (spec.numbers) parts += buildString {
            repeat(spec.numberCount.coerceIn(1, 10)) { append(DIGITS[rnd.nextInt(DIGITS.length)]) }
        }
        return parts.joinToString(" ")
    }
}