package com.ihooyah.flume.sink;

import org.apache.flume.*;
import org.apache.flume.conf.BatchSizeSupported;
import org.apache.flume.conf.Configurable;
import org.apache.flume.instrumentation.SinkCounter;
import org.apache.flume.sink.AbstractSink;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class LocalUploadSink extends AbstractSink implements Configurable, BatchSizeSupported {

    private static final Logger logger = LoggerFactory.getLogger(LocalUploadSink.class);

    private String uploadDir;

    private int batchSize;

    private static final int DEFAULT_BATCH_SIZE = 1000;

    private SinkCounter sinkCounter;

    private String fileReadName;

    private CounterGroup counterGroup = new CounterGroup();

    @Override
    public Status process() throws EventDeliveryException {
        Map<String , List<Event>> eventMap=new HashMap<>();
        Status status = Status.READY;
        Channel channel = getChannel();
        Transaction transaction = channel.getTransaction();
        transaction.begin();
        try {
            for (int i = 0; i < batchSize; i++) {
                Event event = channel.take();
                if (event == null) {
                    break;
                }
                fileReadName=event.getHeaders().get("fileName").toString();
                //System.out.println("读取文件名:"+fileReadName+" 文件大小:"+event.getBody().length);
                this.counterGroup.incrementAndGet("event.file");
                List<Event> eventList=eventMap.get(fileReadName);
                if (eventList==null){
                    eventList=new ArrayList<>();
                }
                eventList.add(event);
                eventMap.put(fileReadName,eventList);
            }
            for (Map.Entry<String , List<Event>> data:eventMap.entrySet()) {
                //System.out.println("文件名称:"+data.getKey()+"集合大小:"+data.getValue().size());
                dealData(data.getKey(),data.getValue());
            }
            transaction.commit();
        } catch (Exception e) {
            transaction.rollback();
            e.printStackTrace();
        } finally {
            transaction.close();
        }
        return status;
    }

    @Override
    public void configure(Context context) {
        uploadDir = context.getString("uploadDir", "/home/www/upload/hza/file/");
        batchSize = context.getInteger("batchSize", DEFAULT_BATCH_SIZE);
    }

    public void dealData(String fileName,List<Event> eventList) throws IOException {
        if (eventList.size() == 0) {
            return;
        }

        logger.info("deal data numbers: {}", eventList.size());
        String filename = System.currentTimeMillis() +fileName;
        File localFile = new File(uploadDir + filename);
        FileOutputStream fos = new FileOutputStream(localFile);
        for (int i = 0; i < eventList.size(); i++) {
            fos.write(eventList.get(i).getBody());
        }

        fos.flush();
        fos.close();
    }

    @Override
    public long getBatchSize() {
        return batchSize;
    }
}
