package com.lumaserv.bgp.protocol.message;

import com.lumaserv.bgp.protocol.DataBuilder;
import com.lumaserv.bgp.protocol.IPPrefix;
import com.lumaserv.bgp.protocol.attribute.PathAttribute;
import lombok.Getter;
import lombok.NoArgsConstructor;
import lombok.Setter;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;

@NoArgsConstructor
@Getter
@Setter
public class BGPUpdate implements DataBuilder {

    List<IPPrefix> withdrawnPrefixes = new ArrayList<>();
    List<PathAttribute> attributes = new ArrayList<>();
    List<IPPrefix> prefixes = new ArrayList<>();

    public BGPUpdate(ByteBuffer message) {
        int routesLength = message.getShort() & 0xFFFF;
        int offset = 2;
        int offsetOffset = 2;
        while ((offset - offsetOffset) < routesLength) {
            IPPrefix prefix = new IPPrefix()
                    .setLength(message.get())
                    .setAddress(new byte[4]);
            offset++;
            int addressLen = (int) Math.ceil(prefix.getLength() / 8d);
            message.get(prefix.getAddress(), 0, addressLen);
            offset += addressLen;
            withdrawnPrefixes.add(prefix);
        }
        int attributesLength = message.getShort() & 0xFFFF;
        offset += 2;
        offsetOffset = offset;
        while ((offset - offsetOffset) < attributesLength) {
            byte flags = message.get();
            byte typeCode = message.get();
            int length = message.get() & 0xFF;
            offset += 3;
            if((flags & 0b0001_0000) > 0) {
                length <<= 8;
                length |= message.get() & 0xFF;
                offset++;
            }
            byte[] data = new byte[length];
            message.get(data);
            offset += length;
            PathAttribute attr = PathAttribute.from(typeCode, data);
            attributes.add(attr);
        }
        while (message.hasRemaining()) {
            IPPrefix prefix = new IPPrefix()
                    .setLength(message.get())
                    .setAddress(new byte[4]);
            offset++;
            int addressLen = (int) Math.ceil(prefix.getLength() / 8d);
            message.get(prefix.getAddress(), 0, addressLen);
            offset += addressLen;
            prefixes.add(prefix);
        }
    }

    public void build(ByteBuffer message) {
        int first = message.position();
        message.putShort((short)0); // Withdrawn
        message.putShort((short)0); // Attributes
        // TODO
        int pos = message.position();
        for (PathAttribute attr: attributes) {
// FIXME add header
            byte flags = 0b0100_0000; // TODO
            int start = message.position();
            message.put(flags);
            message.put(attr.getTypeCode());
            message.put((byte)0); // Dummy length
            attr.build(message);
            int length = message.position() - (start + 3);
            if (length > 255)
                throw new RuntimeException("attribute length > 255");
            message.put(start + 2, (byte)(length & 0xFF));
        }
        int attributesLength = message.position() - pos;
        message.putShort(first + 2, (short)attributesLength);
        for (IPPrefix prefix : prefixes) {
            message.put((byte)prefix.getLength());
            int addressLen = (int) Math.ceil(prefix.getLength() / 8d);
            message.put(prefix.getAddress(), 0, addressLen);
        }
    }
}