package org.asamk.signal.storage.protocol;

import com.fasterxml.jackson.core.JsonGenerator;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.databind.DeserializationContext;
import com.fasterxml.jackson.databind.JsonDeserializer;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.JsonSerializer;
import com.fasterxml.jackson.databind.SerializerProvider;

import org.whispersystems.libsignal.InvalidKeyIdException;
import org.whispersystems.libsignal.state.PreKeyRecord;
import org.whispersystems.libsignal.state.PreKeyStore;
import org.whispersystems.util.Base64;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

class JsonPreKeyStore implements PreKeyStore {

    private final Map<Integer, byte[]> store = new HashMap<>();

    public JsonPreKeyStore() {

    }

    private void addPreKeys(Map<Integer, byte[]> preKeys) {
        store.putAll(preKeys);
    }

    @Override
    public PreKeyRecord loadPreKey(int preKeyId) throws InvalidKeyIdException {
        try {
            if (!store.containsKey(preKeyId)) {
                throw new InvalidKeyIdException("No such prekeyrecord!");
            }

            return new PreKeyRecord(store.get(preKeyId));
        } catch (IOException e) {
            throw new AssertionError(e);
        }
    }

    @Override
    public void storePreKey(int preKeyId, PreKeyRecord record) {
        store.put(preKeyId, record.serialize());
    }

    @Override
    public boolean containsPreKey(int preKeyId) {
        return store.containsKey(preKeyId);
    }

    @Override
    public void removePreKey(int preKeyId) {
        store.remove(preKeyId);
    }

    public static class JsonPreKeyStoreDeserializer extends JsonDeserializer<JsonPreKeyStore> {

        @Override
        public JsonPreKeyStore deserialize(JsonParser jsonParser, DeserializationContext deserializationContext) throws IOException {
            JsonNode node = jsonParser.getCodec().readTree(jsonParser);

            Map<Integer, byte[]> preKeyMap = new HashMap<>();
            if (node.isArray()) {
                for (JsonNode preKey : node) {
                    Integer preKeyId = preKey.get("id").asInt();
                    try {
                        preKeyMap.put(preKeyId, Base64.decode(preKey.get("record").asText()));
                    } catch (IOException e) {
                        System.err.println(String.format("Error while decoding prekey for: %s", preKeyId));
                    }
                }
            }

            JsonPreKeyStore keyStore = new JsonPreKeyStore();
            keyStore.addPreKeys(preKeyMap);

            return keyStore;

        }
    }

    public static class JsonPreKeyStoreSerializer extends JsonSerializer<JsonPreKeyStore> {

        @Override
        public void serialize(JsonPreKeyStore jsonPreKeyStore, JsonGenerator json, SerializerProvider serializerProvider) throws IOException {
            json.writeStartArray();
            for (Map.Entry<Integer, byte[]> preKey : jsonPreKeyStore.store.entrySet()) {
                json.writeStartObject();
                json.writeNumberField("id", preKey.getKey());
                json.writeStringField("record", Base64.encodeBytes(preKey.getValue()));
                json.writeEndObject();
            }
            json.writeEndArray();
        }
    }
}
