package invalid.lena.rsend

import android.content.Context
import android.util.AtomicFile
import java.io.File

// Keys manages the on-device ed25519 identity and the pinned host keys. The
// private key is stored Keystore-encrypted (KeyVault) and only ever exists in
// plaintext in memory, passed to rsh through RSH_KEY_DATA.
object Keys {

    private fun keyEnc(ctx: Context): File = File(ctx.filesDir, "id_ed25519.enc")
    private fun publicKey(ctx: Context): File = File(ctx.filesDir, "id_ed25519.pub")
    const val MAX_KEY_BYTES = 64 * 1024

    fun knownHosts(ctx: Context): File = File(ctx.filesDir, "known_hosts")

    fun exists(ctx: Context): Boolean = keyEnc(ctx).exists()

    // generate creates a key pair via rsh -keygen (the plaintext key never
    // touches disk), stores it encrypted, and returns the public key in
    // authorized_keys format. Storage and validation reuse importKey.
    fun generate(ctx: Context): String {
        val r = Native.run(Native.rsh(ctx), listOf("-keygen"))
        if (r.code != 0) throw IllegalStateException(r.output.trim().ifEmpty { "keygen failed" })
        importKey(ctx, r.output.toByteArray())
        return publicKeyText(ctx)
    }

    // importKey stores a user-supplied private key, replacing any current one.
    // The bundled rsh derives and validates the public key (-pubkey uses the
    // same parser as the transport), so a key that imports here will also sync.
    // Throws IllegalArgumentException with rsh's message if the key is unusable
    // (wrong type, or passphrase-protected: rsend needs an unencrypted key).
    fun importKey(ctx: Context, pem: ByteArray) {
        require(pem.size <= MAX_KEY_BYTES) { "private key exceeds $MAX_KEY_BYTES bytes" }
        val r = Native.run(Native.rsh(ctx), listOf("-pubkey"), mapOf("RSH_KEY_DATA" to String(pem)))
        if (r.code != 0) throw IllegalArgumentException(r.output.trim().ifEmpty { "invalid private key" })
        atomicWrite(keyEnc(ctx), KeyVault.encrypt(pem))
        atomicWrite(publicKey(ctx), (r.output.trim() + "\n").toByteArray())
    }

    fun publicKeyText(ctx: Context): String =
        if (publicKey(ctx).exists()) publicKey(ctx).readText().trim() else ""

    // env returns the RSH_* environment rsh needs. The decrypted key is passed
    // by value so it never touches the filesystem.
    fun env(ctx: Context, port: Int): Map<String, String> = mapOf(
        "RSH_KEY_DATA" to String(KeyVault.decrypt(keyEnc(ctx).readBytes())),
        "RSH_KNOWN_HOSTS" to knownHosts(ctx).absolutePath,
        "RSH_PORT" to port.toString(),
    )

    data class Scan(val ok: Boolean, val fingerprint: String, val line: String, val error: String)

    // scan connects (proving the key is installed), captures the host key, and
    // returns its fingerprint and the known_hosts line to pin. pin() writes it
    // once the user accepts.
    fun scan(ctx: Context, remote: Remote): Scan {
        val r = Native.run(
            Native.rsh(ctx),
            listOf("-scan", "${remote.user}@${remote.host}"),
            env(ctx, remote.port),
        )
        if (r.code != 0) return Scan(false, "", "", r.output.trim())
        val lines = r.output.trim().lines()
        if (lines.size < 2) return Scan(false, "", "", "unexpected scan output:\n${r.output}")
        return Scan(true, lines[0], lines[1], "")
    }

    fun pin(ctx: Context, line: String) {
        atomicWrite(knownHosts(ctx), (line + "\n").toByteArray())
    }

    fun clearPin(ctx: Context) {
        knownHosts(ctx).delete()
    }

    // AtomicFile keeps the previous complete file if a write is interrupted.
    private fun atomicWrite(f: File, bytes: ByteArray) {
        val file = AtomicFile(f)
        val out = file.startWrite()
        try {
            out.write(bytes)
            file.finishWrite(out)
        } catch (e: Exception) {
            file.failWrite(out)
            throw e
        }
    }

    fun pinned(ctx: Context): Boolean = knownHosts(ctx).let { it.exists() && it.length() > 0 }
}
