package com.alterdekim.game.component.rtmp;


import com.alterdekim.game.crypto.HMAC;
import com.alterdekim.game.message.ByteMessage;
import com.alterdekim.game.message.SerializerUtils;
import com.alterdekim.game.message.serializer.BinaryMessageSerializer;
import com.alterdekim.game.utils.ByteUtils;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.Socket;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Random;
import java.util.stream.IntStream;

@Slf4j
public class HandshakeProcessor extends ConnectionProcessor {

    private HandshakeState handshakeState;

    private static final int PROTOCOL_VERSION = 0x03;
    private static final int HANDSHAKE_SIZE = 1536;
    private static final int DIGEST_OFFSET_INDICATOR_POS = 772;
    private static final int SHA256_DIGEST_SIZE = 32;

    private static final byte[] GENUINE_FP_KEY = {
            (byte) 0x47, (byte) 0x65, (byte) 0x6E, (byte) 0x75, (byte) 0x69, (byte) 0x6E, (byte) 0x65, (byte) 0x20,
            (byte) 0x41, (byte) 0x64, (byte) 0x6F, (byte) 0x62, (byte) 0x65, (byte) 0x20, (byte) 0x46, (byte) 0x6C,
            (byte) 0x61, (byte) 0x73, (byte) 0x68, (byte) 0x20, (byte) 0x50, (byte) 0x6C, (byte) 0x61, (byte) 0x79,
            (byte) 0x65, (byte) 0x72, (byte) 0x20, (byte) 0x30, (byte) 0x30, (byte) 0x31, // Genuine Adobe Flash Player 001
            (byte) 0xF0, (byte) 0xEE, (byte) 0xC2, (byte) 0x4A, (byte) 0x80, (byte) 0x68, (byte) 0xBE, (byte) 0xE8,
            (byte) 0x2E, (byte) 0x00, (byte) 0xD0, (byte) 0xD1, (byte) 0x02, (byte) 0x9E, (byte) 0x7E, (byte) 0x57,
            (byte) 0x6E, (byte) 0xEC, (byte) 0x5D, (byte) 0x2D, (byte) 0x29, (byte) 0x80, (byte) 0x6F, (byte) 0xAB,
            (byte) 0x93, (byte) 0xB8, (byte) 0xE6, (byte) 0x36, (byte) 0xCF, (byte) 0xEB, (byte) 0x31, (byte) 0xAE};

    private final Random random;

    public HandshakeProcessor(InputStream inputStream, OutputStream outputStream, Socket sock) {
        super(inputStream, outputStream, sock);
        this.random = new Random();
        this.handshakeState = HandshakeState.C0S0;
    }

    @Override
    public ConnectionState getState() {
        return ConnectionState.HANDSHAKE;
    }

    @Override
    public ConnectionState process() throws IOException {
        switch (handshakeState) {
            case C0S0 -> processC0();
            case C1S1 -> processC1();
            case C2S2 -> {
                return processThirdMessage();
            }
        }
        return getState();
    }

    private void processC0() throws IOException {
        int b = this.getInputStream().read();
        if( b != PROTOCOL_VERSION ) {close(); return;}
        this.getOutputStream().write(b);
        this.handshakeState = HandshakeState.C1S1;
    }

    private void processC1() throws IOException {
        byte[] b = new byte[HANDSHAKE_SIZE];
        if( this.getInputStream().read(b) == -1 ) { close(); return; }
        this.getOutputStream().write(b);
        this.handshakeState = HandshakeState.C2S2;
    }

    private ConnectionState processThirdMessage() throws IOException {
        byte[] b = new byte[HANDSHAKE_SIZE];
        if( this.getInputStream().read(b) == -1 ) { close(); return ConnectionState.HANDSHAKE; }

        final int digestOffset = random.nextInt(HANDSHAKE_SIZE - DIGEST_OFFSET_INDICATOR_POS - 4 - 8 - SHA256_DIGEST_SIZE);
        final int absoluteDigestOffset = ((digestOffset % 728) + DIGEST_OFFSET_INDICATOR_POS + 4);

        int remaining = digestOffset;
        final byte[] digestOffsetBytes = new byte[4];
        for (int i = 3; i >= 0; i--) {
            if (remaining > 255) {
                digestOffsetBytes[i] = (byte)255;
                remaining -= 255;
            } else {
                digestOffsetBytes[i] = (byte)remaining;
                remaining -= remaining;
            }
        }

        byte[] partBeforeDigest = new byte[absoluteDigestOffset];
        random.nextBytes(partBeforeDigest);

        byte[] timeStamp = SerializerUtils.bytesToPrimitive(BinaryMessageSerializer.serializeInteger((int) (System.currentTimeMillis() / 1000)));
        System.arraycopy(timeStamp, 0, partBeforeDigest, 0, 4);
        System.arraycopy(new byte[]{(byte) 0x80, 0x00, 0x07, 0x02}, 0, partBeforeDigest, 4, 4);

        byte[] partAfterDigest = new byte[HANDSHAKE_SIZE - absoluteDigestOffset - SHA256_DIGEST_SIZE];
        random.nextBytes(partAfterDigest);
        System.arraycopy(digestOffsetBytes, 0, partBeforeDigest, DIGEST_OFFSET_INDICATOR_POS, 4);

        byte[] tempBuffer = new byte[HANDSHAKE_SIZE - SHA256_DIGEST_SIZE];
        System.arraycopy(partBeforeDigest, 0, tempBuffer, 0, partBeforeDigest.length);
        System.arraycopy(partAfterDigest, 0, tempBuffer, partBeforeDigest.length, partAfterDigest.length);

        HMAC crypto = new HMAC();
        byte[] digest = crypto.calculateHmacSHA256(tempBuffer, GENUINE_FP_KEY, 30);

        this.getOutputStream().write(partBeforeDigest);
        this.getOutputStream().write(digest);
        this.getOutputStream().write(partAfterDigest);
        this.getOutputStream().flush();

        return ConnectionState.CONNECTED;
    }

    private enum HandshakeState {
        C0S0,
        C1S1,
        C2S2
    }
}
