package invalid.lena.rsend

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

// 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 knownHosts(ctx: Context): File = File(ctx.filesDir, "known_hosts")
    const val MAX_KEY_BYTES = 64 * 1024
    private const val MAX_ENCRYPTED_KEY_BYTES = MAX_KEY_BYTES + 64
    private val keyLock = Any()

    fun exists(ctx: Context): Boolean = synchronized(keyLock) {
        try {
            AtomicFile(keyEnc(ctx)).openRead().use { }
            true
        } catch (_: FileNotFoundException) {
            false
        }
    }

    // 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 = synchronized(keyLock) {
        val result = Native.run(Native.rsh(ctx), listOf("-keygen"))
        if (result.code != 0) {
            throw IllegalStateException(result.output.trim().ifEmpty { "keygen failed" })
        }
        validateAndStore(ctx, result.output.toByteArray(Charsets.UTF_8))
    }

    // 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): String = synchronized(keyLock) {
        validateAndStore(ctx, pem)
    }

    private fun validateAndStore(ctx: Context, pem: ByteArray): String {
        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, Charsets.UTF_8)),
        )
        if (r.code != 0) throw IllegalArgumentException(r.output.trim().ifEmpty { "invalid private key" })
        atomicWrite(keyEnc(ctx), KeyVault.encrypt(pem))
        return r.output.trim()
    }

    fun publicKeyText(ctx: Context): String = synchronized(keyLock) {
        if (!exists(ctx)) return ""
        val r = Native.run(Native.rsh(ctx), listOf("-pubkey"), keyEnvironment(ctx))
        if (r.code != 0) throw IllegalStateException(r.output.trim().ifEmpty { "could not read public key" })
        r.output.trim()
    }

    // env derives a one-pin known_hosts file for this connection. A shared file
    // containing every remote's key would let one configured alias authorize a
    // different alias of the same endpoint.
    fun env(ctx: Context, remote: Remote): Map<String, String> = synchronized(keyLock) {
        require(RemoteRules.hostKeyAllowed(remote)) { "remote pin is invalid" }
        require(remote.pinned()) { "remote is not pinned" }
        val hosts = knownHosts(ctx)
        atomicWrite(hosts, (remote.hostKey.trim() + "\n").toByteArray(Charsets.UTF_8))
        keyEnvironment(ctx) + mapOf(
            "RSH_KNOWN_HOSTS" to hosts.absolutePath,
            "RSH_PORT" to remote.port.toString(),
        )
    }

    // legacyPin returns the known_hosts line 0.1.x kept beside the config, back
    // when one pin was all there was. Config reads it once while migrating a
    // 0.1.x config; nothing else should, because this path is now a scratch
    // file rewritten per connection. Delete with that migration in 0.4.0.
    fun legacyPin(ctx: Context): String {
        val f = knownHosts(ctx)
        if (!f.exists() || f.length() > MAX_LEGACY_PIN_BYTES) return ""
        return try {
            f.readText(Charsets.UTF_8).lineSequence().firstOrNull { it.isNotBlank() }?.trim().orEmpty()
        } catch (_: Exception) {
            ""
        }
    }

    private const val MAX_LEGACY_PIN_BYTES = 64L * 1024

    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 = synchronized(keyLock) {
        // Point the scan at this remote's own pin, so rsh offers the key type
        // already pinned and an unchanged server reproduces its stored line. A
        // config pinned by an older rsend holds whichever key that version
        // negotiated, and re-scanning without this would return a different key
        // type and look like a host-key change on an untouched server.
        val scanHosts = File(ctx.cacheDir, "scan_known_hosts")
        atomicWrite(
            scanHosts,
            if (remote.pinned()) (remote.hostKey + "\n").toByteArray(Charsets.UTF_8) else ByteArray(0),
        )
        val r = Native.run(
            Native.rsh(ctx),
            listOf("-scan", "${remote.user}@${remote.host}"),
            keyEnvironment(ctx) + mapOf(
                "RSH_KNOWN_HOSTS" to scanHosts.absolutePath,
                "RSH_PORT" to remote.port.toString(),
            ),
        )
        if (r.code != 0) return@synchronized Scan(false, "", "", r.output.trim())
        val lines = r.output.trim().lines()
        if (lines.size < 2) {
            return@synchronized Scan(false, "", "", "unexpected scan output:\n${r.output}")
        }
        Scan(true, lines[0], lines[1], "")
    }

    private fun keyEnvironment(ctx: Context): Map<String, String> =
        mapOf("RSH_KEY_DATA" to String(KeyVault.decrypt(encryptedKey(ctx)), Charsets.UTF_8))

    private fun encryptedKey(ctx: Context): ByteArray {
        val out = ByteArrayOutputStream()
        AtomicFile(keyEnc(ctx)).openRead().use { input ->
            val buf = ByteArray(4096)
            while (true) {
                val n = input.read(buf)
                if (n < 0) break
                require(out.size() + n <= MAX_ENCRYPTED_KEY_BYTES) { "encrypted private key is too large" }
                out.write(buf, 0, n)
            }
        }
        return out.toByteArray()
    }

    // 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
        }
    }
}
