package me.proxylink.common;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.security.GeneralSecurityException;
import java.security.MessageDigest;
import java.security.SecureRandom;
import java.util.Arrays;
import java.util.Objects;

public final class AuthPayloads {
    private static final SecureRandom SECURE_RANDOM = new SecureRandom();

    private AuthPayloads() {
    }

    public record Hello(int protocolVersion, String role, String nodeName, byte[] clientNonce) {
        public Hello {
            Objects.requireNonNull(role, "role");
            Objects.requireNonNull(nodeName, "nodeName");
            clientNonce = copyNonce(clientNonce, "clientNonce");
        }

        @Override
        public byte[] clientNonce() {
            return Arrays.copyOf(clientNonce, clientNonce.length);
        }
    }

    public record Challenge(byte[] serverNonce) {
        public Challenge {
            serverNonce = copyNonce(serverNonce, "serverNonce");
        }

        @Override
        public byte[] serverNonce() {
            return Arrays.copyOf(serverNonce, serverNonce.length);
        }
    }

    public record Response(byte[] hmac) {
        public Response {
            if (hmac == null || hmac.length != ProtocolConstants.AUTH_RESPONSE_BYTES) {
                throw new IllegalArgumentException("hmac must be " + ProtocolConstants.AUTH_RESPONSE_BYTES + " bytes");
            }
            hmac = Arrays.copyOf(hmac, hmac.length);
        }

        @Override
        public byte[] hmac() {
            return Arrays.copyOf(hmac, hmac.length);
        }
    }

    public static byte[] randomNonce() {
        byte[] nonce = new byte[ProtocolConstants.AUTH_NONCE_BYTES];
        SECURE_RANDOM.nextBytes(nonce);
        return nonce;
    }

    public static byte[] encodeHello(Hello hello) throws IOException {
        ByteArrayOutputStream bytes = new ByteArrayOutputStream();
        DataOutputStream output = new DataOutputStream(bytes);
        output.writeInt(hello.protocolVersion());
        output.writeUTF(hello.role());
        output.writeUTF(hello.nodeName());
        output.write(hello.clientNonce());
        output.flush();
        return bytes.toByteArray();
    }

    public static Hello decodeHello(byte[] payload) throws IOException {
        DataInputStream input = new DataInputStream(new ByteArrayInputStream(payload));
        int protocolVersion = input.readInt();
        String role = input.readUTF();
        String nodeName = input.readUTF();
        byte[] clientNonce = input.readNBytes(ProtocolConstants.AUTH_NONCE_BYTES);
        if (clientNonce.length != ProtocolConstants.AUTH_NONCE_BYTES || input.available() != 0) {
            throw new ProtocolException("Invalid HELLO payload");
        }
        return new Hello(protocolVersion, role, nodeName, clientNonce);
    }

    public static byte[] encodeChallenge(Challenge challenge) throws IOException {
        return challenge.serverNonce();
    }

    public static Challenge decodeChallenge(byte[] payload) throws IOException {
        if (payload.length != ProtocolConstants.AUTH_NONCE_BYTES) {
            throw new ProtocolException("Invalid AUTH_CHALLENGE payload");
        }
        return new Challenge(payload);
    }

    public static byte[] encodeResponse(Response response) {
        return response.hmac();
    }

    public static Response decodeResponse(byte[] payload) throws IOException {
        if (payload.length != ProtocolConstants.AUTH_RESPONSE_BYTES) {
            throw new ProtocolException("Invalid AUTH_RESPONSE payload");
        }
        return new Response(payload);
    }

    public static byte[] computeResponse(String authToken, Hello hello, Challenge challenge) throws ProtocolException {
        try {
            Mac mac = Mac.getInstance("HmacSHA256");
            mac.init(new SecretKeySpec(authToken.getBytes(StandardCharsets.UTF_8), "HmacSHA256"));
            updateUtf8(mac, ProtocolConstants.AUTH_CONTEXT);
            updateUtf8(mac, hello.role());
            updateUtf8(mac, hello.nodeName());
            mac.update(hello.clientNonce());
            mac.update(challenge.serverNonce());
            return mac.doFinal();
        } catch (GeneralSecurityException e) {
            throw new ProtocolException("Unable to compute authentication response", e);
        }
    }

    public static boolean constantTimeEquals(byte[] expected, byte[] actual) {
        return MessageDigest.isEqual(expected, actual);
    }

    private static void updateUtf8(Mac mac, String value) {
        byte[] bytes = value.getBytes(StandardCharsets.UTF_8);
        mac.update((byte) ((bytes.length >>> 8) & 0xff));
        mac.update((byte) (bytes.length & 0xff));
        mac.update(bytes);
    }

    private static byte[] copyNonce(byte[] nonce, String name) {
        if (nonce == null || nonce.length != ProtocolConstants.AUTH_NONCE_BYTES) {
            throw new IllegalArgumentException(name + " must be " + ProtocolConstants.AUTH_NONCE_BYTES + " bytes");
        }
        return Arrays.copyOf(nonce, nonce.length);
    }
}
