package model.asn1;

import model.asn1.exceptions.ParseException;
import model.asn1.parsing.BytesReader;
import ui.Utils;

import java.math.BigInteger;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * Represents an X.680 OID, which is a global-unique multi-component int value (with a registry managing all OIDs).
 */
public class ObjectIdentifier extends ASN1Object {
    /**
     * The X.680 universal class tag assignment.
     */
    public static final Tag TAG = new Tag(TagClass.UNIVERSAL, false, 0x6);

    public static final Integer[] OID_CN = new Integer[]{2, 5, 4, 3};
    public static final Integer[] OID_SN = new Integer[]{2, 5, 4, 4};
    public static final Integer[] OID_C = new Integer[]{2, 5, 4, 6};
    public static final Integer[] OID_L = new Integer[]{2, 5, 4, 7};
    public static final Integer[] OID_O = new Integer[]{2, 5, 4, 10};
    public static final Integer[] OID_OU = new Integer[]{2, 5, 4, 11};
    public static final Integer[] OID_DC = new Integer[]{0, 9, 2342, 19200300, 100, 1, 25};

    public static final Integer[] OID_EXTENSION_REQUEST =
            new Integer[]{1, 2, 840, 113549, 1, 9, 14};

    public static final Integer[] OID_RSA_ENCRYPTION =
            new Integer[]{1, 2, 840, 113549, 1, 1, 1};
    public static final Integer[] OID_SHA256_WITH_RSA_ENCRYPTION =
            new Integer[]{1, 2, 840, 113549, 1, 1, 11};

    public static final Integer[] OID_EC_PUBLIC_KEY =
            new Integer[]{1, 2, 840, 10045, 2, 1};
    public static final Integer[] OID_ECDSA_WITH_SHA256 =
            new Integer[]{1, 2, 840, 10045, 4, 3, 2};
    public static final Integer[] OID_ECDSA_WITH_SHA512 =
            new Integer[]{1, 2, 840, 10045, 4, 3, 4};
    public static final Integer[] OID_PRIME256_V1 =
            new Integer[]{1, 2, 840, 10045, 3, 1, 7};

    public static final Integer[] OID_SUBJECT_KEY_IDENTIFIER =
            new Integer[]{2, 5, 29, 14};
    public static final Integer[] OID_KEY_USAGE =
            new Integer[]{2, 5, 29, 15};
    public static final Integer[] OID_BASIC_CONSTRAINTS =
            new Integer[]{2, 5, 29, 19};
    public static final Integer[] OID_AUTHORITY_KEY_IDENTIFIER =
            new Integer[]{2, 5, 29, 35};
    public static final Integer[] OID_CRL_DISTRIBUTION_POINTS =
            new Integer[]{2, 5, 29, 31};
    public static final Integer[] OID_AUTHORITY_INFO_ACCESS =
            new Integer[]{1, 3, 6, 1, 5, 5, 7, 1, 1};

    public static final Integer[] OID_CURVED_25519 =
            new Integer[]{1, 3, 101, 112};

    public static final Integer[] OID_CRL_REASON =
            new Integer[]{2, 5, 29, 21};

    private final Integer[] ints;

    /**
     * EFFECTS: Construct the OID object with the given array of OID numbers. For the tag and parentTag,
     * consult {@link ASN1Object}.
     * REQUIRES: The ints array must have at least two elements, and the first element must be 0, 1, or 2.
     * If the first element is 0 or 1, the second element must be < 40. For the tag and parentTag,
     * consult {@link ASN1Object}.
     */
    public ObjectIdentifier(Tag tag, Tag parentTag, Integer[] ints) {
        super(tag, parentTag);
        this.ints = ints;
    }

    /**
     * EFFECTS: Parse the input DER.
     * Throws {@link ParseException} if the input is invalid:
     * - Zero bytes long
     * - A multibyte integer is unterminated until the end of input
     */
    public ObjectIdentifier(BytesReader encoded, boolean hasParentTag) throws ParseException {
        super(encoded, hasParentTag);
        if (getLength() < 1) {
            throw new ParseException("Invalid OID");
        }
        List<Integer> nums = new ArrayList<>();
        final Byte[] raw = encoded.require(getLength(), true);
        Byte first = raw[0];
        if (first >= 80) {
            nums.add(2);
            nums.add(first - 80);
        } else if (first >= 40) {
            nums.add(1);
            nums.add(first - 40);
        } else {
            nums.add(0);
            nums.add((int) first);
        }
        nums.addAll(parse(raw));
        this.ints = nums.toArray(new Integer[0]);
    }

