/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 * http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.nifi.processors.network;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import org.apache.nifi.annotation.behavior.InputRequirement;
import org.apache.nifi.annotation.behavior.InputRequirement.Requirement;
import org.apache.nifi.annotation.behavior.ReadsAttribute;
import org.apache.nifi.annotation.behavior.ReadsAttributes;
import org.apache.nifi.annotation.behavior.SideEffectFree;
import org.apache.nifi.annotation.behavior.SupportsBatching;
import org.apache.nifi.annotation.behavior.WritesAttribute;
import org.apache.nifi.annotation.behavior.WritesAttributes;
import org.apache.nifi.annotation.documentation.CapabilityDescription;
import org.apache.nifi.annotation.documentation.Tags;
import org.apache.nifi.annotation.lifecycle.OnScheduled;
import org.apache.nifi.components.AllowableValue;
import org.apache.nifi.components.PropertyDescriptor;
import org.apache.nifi.flowfile.FlowFile;
import org.apache.nifi.flowfile.attributes.CoreAttributes;
import org.apache.nifi.logging.ComponentLog;
import org.apache.nifi.processor.AbstractProcessor;
import org.apache.nifi.processor.ProcessContext;
import org.apache.nifi.processor.ProcessSession;
import org.apache.nifi.processor.Relationship;
import org.apache.nifi.processor.exception.ProcessException;
import org.apache.nifi.processor.io.OutputStreamCallback;
import org.apache.nifi.processors.network.parser.Netflowv5Parser;

import java.io.BufferedOutputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.OptionalInt;
import java.util.Set;

import static org.apache.nifi.processors.network.parser.Netflowv5Parser.getHeaderFields;
import static org.apache.nifi.processors.network.parser.Netflowv5Parser.getRecordFields;

@SideEffectFree
@SupportsBatching
@InputRequirement(Requirement.INPUT_REQUIRED)
@Tags({ "network", "netflow", "attributes", "datagram", "v5", "packet", "byte" })
@CapabilityDescription("Parses netflowv5 byte ingest and add to NiFi flowfile as attributes or JSON content.")
@ReadsAttributes({ @ReadsAttribute(attribute = "udp.port", description = "Optionally read if packets are received from UDP datagrams.") })
@WritesAttributes({ @WritesAttribute(attribute = "netflowv5.header.*", description = "The key and value generated by the parsing of the header fields."),
        @WritesAttribute(attribute = "netflowv5.record.*", description = "The key and value generated by the parsing of the record fields.") })

public class ParseNetflowv5 extends AbstractProcessor {
    private String destination;
    // Add mapper
    private static final ObjectMapper mapper = new ObjectMapper();

    public static final String FLOWFILE_CONTENT = "flowfile-content";
    public static final String FLOWFILE_ATTRIBUTE = "flowfile-attribute";
    public static final AllowableValue DESTINATION_CONTENT = new AllowableValue(FLOWFILE_CONTENT, FLOWFILE_CONTENT,
            "Parsed data routes as flowfile JSON content");
    public static final AllowableValue DESTINATION_ATTRIBUTES = new AllowableValue(FLOWFILE_ATTRIBUTE, FLOWFILE_ATTRIBUTE,
            "Parsed data routes as flowfile attributes");
    public static final PropertyDescriptor FIELDS_DESTINATION = new PropertyDescriptor.Builder().name("FIELDS_DESTINATION").displayName("Parsed fields destination")
            .description("Indicates whether the results of the parser are written " + "to the FlowFile content or a FlowFile attribute; if using " + DESTINATION_ATTRIBUTES
                    + ", fields will be populated as attributes. If set to " + DESTINATION_CONTENT + ", the netflowv5 field will be converted into a flat JSON object.")
            .required(true).allowableValues(DESTINATION_CONTENT, DESTINATION_ATTRIBUTES).defaultValue(DESTINATION_CONTENT.getDisplayName()).build();

    public static final Relationship REL_FAILURE = new Relationship.Builder().name("failure")
            .description("Any FlowFile that could not be parsed as a netflowv5 message will be transferred to this Relationship without any attributes being added").build();
    public static final Relationship REL_ORIGINAL = new Relationship.Builder().name("original").description("The original raw content").build();
    public static final Relationship REL_SUCCESS = new Relationship.Builder().name("success")
            .description("Any FlowFile that is successfully parsed as a netflowv5 data will be transferred to this Relationship.").build();

    public static final List<PropertyDescriptor> PROPERTIES = Collections.unmodifiableList(Arrays.asList(FIELDS_DESTINATION));
    public static final Set<Relationship> RELATIONSHIPS = Collections.unmodifiableSet(new HashSet<>(Arrays.asList(REL_FAILURE, REL_ORIGINAL, REL_SUCCESS)));

    @Override
    public Set<Relationship> getRelationships() {
        return RELATIONSHIPS;
    }

    @Override
    public final List<PropertyDescriptor> getSupportedPropertyDescriptors() {
        return PROPERTIES;
    }

    @OnScheduled
    public void onScheduled(final ProcessContext context) {
        destination = context.getProperty(FIELDS_DESTINATION).getValue();
    }

    @Override
    public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
        final ComponentLog logger = getLogger();
        FlowFile flowFile = session.get();
        if (flowFile == null) {
            return;
        }

