package persistence;

import model.asn1.ASN1Object;
import model.asn1.ASN1Time;
import model.asn1.Int;
import model.asn1.exceptions.InvalidCAException;
import model.asn1.exceptions.InvalidDBException;
import model.asn1.exceptions.ParseException;
import model.asn1.parsing.BytesReader;
import model.ca.AuditLogEntry;
import model.ca.CertificationAuthority;
import model.ca.Template;
import model.pki.cert.Certificate;
import model.pki.crl.Reason;
import model.pki.crl.RevokedCertificate;
import model.x501.Name;
import org.json.JSONArray;
import org.json.JSONException;
import org.json.JSONObject;
import ui.Utils;

import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.security.NoSuchAlgorithmException;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.time.ZonedDateTime;
import java.time.format.DateTimeFormatter;
import java.time.format.DateTimeParseException;
import java.util.Collections;
import java.util.List;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

/**
 * Util class that encodes / decodes the CA into JSON and vice-versa.
 */
public class Decoder {
    /**
     * EFFECTS: Convert the list into a JSONArray using the supplied mapper and add to the JSONObject. Nothing is added
     * if the given list is empty.
     */
    private static <T> void encodeList(JSONObject o, String key, List<T> list, Function<? super T, ?> mapper) {
        if (list.isEmpty()) {
            return;
        }
        o.put(key, new JSONArray(list.stream().map(mapper).collect(Collectors.toList())));
    }

    /**
     * EFFECTS: Decode the JSONArray into a list of specific objects, using the given mapper. Return empty list if the
     * o does not contain key.
     * REQUIRES: If o contains key, it must be a JSONArray.
     */
    private static <T> List<T> decodeListFromString(final JSONObject o, String key,
                                          Function<? super String, ? extends T> mapper) {
        return o.keySet().contains(key)
                ? StreamSupport.stream(o.getJSONArray(key).spliterator(), false)
                        .map(obj -> mapper.apply((String) obj))
                        .collect(Collectors.toList())
                : Collections.emptyList();
    }

    /**
     * EFFECTS: Decode the JSONArray into a list of specific objects, using the given mapper. Return empty list if the
     * o does not contain key.
     * REQUIRES: If o contains key, it must be a JSONArray.
     */
    private static <T> List<T> decodeListFromObject(final JSONObject o, String key,
                                             Function<? super JSONObject, ? extends T> mapper) {
        return o.keySet().contains(key)
                ? StreamSupport.stream(o.getJSONArray(key).spliterator(), false)
                        .map(obj -> mapper.apply((JSONObject) obj))
                        .collect(Collectors.toList())
                : Collections.emptyList();
    }

    /**
     * EFFECTS: Decode the JSON object into CertificationAuthority
     * Throws {@link InvalidDBException} if the JSON is invalid.
     * Throws {@link NoSuchAlgorithmException} if RSA is not supported.
     * JSON should be <pre>
     *     {
     *         "key": {
     *             "n": "octet string",
     *             "p": "octet string",
     *             "e": "octet string"
     *         },
     *         "cert": "-----BEGIN CERTIFICATE-----\n...",
     *         "templates": [
     *              {
     *                  "name": "123",
     *                  "enabled": true,
     *                  "subject": "-----BEGIN DISTINGUISHED NAME-----\n...",
     *                  "validity": 100
     *              }
     *          ],
     *          "signed": [
     *              "-----BEGIN CERTIFICATE-----\n...",
     *          ],
     *          "revoked": [
     *              {
     *                  "serial": "octet string",
     *                  "date": "-----BEGIN DATE-----\n",
     *                  "reason": "KEY_COMPROMISED"
     *              }
     *          ],
     *          "logs": [
     *              {
     *                  "user": "",
     *                  "time": "ISO 8601",
     *                  "action": ""
     *              }
     *          ],
     *          "serial": 1
     *     }
     * </pre>
     * Missing fields will be default.
     */
    public static CertificationAuthority decodeCA(JSONObject o) throws InvalidDBException, NoSuchAlgorithmException {
        try {
            Certificate caCert = o.keySet().contains("cert") ? decodeCertificate(o.getString("cert")) : null;
            final List<Template> templates = decodeListFromObject(o, "templates", Decoder::decodeTemplate);
            final List<Certificate> signed = decodeListFromString(o, "signed", Decoder::decodeCertificate);
            final List<RevokedCertificate> revoked =
                    decodeListFromObject(o, "revoked", Decoder::decodeRevokedCertificate);
            final List<AuditLogEntry> logs = decodeListFromObject(o, "logs", Decoder::decodeLog);
            final int serial = o.keySet().contains("serial") ? o.getInt("serial") :
                    CertificationAuthority.SERIAL_DEFAULT;
            final BigInteger n = o.keySet().contains("key")
                    ? new BigInteger(o.getJSONObject("key").getString("n"), 16) : null;
            final BigInteger p = o.keySet().contains("key")
                    ? new BigInteger(o.getJSONObject("key").getString("p"), 16) : null;
            final BigInteger e = o.keySet().contains("key")
                    ? new BigInteger(o.getJSONObject("key").getString("e"), 16) : null;
            return new CertificationAuthority(n, p, e, caCert, signed, serial, revoked, templates, logs);
        } catch (InvalidKeySpecException | NullPointerException | ParseException
                 | InvalidCAException | NumberFormatException | JSONException ex) {
            throw new InvalidDBException("Invalid JSON format", ex);
        }
    }

