package com.okithaki.autoapp;

import android.util.Base64;

import org.json.JSONObject;

import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.SecureRandom;
import java.security.spec.KeySpec;

import javax.crypto.Cipher;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;

/** AES-256-GCM envelope for the local Wi-Fi/Hotspot Neighborhood Mesh. */
public final class MeshCryptoV249435 {
    public static final String PROTOCOL = "ithaki_mesh_v249435";
    private static final byte[] KDF_SALT = "OK-ITHAKI-MESH-V249435".getBytes(StandardCharsets.UTF_8);
    private static final SecureRandom RANDOM = new SecureRandom();

    private MeshCryptoV249435() { }

    private static SecretKeySpec key(String sharedCode) throws Exception {
        String value = sharedCode == null ? "" : sharedCode.trim();
        if (value.length() < 12) throw new IllegalArgumentException("mesh_code_min_12");
        KeySpec spec = new PBEKeySpec(value.toCharArray(), KDF_SALT, 150000, 256);
        byte[] raw = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA1").generateSecret(spec).getEncoded();
        return new SecretKeySpec(raw, "AES");
    }

    private static byte[] aad(String packetId) {
        return (PROTOCOL + ":" + (packetId == null ? "" : packetId)).getBytes(StandardCharsets.UTF_8);
    }

    public static JSONObject encrypt(String sharedCode, JSONObject payload) throws Exception {
        String packetId = payload.optString("packet_id", "");
        if (packetId.trim().isEmpty()) throw new IllegalArgumentException("packet_id_required");
        byte[] iv = new byte[12]; RANDOM.nextBytes(iv);
        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        cipher.init(Cipher.ENCRYPT_MODE, key(sharedCode), new GCMParameterSpec(128, iv));
        cipher.updateAAD(aad(packetId));
        byte[] encrypted = cipher.doFinal(payload.toString().getBytes(StandardCharsets.UTF_8));
        JSONObject envelope = new JSONObject();
        envelope.put("protocol", PROTOCOL);
        envelope.put("packet_id", packetId);
        envelope.put("from_id", payload.optString("from_id", ""));
        envelope.put("to_id", payload.optString("to_id", ""));
        envelope.put("created_at", payload.optDouble("created_at", System.currentTimeMillis()/1000.0));
        envelope.put("ttl", 6);
        envelope.put("hops", 0);
        envelope.put("iv", Base64.encodeToString(iv, Base64.NO_WRAP));
        envelope.put("cipher_text", Base64.encodeToString(encrypted, Base64.NO_WRAP));
        envelope.put("kdf", "PBKDF2-HMAC-SHA1-150000");
        envelope.put("cipher", "AES-256-GCM");
        return envelope;
    }

    public static JSONObject decrypt(String sharedCode, JSONObject envelope) throws Exception {
        if (!PROTOCOL.equals(envelope.optString("protocol", ""))) throw new IllegalArgumentException("protocol_mismatch");
        String packetId = envelope.optString("packet_id", "");
        byte[] iv = Base64.decode(envelope.optString("iv", ""), Base64.DEFAULT);
        byte[] encrypted = Base64.decode(envelope.optString("cipher_text", ""), Base64.DEFAULT);
        if (iv.length != 12 || encrypted.length < 17) throw new IllegalArgumentException("invalid_envelope");
        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        cipher.init(Cipher.DECRYPT_MODE, key(sharedCode), new GCMParameterSpec(128, iv));
        cipher.updateAAD(aad(packetId));
        JSONObject payload = new JSONObject(new String(cipher.doFinal(encrypted), StandardCharsets.UTF_8));
        if (!packetId.equals(payload.optString("packet_id", ""))) throw new SecurityException("packet_id_mismatch");
        return payload;
    }

    public static String fingerprint(String sharedCode) {
        try {
            byte[] d = MessageDigest.getInstance("SHA-256").digest((sharedCode == null ? "" : sharedCode.trim()).getBytes(StandardCharsets.UTF_8));
            StringBuilder b = new StringBuilder();
            for (int i=0; i<6; i++) b.append(String.format("%02X", d[i]));
            return b.toString();
        } catch (Exception ignored) { return ""; }
    }
}