        final OptionalInt portNumber = resolvePort(flowFile);
        final Netflowv5Parser parser = new Netflowv5Parser(portNumber);
        final ByteArrayOutputStream baos = new ByteArrayOutputStream();
        session.exportTo(flowFile, baos);
        final byte[] buffer = baos.toByteArray();

        final int processedRecord;
        try {
            processedRecord = parser.parse(buffer);
            if (logger.isDebugEnabled()) {
                logger.debug("Parsed {} records from the packet", processedRecord);
            }
        } catch (Throwable e) {
            logger.error("Routing to failure since while processing {}, parser returned unexpected Exception", flowFile, e);
            session.transfer(flowFile, REL_FAILURE);
            return;
        }

        try {
            final List<FlowFile> multipleRecords = new ArrayList<>();
            switch (destination) {
                case FLOWFILE_ATTRIBUTE:
                    final Map<String, String> attributes = new HashMap<>();
                    generateKV(multipleRecords, session, flowFile, attributes, parser, processedRecord);
                    break;
                case FLOWFILE_CONTENT:
                    generateJSON(multipleRecords, session, flowFile, parser, processedRecord);
                    break;
            }

            // Ready to transfer and commit
            session.transfer(flowFile, REL_ORIGINAL);
            session.transfer(multipleRecords, REL_SUCCESS);
            session.adjustCounter("Records Processed", processedRecord, false);
        } catch (Exception e) {
            // The flowfile has failed parsing & validation, routing to failure
            logger.error("Routing to failure since failed to parse {} as a netflowv5 message", flowFile, e);

            // Create a provenance event recording the routing to failure
            session.transfer(flowFile, REL_FAILURE);
        }
    }

    private void generateJSON(final List<FlowFile> multipleRecords, final ProcessSession session, final FlowFile flowFile, final Netflowv5Parser parser, final int processedRecord)
            throws JsonProcessingException {
        int numberOfRecords = processedRecord;
        FlowFile recordFlowFile = flowFile;
        int record = 0;
        while (numberOfRecords-- > 0) {
            ObjectNode results = mapper.createObjectNode();
            // Add Port number and message format
            results.set("port", mapper.valueToTree(parser.getPortNumber()));
            results.set("format", mapper.valueToTree("netflowv5"));

            recordFlowFile = session.create(flowFile);
            // Add JSON Objects
            generateJSONUtil(results, parser, record++);

            recordFlowFile = session.write(recordFlowFile, new OutputStreamCallback() {
                @Override
                public void process(OutputStream out) throws IOException {
                    try (OutputStream outputStream = new BufferedOutputStream(out)) {
                        outputStream.write(mapper.writeValueAsBytes(results));
                    }
                }
            });
            // Adjust the FlowFile mime.type attribute
            recordFlowFile = session.putAttribute(recordFlowFile, CoreAttributes.MIME_TYPE.key(), "application/json");
            multipleRecords.add(recordFlowFile);
        }
    }

    private void generateKV(final List<FlowFile> multipleRecords, final ProcessSession session, final FlowFile flowFile, final Map<String, String> attributes, final Netflowv5Parser parser,
            final int processedRecord) {
        int numberOfRecords = processedRecord;
        generateHeaderAttributes(attributes, parser);

        final String[] fieldname = getRecordFields();
        int record = 0;
        FlowFile recordFlowFile = flowFile;
        while (numberOfRecords-- > 0) {
            // Process KVs of the Flow Record fields
            final Object[] fieldvalue = parser.getRecordData()[record++];
            for (int i = 0; i < fieldname.length; i++) {
                attributes.put("netflowv5.record." + fieldname[i], String.valueOf(fieldvalue[i]));
            }
            recordFlowFile = session.create(flowFile);
            recordFlowFile = session.putAllAttributes(recordFlowFile, attributes);
            multipleRecords.add(recordFlowFile);
        }
    }

    private OptionalInt resolvePort(final FlowFile flowFile) {
        final String port;
        if ((port = flowFile.getAttribute("udp.port")) != null) {
            return OptionalInt.of(Integer.parseInt(port));
        }
        return OptionalInt.empty();
    }

    private void generateJSONUtil(final ObjectNode results, final Netflowv5Parser parser, final int record) {
        final ObjectNode header = mapper.createObjectNode();

        // Process KVs of the Flow Header fields
        String fieldname[] = getHeaderFields();
        Object fieldvalue[] = parser.getHeaderData();
        for (int i = 0; i < fieldname.length; i++) {
            header.set(fieldname[i], mapper.valueToTree(fieldvalue[i]));
        }
        results.set("header", header);

        final ObjectNode data = mapper.createObjectNode();
        // Process KVs of the Flow Record fields
        fieldname = getRecordFields();
        fieldvalue = parser.getRecordData()[record];
        for (int i = 0; i < fieldname.length; i++) {
            data.set(fieldname[i], mapper.valueToTree(fieldvalue[i]));
        }
        results.set("record", data);
    }

    private void generateHeaderAttributes(final Map<String, String> attributes, final Netflowv5Parser parser) {
        // Process KVs of the Flow Header fields
        final String[] hFieldName = getHeaderFields();
        final Object[] hFieldValue = parser.getHeaderData();

        for (int i = 0; i < hFieldName.length; i++) {
            attributes.put("netflowv5.header." + hFieldName[i], String.valueOf(hFieldValue[i]));
        }
    }
}
