package madynes.flowoid.exporter;

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


public class DataFlowSet implements Exportable {
	
	// See RFC3954 5.3 > format
    private final short headerSize = 4;
	private ArrayList<DataRecord> dataRecords;
	private TemplateRecord recordTemplate;
	
	/**
	 * Defines a NetFlow v9 FlowSet associated to a template record.
	 * According to the RFC, this FlowSet may then contain only data
	 * records defined with this template. See RFC3954 5.3.
	 * 
	 * @param recordTemplate template record then accepted
	 */
	public DataFlowSet(TemplateRecord recordTemplate) {
		dataRecords = new ArrayList<DataRecord>();
		this.recordTemplate = recordTemplate;
	}
	
	/**
	 * Adds a data record inevitably defined with the template given to
	 * the constructor. See RFC3954 5.3 > FlowSet ID = Template ID.
	 * 
	 * @param dataRecord data record to add inside
	 * @throws TemplateMismatchException launched if the record was not
	 * defined with the same template that the one given to the constructor 
	 */
	public void addDataRecord(DataRecord dataRecord) throws TemplateMismatchException {
		if(!dataRecord.getTemplate().equals(recordTemplate)) {
			throw new TemplateMismatchException("You CANNOT add a record into a"
					+ " FlowSet if this one was not defined with the same template"
					+ " record that the FlowSet (here " + recordTemplate.getTemplateID()
					+ " instead of " + recordTemplate.getTemplateID() + ").");
		}
			
		dataRecords.add(dataRecord);
	}

	/**
	 * Returns the FlowSetID corresponding to the record template ID accepted.
	 *  
	 * @return FlowSetID
	 */
	public short getFlowSetID() {
		return recordTemplate.getTemplateID();
	}
	
	/**
	 * Returns the number of data records added in this FlowSet.
	 * 
	 * @return Number of data records added in this FlowSet
	 */
	public short getCount() {
		return (short) dataRecords.size();
	}

	/**
	 * See RFC3954 5.3 > Padding. 
	 * 
	 * @return Padding size to add at the end of this FlowSet in
	 * the corresponding export packet.
	 */
	private int getPaddingSize() {
		return 4 - (getStrictTotalLength() % 4);
	}
	
	/**
	 * See getPaddingSize().
	 * 
	 * @return Padding to add.
	 */
	public ArrayList<Byte> getPadding() {
		ArrayList<Byte> padding = new ArrayList<Byte>();
		int paddingSize = getPaddingSize();

		// RFC3954 5.3 > Padding
		for(int i = 0; i < paddingSize; i++)
			padding.add((byte) 0);
		
		return padding;
	}

	/**
	 * Returns the total length of this FlowSet without headers.
	 * 
	 * @return total length without headers
	 */
	public short getStrictTotalLength() {
		short length = 0;
		
		for(DataRecord data : dataRecords) {
			length += data.getTotalLength();
		}
		
		return (short) (headerSize + length);
	}
	
	@Override
	public short getTotalLength() {
		return (short) (getStrictTotalLength() + getPaddingSize());
	}

	@Override
	public ArrayList<Byte> toBytes() throws TemplateMismatchException {
		ArrayList<Byte> bytes = new ArrayList<Byte>();
		
    	ByteBuffer headers = ByteBuffer.allocate(headerSize);
    	headers.order(ByteOrder.BIG_ENDIAN);
    	
    	headers.putShort(getFlowSetID());
    	
		// RFC3954 5.3 > Length
    	headers.putShort(getTotalLength());

    	for(byte b : headers.array()) {
    		bytes.add((Byte) b);
    	}
		
		for(DataRecord record : dataRecords) {
			bytes.addAll(record.toBytes());
		}
		
		for(Byte b : getPadding()) {
			bytes.add(b);
		}
			
		return bytes;
	}
}