    /**
     * EFFECTS: Encode the CertificationAuthority to JSON.
     */
    public static JSONObject encodeCA(CertificationAuthority ca) {
        JSONObject obj = new JSONObject();
        if (ca.getKey() != null) {
            obj.put("key", encodeKey(ca.getKey(), ca.getPublicKey()));
        }
        if (ca.getCertificate() != null) {
            obj.put("cert", encodeCertificate(ca.getCertificate()));
        }
        if (ca.getSerial() != CertificationAuthority.SERIAL_DEFAULT) {
            obj.put("serial", ca.getSerial());
        }
        encodeList(obj, "templates", ca.getTemplates(), Decoder::encodeTemplate);
        encodeList(obj, "signed", ca.getSigned(), Decoder::encodeCertificate);
        encodeList(obj, "revoked", ca.getRevoked(), Decoder::encodeRevokedCertificate);
        encodeList(obj, "logs", ca.getLogs(), Decoder::encodeLog);
        return obj;
    }

    /**
     * EFFECTS: Encode the n / p / e modulus and exponents into JSON.
     */
    private static JSONObject encodeKey(RSAPrivateKey privateKey, RSAPublicKey publicKey) {
        JSONObject obj = new JSONObject();
        obj.put("n", privateKey.getModulus().toString(16));
        obj.put("p", privateKey.getPrivateExponent().toString(16));
        obj.put("e", publicKey.getPublicExponent().toString(16));
        return obj;
    }

    /**
     * EFFECTS: Encode the log entry into a JSON object.
     */
    private static JSONObject encodeLog(AuditLogEntry entry) {
        JSONObject o = new JSONObject();
        o.put("user", entry.getUser());
        o.put("time", entry.getTime().format(DateTimeFormatter.ISO_ZONED_DATE_TIME));
        o.put("action", entry.getAction());
        return o;
    }

    /**
     * EFFECTS: Decode the log entry from JSON object.
     * Throws {@link InvalidDBException} if the input is invalid.
     * REQUIRES: { "user": "", "time": "ISO_LOCAL_DATE_TIME", "action": "" }
     */
    private static AuditLogEntry decodeLog(JSONObject o) throws InvalidDBException {
        try {
            return new AuditLogEntry(o.getString("user"),
                    ZonedDateTime.parse(o.getString("time"), DateTimeFormatter.ISO_ZONED_DATE_TIME),
                    o.getString("action"));
        } catch (DateTimeParseException e) {
            throw new InvalidDBException("Invalid date", e);
        }
    }

    /**
     * EFFECTS: Encode the templates into a JSON array.
     */
    private static JSONObject encodeTemplate(Template template) {
        JSONObject o = new JSONObject();
        o.put("name", template.getName());
        o.put("enabled", template.isEnabled());
        if (template.getSubject() != null) {
            o.put("subject", Utils.toPEM(template.getSubject().encodeDER(), "DISTINGUISHED NAME"));
        }
        o.put("validity", template.getValidity());
        return o;
    }

    /**
     * EFFECTS: Decode the JSON object into a Template.
     * Throws {@link InvalidDBException} if the subject is invalid.
     * REQUIRES: { "name": "string", "enabled": boolean, "name": "-----BEGIN DISTINGUISHED NAME-----", "validity": long}
     */
    private static Template decodeTemplate(JSONObject o) throws InvalidDBException {
        try {
            return new Template(o.getString("name"),
                    o.getBoolean("enabled"),
                    o.keySet().contains("subject") ? new Name(new BytesReader(Utils.parsePEM(
                            Utils.byteToByte(o.getString("subject").getBytes()), "DISTINGUISHED NAME")),
                            false) : null,
                    o.getLong("validity"));
        } catch (ParseException e) {
            throw new InvalidDBException("Invalid template", e);
        }
    }

    /**
     * EFFECTS: Encode the certificate to PEM.
     */
    private static String encodeCertificate(Certificate cert) {
        return Utils.toPEM(cert.encodeDER(), "CERTIFICATE");
    }

    /**
     * EFFECTS: Decode the JSON pem into Certificate.
     * Throws {@link InvalidDBException} if the value cannot be parsed.
     * REQUIRES: -----BEGIN CERTIFICATE----- ...
     */
    private static Certificate decodeCertificate(String pem) {
        try {
            return new Certificate(new BytesReader(
                    Utils.parsePEM(Utils.byteToByte(pem.getBytes(StandardCharsets.UTF_8)), "CERTIFICATE")),
                    false);
        } catch (ParseException e) {
            throw new InvalidDBException("Invalid certificate", e);
        }
    }

    /**
     * EFFECTS: Encode the RevokedCertificate into a JSON object.
     */
    private static JSONObject encodeRevokedCertificate(RevokedCertificate rev) {
        JSONObject o = new JSONObject();
        o.put("serial", rev.getSerialNumber().getValue().toString(16));
        o.put("date", Utils.toPEM(rev.getRevocationDate().encodeDER(), "DATE"));
        o.put("reason", rev.getReason().name());
        return o;
    }

    /**
     * EFFECTS: Decode the RevokedCertificate from JSON.
     * Throws {@link InvalidDBException} if the value cannot be parsed.
     * REQUIRES: { "serial": "12ab", "date": "-----BEGIN DATE-----", "reason": "KEY_COMPROMISED" }
     */
    private static RevokedCertificate decodeRevokedCertificate(JSONObject o) {
        try {
            return new RevokedCertificate(ASN1Object.TAG_SEQUENCE, null,
                    new Int(Int.TAG, null, new BigInteger(o.getString("serial"), 16)),
                    (ASN1Time) ASN1Object.parse(new BytesReader(Utils.parsePEM(Utils.byteToByte(o.getString("date")
                            .getBytes(StandardCharsets.UTF_8)), "DATE")), false),
                    Reason.valueOf(o.getString("reason")));
        } catch (ParseException | IllegalArgumentException e) {
            throw new InvalidDBException("Invalid revoked certificate", e);
        }
    }
}