    /**
     * EFFECTS: Get OID from a known part, case insensitive. Currently known: C CN OU O L DC.
     *          Throws {@link ParseException} if the name is unsupported.
     */
    public static Integer[] getKnown(String name) throws ParseException {
        switch (name.toUpperCase()) {
            case "C": return OID_C;
            case "CN": return OID_CN;
            case "OU": return OID_OU;
            case "O": return OID_O;
            case "L": return OID_L;
            case "DC": return OID_DC;
            default: throw new ParseException("Unsupported DN part: " + name);
        }
    }

    /**
     * EFFECTS: Parse input OID bytes.
     * REQUIRES: raw.length >= 1
     */
    private static List<Integer> parse(Byte[] raw) throws ParseException {
        List<Integer> nums = new ArrayList<>();
        List<BitSet> num = new ArrayList<>();
        for (int i = 1; i < raw.length; i++) {
            Byte b = raw[i];
            num.add(BitSet.valueOf(new byte[]{(byte) (b & 127)}));
            if ((b & -128) == 0) {
                BitSet bitSet = new BitSet(num.size() * 7);
                int z = 0;

                for (int j = num.size() - 1; j >= 0; j--) {
                    for (int k = 0; k < 7; k++) {
                        bitSet.set(z++, num.get(j).get(k));
                    }
                }

                List<Byte> bs1 = Arrays.asList(Utils.byteToByte(bitSet.toByteArray()));
                Collections.reverse(bs1);
                nums.add(new BigInteger(Utils.byteToByte(bs1.toArray(new Byte[0]))).intValueExact());
                num.clear();
            }
        }
        if (!num.isEmpty()) {
            throw new ParseException("Unterminated byte.");
        }
        return nums;
    }

    /**
     * EFFECTS: Generate a human-readable output of that OID, in the format of 0.1.2. In case of a well-known OID, its
     * name is returned.
     */
    @Override
    public String toString() {
        if (Arrays.equals(ints, OID_CN)) {
            return "CN";
        } else if (Arrays.equals(ints, OID_SN)) {
            return "SN";
        } else if (Arrays.equals(ints, OID_C)) {
            return "C";
        } else if (Arrays.equals(ints, OID_L)) {
            return "L";
        } else if (Arrays.equals(ints, OID_O)) {
            return "O";
        } else if (Arrays.equals(ints, OID_OU)) {
            return "OU";
        } else if (Arrays.equals(ints, OID_DC)) {
            return "DC";
        }
        return Arrays.stream(ints)
                .map(i -> Integer.toString(i))
                .collect(Collectors.joining("."));
    }

    /**
     * EFFECTS: Encode the OID into DER bytes, following the DER rules as follows:
     * - First two ints: first * 40 + second
     * - Remaining: Int components are encoded as-is if they are <= 127. Otherwise, they are encoded into multiple 7bit
     * bytes, with the MSB set on every byte except for the last (rightmost byte) of each component.
     * - Integers are in big-endian.
     */
    @Override
    public Byte[] encodeValueDER() {
        return Stream.of(
                        Arrays.asList(Utils.valToByte(ints[0] * 40 + ints[1])),
                        Stream.of(ints)
                                .skip(2)
                                .map(ObjectIdentifier::encodeSingleInt)
                                .flatMap(Collection::stream)
                                .collect(Collectors.toList())
                ).flatMap(Collection::stream)
                .toArray(Byte[]::new);
    }

    /**
     * EFFECTS: Encode a single int component into OID-format bytes.
     */
    private static List<Byte> encodeSingleInt(int i) {
        BigInteger bi = BigInteger.valueOf(i);
        List<Byte> bs = Arrays.asList(Utils.byteToByte(bi.toByteArray()));
        Collections.reverse(bs);
        final BitSet bitSetOriginal = BitSet.valueOf(Utils.byteToByte(bs.toArray(new Byte[0])));
        BitSet bitSet = new BitSet(bs.size() * 16);
        int k = 0;
        for (int j = 0; j < bs.size() * 8; j++) {
            if (j == 0 || j % 7 != 0) {
                bitSet.set(k++, bitSetOriginal.get(j));
            } else {
                bitSet.set(k++, j != 7);
                bitSet.set(k++, bitSetOriginal.get(j));
            }
        }
        byte[] bs1 = bitSet.toByteArray();
        List<Byte> res =
                Arrays.asList(Utils.byteToByte(bs1));
        Collections.reverse(res);
        return res;
    }

    public Integer[] getInts() {
        return ints;
    }
}