package com.jscape.inet.ssh.protocol.v2.marshaling;

import com.jscape.inet.ssh.protocol.messages.Message;
import com.jscape.inet.ssh.protocol.v2.marshaling.algorithms.Compression;
import com.jscape.inet.ssh.protocol.v2.marshaling.algorithms.Encryption;
import com.jscape.inet.ssh.protocol.v2.marshaling.algorithms.Mac;
import com.jscape.util.ArrayTools;
import com.jscape.util.BigEndianBits;
import com.jscape.util.IOTools;
import com.jscape.util.SuppressObfuscation;
import com.jscape.util.io.ByteArrayInputStream;
import com.jscape.util.io.ByteArrayOutputStream;
import com.jscape.util.io.DataCodec;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Arrays;
import java.util.Random;

import static com.jscape.util.Parameters.assertNotNull;

/**
 * @author Alex Usun
 */
public class PacketCodec
        implements DataCodec<Message> {

    private static final int PACKET_HEADER_LENGTH = 5;
    private static final int PADDING_LENGTH_OFFSET = 4;
    private static final int MIN_LENGTH_DIVISOR = 8;
    private static final int MIN_PADDING_LENGTH = 4;
    private static final int MAX_PADDING_LENGTH = 255;
    private static final int INITIAL_BUFFER_SIZE = 35000;

    private final DataCodec<Message> messageCodec;
    private final SequenceNumber readingSequenceNumber;
    private final SequenceNumber writingSequenceNumber;
    private final ReadingBuffer readingBuffer;
    private final ByteArrayOutputStream messageBuffer;
    private final ByteArrayOutputStream writingBuffer;
    private final Random paddingRandom;
    private Encryption readingEncryption;
    private Encryption writingEncryption;
    private Mac readingMac;
    private Mac writingMac;
    private Compression readingCompression;
    private Compression writingCompression;
    private int readingDivisor;
    private int writingDivisor;
    private byte[] macBuffer;
    private volatile long readBytes;
    private volatile long writtenBytes;
    private volatile long readMessages;
    private volatile long writtenMessages;

    public PacketCodec(DataCodec<Message> messageCodec) {
        assertNotNull(messageCodec);
        this.messageCodec = messageCodec;

        this.readingSequenceNumber = new SequenceNumber();
        this.writingSequenceNumber = new SequenceNumber();
        this.readingBuffer = new ReadingBuffer(INITIAL_BUFFER_SIZE);
        this.messageBuffer = new ByteArrayOutputStream(INITIAL_BUFFER_SIZE);
        this.writingBuffer = new ByteArrayOutputStream(INITIAL_BUFFER_SIZE);
        this.paddingRandom = new Random();
        this.readingEncryption = Encryption.NULL;
        this.writingEncryption = Encryption.NULL;
        this.readingMac = Mac.NULL;
        this.writingMac = Mac.NULL;
        this.readingCompression = Compression.NULL;
        this.writingCompression = Compression.NULL;
        this.readingDivisor = MIN_LENGTH_DIVISOR;
        this.writingDivisor = MIN_LENGTH_DIVISOR;
        this.macBuffer = new byte[this.readingMac.digestLength()];
    }

    @Override
    public Message read(InputStream in)
            throws IOException {
        nextIncomingMessage();
        readFirstBlock(in);
        int remainingPacketLength = remainingPacketLength();
        readPacket(remainingPacketLength, in);
        assertMacValid();

        return readMessage();
    }

    @Override
    public void write(Message value, OutputStream out)
            throws IOException {
        nextOutgoingMessage();
        packetFor(value);
        applyMac();
        writePacket(out);
    }

    public void readingEncryption(Encryption encryption) {
        assertNotNull(encryption);
        this.readingEncryption = encryption;

        this.readingDivisor = Math.max(this.readingEncryption.blockLength(), MIN_LENGTH_DIVISOR);
    }

    public void writingEncryption(Encryption encryption) {
        assertNotNull(encryption);
        this.writingEncryption = encryption;

        this.writingDivisor = Math.max(this.writingEncryption.blockLength(), MIN_LENGTH_DIVISOR);
    }

    public void readingMac(Mac mac) {
        assertNotNull(mac);
        this.readingMac = mac;

        this.macBuffer = new byte[this.readingMac.digestLength()];
    }

    public void writingMac(Mac mac) {
        assertNotNull(mac);
        this.writingMac = mac;
    }

    public void readingCompression(Compression compression) {
        assertNotNull(compression);
        this.readingCompression = compression;
    }

    public void writingCompression(Compression compression) {
        assertNotNull(compression);
        this.writingCompression = compression;
    }

    public long readBytes() {
        return this.readBytes;
    }

    public long writtenBytes() {
        return this.writtenBytes;
    }

    public long readMessages() {
        return this.readMessages;
    }

    public long writtenMessages() {
        return this.writtenMessages;
    }

    @Override
    public String toString() {
        return "PacketCodec {"
                + "messageCodec=" + this.messageCodec
                + ", readingSequenceNumber=" + this.readingSequenceNumber
                + ", writingSequenceNumber=" + this.writingSequenceNumber
                + ", readingEncryption=" + this.readingEncryption
                + ", writingEncryption=" + this.writingEncryption
                + ", readingMac=" + this.readingMac
                + ", writingMac=" + this.writingMac
                + ", readingCompression=" + this.readingCompression
                + ", writingCompression=" + this.writingCompression
                + ", readingDivisor=" + this.readingDivisor
                + ", writingDivisor=" + this.writingDivisor
                + ", readBytes=" + this.readBytes
                + ", writtenBytes=" + this.writtenBytes
                + ", readMessages=" + this.readMessages
                + ", writtenMessages=" + this.writtenMessages
                + '}';
    }

    private void nextIncomingMessage() {
        this.readingSequenceNumber.next();
    }

    private void readFirstBlock(InputStream in)
            throws IOException {
        this.readingBuffer.reset();
        this.readingBuffer.fill(in, this.readingDivisor);
        this.readingEncryption.apply(this.readingBuffer.buffer(), 0, this.readingDivisor, this.readingBuffer.buffer(), 0);
    }

    private int remainingPacketLength()
            throws IOException {
        int packetLength = BigEndianBits.INSTANCE.getInt(this.readingBuffer.buffer(), 0);

        if (packetLength < 0 || packetLength > 1048575) {
            throw new IOException(String.format("Unsupported SSH packet length: %s. Possible decryption error.", packetLength));
        }

        return PACKET_HEADER_LENGTH - 1 + packetLength - this.readingDivisor;
    }

    private void readPacket(int remainingPacketLength, InputStream in)
            throws IOException {
        this.readingBuffer.fill(in, remainingPacketLength + this.readingMac.digestLength());
        this.readingEncryption.apply(this.readingBuffer.buffer(), this.readingDivisor, remainingPacketLength, this.readingBuffer.buffer(), this.readingDivisor);
        ArrayTools.copy(this.readingBuffer.buffer(), this.readingDivisor + remainingPacketLength, this.macBuffer);
        this.readBytes += this.readingBuffer.position();
        ++this.readMessages;
    }

    private void assertMacValid()
            throws IOException {
        try {
            this.readingMac.assertMacValid(this.readingSequenceNumber.valueBytes(), this.readingBuffer.buffer(), 0, this.readingBuffer.position() - this.readingMac.digestLength(), this.macBuffer);
        } catch (Mac.InvalidMacException e) {
            throw new InvalidMacException();
        }
    }

    private Message readMessage()
            throws IOException {
        int paddingLength = this.readingBuffer.at(PADDING_LENGTH_OFFSET);
        int payloadLength = this.readingBuffer.position() - this.readingMac.digestLength() - PACKET_HEADER_LENGTH - paddingLength;
        int decompressedPayloadLength = this.readingCompression.apply(this.readingBuffer.buffer(), PACKET_HEADER_LENGTH, payloadLength, this.readingBuffer.asOutputStream(0));
        InputStream in = this.readingBuffer.asInputStream(0, decompressedPayloadLength);

        Message message = this.messageCodec.read(in);
        message.number = this.readingSequenceNumber.value();

        return message;
    }

    private void nextOutgoingMessage() {
        this.writingSequenceNumber.next();
    }

    private void packetFor(Message message)
            throws IOException {
        message.number = this.writingSequenceNumber.value();
        payloadFor(message);
        byte paddingLength = addPadding();
        int packetLength = this.writingBuffer.position() - PACKET_HEADER_LENGTH + 1;
        BigEndianBits.INSTANCE.putInt(packetLength, this.writingBuffer.buffer(), 0);
        this.writingBuffer.buffer()[PADDING_LENGTH_OFFSET] = paddingLength;
    }

    private void applyMac()
            throws IOException {
        this.writingMac.apply(this.writingSequenceNumber.valueBytes(), this.writingBuffer.buffer(), 0, this.writingBuffer.position(), this.writingBuffer);
    }

    private void writePacket(OutputStream out)
            throws IOException {
        this.writingEncryption.apply(this.writingBuffer.buffer(), 0, this.writingBuffer.position() - this.writingMac.digestLength(), this.writingBuffer.buffer(), 0);
        this.writingBuffer.writeTo(out);
        this.writtenBytes += this.writingBuffer.position();
        ++this.writtenMessages;
    }

    private void payloadFor(Message message)
            throws IOException {
        this.messageBuffer.reset();
        this.messageCodec.write(message, this.messageBuffer);

        this.writingBuffer.position(PACKET_HEADER_LENGTH);
        this.writingCompression.apply(this.messageBuffer.buffer(), 0, this.messageBuffer.position(), this.writingBuffer);
    }

    private byte addPadding()
            throws IOException {
        int minLength = this.writingBuffer.position() + MIN_PADDING_LENGTH;
        int reminder = minLength % this.writingDivisor;
        int length = reminder > 0 ? MIN_PADDING_LENGTH + this.writingDivisor - reminder : MIN_PADDING_LENGTH;
        byte[] padding = new byte[length];
        this.paddingRandom.nextBytes(padding);
        this.writingBuffer.write(padding);

        return (byte) length;
    }

    public static class SequenceNumber {

        private final byte[] valueBytes;
        private int value;

        public SequenceNumber() {
            this.valueBytes = new byte[4];
            this.value = -1;
        }

        public byte[] valueBytes() {
            return this.valueBytes;
        }

        public int value() {
            return this.value;
        }

        public void next() {
            ++this.value;
            BigEndianBits.INSTANCE.putInt(this.value, this.valueBytes, 0);
        }

        @Override
        public String toString() {
            return "SequenceNumber {"
                    + "valueBytes=" + Arrays.toString(this.valueBytes)
                    + ", value=" + this.value
                    + '}';
        }
    }

    public static class ReadingBuffer {

        private byte[] buffer;
        private int position;
        private final ByteArrayInputStream in;
        private final ByteArrayOutputStream out;

        public ReadingBuffer(int size) {
            this.buffer = new byte[size];
            this.position = 0;
            this.in = new ByteArrayInputStream(this.buffer);
            this.out = new ByteArrayOutputStream(this.buffer);
        }

        public void reset() {
            this.position = 0;
        }

        public void fill(InputStream in, int length)
                throws IOException {
            adjustBuffer(length);
            IOTools.read(this.buffer, this.position, length, in);
            this.position += length;
        }

        public byte[] buffer() {
            return this.buffer;
        }

        public int position() {
            return this.position;
        }

        public int at(int index) {
            return this.buffer[index] & 0xFF;
        }

        public InputStream asInputStream(int offset, int length) {
            this.in.put(this.buffer, offset, length);
            return this.in;
        }

        public OutputStream asOutputStream(int offset) {
            this.out.put(this.buffer, offset);
            return this.out;
        }

        private int available() {
            return this.buffer.length - this.position;
        }

        private void adjustBuffer(int length) {
            if (available() < length) {
                int newLength = Math.max(this.buffer.length * 2, this.position + length);
                this.buffer = ArrayTools.copy(this.buffer, 0, new byte[newLength], 0, this.position);
            }
        }
    }

    @SuppressObfuscation
    public static class InvalidMacException
            extends IOException {
    }
}
