package com.inou.clawdnode.security

import android.content.Context
import android.util.Base64
import android.util.Log
import androidx.security.crypto.EncryptedSharedPreferences
import androidx.security.crypto.MasterKey
import net.i2p.crypto.eddsa.EdDSAEngine
import net.i2p.crypto.eddsa.EdDSAPrivateKey
import net.i2p.crypto.eddsa.EdDSAPublicKey
import net.i2p.crypto.eddsa.KeyPairGenerator
import net.i2p.crypto.eddsa.spec.EdDSANamedCurveTable
import net.i2p.crypto.eddsa.spec.EdDSAPrivateKeySpec
import net.i2p.crypto.eddsa.spec.EdDSAPublicKeySpec
import java.security.MessageDigest
import java.security.SecureRandom

/**
 * Manages device identity using Ed25519 keys.
 * 
 * Keys are stored in EncryptedSharedPreferences for security.
 * Uses pure-Java Ed25519 implementation for compatibility with all Android versions.
 * 
 * The key format matches the Clawdbot gateway protocol:
 * - Public key: 32 bytes raw, base64url-encoded
 * - Signature: 64 bytes raw, base64url-encoded
 * - Payload: "$nonce:$signedAt" (matching gateway's expected format)
 */
class DeviceIdentity(context: Context) {
    
    private val tag = "DeviceIdentity"
    private val prefsName = "clawdnode_device_identity"
    private val keyPrivate = "private_key"
    private val keyPublic = "public_key"
    
    private val ed25519Spec = EdDSANamedCurveTable.getByName("Ed25519")
    
    private val prefs by lazy {
        val masterKey = MasterKey.Builder(context)
            .setKeyScheme(MasterKey.KeyScheme.AES256_GCM)
            .build()
        
        EncryptedSharedPreferences.create(
            context,
            prefsName,
            masterKey,
            EncryptedSharedPreferences.PrefKeyEncryptionScheme.AES256_SIV,
            EncryptedSharedPreferences.PrefValueEncryptionScheme.AES256_GCM
        )
    }
    
    /**
     * Get the device's public key as base64url-encoded raw bytes (32 bytes)
     */
    val publicKey: String by lazy {
        ensureKeyPair()
        val publicKeyBytes = prefs.getString(keyPublic, null)
            ?: throw IllegalStateException("No public key available")
        publicKeyBytes // Already stored as base64url
    }
    
    /**
     * Get device ID (SHA-256 fingerprint of the public key raw bytes)
     */
    val deviceId: String by lazy {
        ensureKeyPair()
        val publicKeyBase64 = prefs.getString(keyPublic, null)
            ?: throw IllegalStateException("No public key available")
        
        // Decode the raw public key bytes
        val publicKeyBytes = base64UrlDecode(publicKeyBase64)
        
        // Create fingerprint using SHA-256 of the raw public key bytes
        val digest = MessageDigest.getInstance("SHA-256")
        val hash = digest.digest(publicKeyBytes)
        
        // Return as hex string (full 32 bytes = 64 chars, matching gateway)
        hash.joinToString("") { "%02x".format(it) }
    }
    
    /**
     * Sign a challenge nonce for gateway authentication
     * 
     * @param nonce The challenge nonce from the gateway
     * @return Signature and timestamp
     */
    fun signChallenge(nonce: String): SignedChallenge {
        ensureKeyPair()
        
        val signedAt = System.currentTimeMillis()
        val payload = "$nonce:$signedAt"
        
        Log.d(tag, "Signing payload: $payload")
        
        // Load private key
        val privateKeyBase64 = prefs.getString(keyPrivate, null)
            ?: throw IllegalStateException("No private key available")
        val privateKeyBytes = base64UrlDecode(privateKeyBase64)
        
        // Create EdDSA private key
        val privateKeySpec = EdDSAPrivateKeySpec(privateKeyBytes, ed25519Spec)
        val privateKey = EdDSAPrivateKey(privateKeySpec)
        
        // Sign the payload
        val signature = EdDSAEngine(MessageDigest.getInstance(ed25519Spec.hashAlgorithm)).apply {
            initSign(privateKey)
            update(payload.toByteArray(Charsets.UTF_8))
        }
        
        val signatureBytes = signature.sign()
        val signatureBase64 = base64UrlEncode(signatureBytes)
        
        Log.d(tag, "Generated signature: ${signatureBase64.take(20)}... (${signatureBytes.size} bytes)")
        
        return SignedChallenge(
            signature = signatureBase64,
            signedAt = signedAt,
            nonce = nonce
        )
    }
    
    /**
     * Ensure the keypair exists, generating if necessary
     */
    private fun ensureKeyPair() {
        try {
            if (prefs.getString(keyPrivate, null) == null) {
                Log.i(tag, "No existing keypair found, generating new Ed25519 keypair")
                generateKeyPair()
            } else {
                Log.d(tag, "Using existing Ed25519 keypair")
            }
        } catch (e: Exception) {
            Log.e(tag, "Failed to check/ensure keypair", e)
            throw e
        }
    }
    
    /**
     * Generate a new Ed25519 keypair and store it
     */
    private fun generateKeyPair() {
        Log.i(tag, "Generating new Ed25519 device keypair")
        
        val keyPairGenerator = KeyPairGenerator()
        keyPairGenerator.initialize(256, SecureRandom())
        val keyPair = keyPairGenerator.generateKeyPair()
        
        val privateKey = keyPair.private as EdDSAPrivateKey
        val publicKey = keyPair.public as EdDSAPublicKey
        
        // Get raw 32-byte keys (matching gateway format)
        // For Ed25519, the "seed" is the 32-byte private key
        val privateKeyBytes = privateKey.seed
        // The "A" point is the 32-byte public key
        val publicKeyBytes = publicKey.abyte
        
        Log.d(tag, "Generated keypair: private=${privateKeyBytes.size} bytes, public=${publicKeyBytes.size} bytes")
        
        // Store as base64url
        prefs.edit()
            .putString(keyPrivate, base64UrlEncode(privateKeyBytes))
            .putString(keyPublic, base64UrlEncode(publicKeyBytes))
            .apply()
        
        Log.i(tag, "Ed25519 device keypair generated and stored successfully")
    }
    
    /**
     * Delete the keypair (for testing/reset)
     */
    fun deleteKeyPair() {
        prefs.edit()
            .remove(keyPrivate)
            .remove(keyPublic)
            .apply()
        Log.i(tag, "Device keypair deleted")
    }
    
    /**
     * Base64url encode (no padding, URL-safe)
     */
    private fun base64UrlEncode(bytes: ByteArray): String {
        return Base64.encodeToString(bytes, Base64.URL_SAFE or Base64.NO_WRAP or Base64.NO_PADDING)
    }
    
    /**
     * Base64url decode (handles both padded and unpadded)
     */
    private fun base64UrlDecode(input: String): ByteArray {
        return Base64.decode(input, Base64.URL_SAFE or Base64.NO_WRAP or Base64.NO_PADDING)
    }
    
    data class SignedChallenge(
        val signature: String,
        val signedAt: Long,
        val nonce: String
    )
}