package com.peachpasswords.vault

import javax.crypto.Mac
import javax.crypto.spec.SecretKeySpec

/**
 * RFC 6238 TOTP / RFC 4226 HOTP. Port of the browser totp.ts behaviour:
 * 6 digits, 30-second period, SHA1 by default. The `totp` field on an entry
 * may hold either a raw base32 secret or a full `otpauth://` URI.
 */
object TOTPGenerator {
    const val PERIOD_SECONDS = 30
    private const val DIGITS = 6
    private const val BASE32 = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567"

    data class Otp(val code: String, val remaining: Int)

    fun generate(totpField: String, nowMillis: Long = System.currentTimeMillis()): Otp? {
        val parsed = parse(totpField) ?: return null
        val key = base32Decode(parsed.secret) ?: return null
        if (key.isEmpty()) return null
        val counter = (nowMillis / 1000L) / PERIOD_SECONDS
        val code = hotp(key, counter, parsed.algorithm)
        return Otp(code, remainingSeconds(nowMillis))
    }

    /** Just the code, or null if the secret is unparseable. */
    fun generateCode(totpField: String, nowMillis: Long = System.currentTimeMillis()): String? =
        generate(totpField, nowMillis)?.code

    fun remainingSeconds(nowMillis: Long = System.currentTimeMillis()): Int =
        PERIOD_SECONDS - ((nowMillis / 1000L) % PERIOD_SECONDS).toInt()

    /** Fraction of the current period elapsed, 0f → 1f. */
    fun periodProgress(nowMillis: Long = System.currentTimeMillis()): Float =
        ((nowMillis / 1000L) % PERIOD_SECONDS).toFloat() / PERIOD_SECONDS

    private data class Parsed(val secret: String, val algorithm: String)

    private fun parse(field: String): Parsed? {
        val trimmed = field.trim()
        if (trimmed.isEmpty()) return null
        if (!trimmed.startsWith("otpauth://", ignoreCase = true)) {
            return Parsed(trimmed.replace(" ", ""), "SHA1")
        }
        // otpauth://totp/label?secret=XXX&algorithm=SHA1&digits=6&period=30
        val query = trimmed.substringAfter('?', "")
        val params = query.split('&').mapNotNull {
            val i = it.indexOf('='); if (i <= 0) null else it.substring(0, i) to it.substring(i + 1)
        }.toMap()
        val secret = params["secret"] ?: return null
        val alg = when (params["algorithm"]?.uppercase()) {
            "SHA256" -> "SHA256"; "SHA512" -> "SHA512"; else -> "SHA1"
        }
        return Parsed(secret.replace(" ", ""), alg)
    }

    private fun hotp(key: ByteArray, counter: Long, algorithm: String): String {
        val msg = ByteArray(8)
        var c = counter
        for (i in 7 downTo 0) { msg[i] = (c and 0xFF).toByte(); c = c shr 8 }
        val macAlg = when (algorithm) { "SHA256" -> "HmacSHA256"; "SHA512" -> "HmacSHA512"; else -> "HmacSHA1" }
        val mac = Mac.getInstance(macAlg).apply { init(SecretKeySpec(key, macAlg)) }
        val hash = mac.doFinal(msg)
        val offset = (hash[hash.size - 1].toInt() and 0x0F)
        val binary = ((hash[offset].toInt() and 0x7F) shl 24) or
            ((hash[offset + 1].toInt() and 0xFF) shl 16) or
            ((hash[offset + 2].toInt() and 0xFF) shl 8) or
            (hash[offset + 3].toInt() and 0xFF)
        val otp = binary % 1_000_000
        return otp.toString().padStart(DIGITS, '0')
    }

    private fun base32Decode(input: String): ByteArray? {
        val clean = input.uppercase().replace("=", "").replace(" ", "")
        if (clean.isEmpty()) return null
        val out = ArrayList<Byte>(clean.length * 5 / 8)
        var buffer = 0
        var bits = 0
        for (ch in clean) {
            val v = BASE32.indexOf(ch)
            if (v < 0) return null
            buffer = (buffer shl 5) or v
            bits += 5
            if (bits >= 8) {
                bits -= 8
                out.add(((buffer shr bits) and 0xFF).toByte())
            }
        }
        return out.toByteArray()
    }
}