package org.apache.storm.hdfs.trident;

import backtype.storm.task.IMetricsContext;
import backtype.storm.topology.FailedException;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.storm.hdfs.trident.format.FileNameFormat;
import org.apache.storm.hdfs.trident.format.RecordFormat;
import org.apache.storm.hdfs.trident.rotation.FileRotationPolicy;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import storm.trident.operation.TridentCollector;
import storm.trident.state.State;
import storm.trident.tuple.TridentTuple;

import java.io.IOException;
import java.io.Serializable;
import java.net.URI;
import java.util.List;
import java.util.Map;

public class HdfsState implements State {

    public static class Options implements Serializable {
        private RecordFormat format;
        private String fsUrl;
        private String path;
        private FileRotationPolicy rotationPolicy;
        private FileNameFormat fileNameFormat;

        public Options(){}

        public Options withRecordFormat(RecordFormat format){
            this.format = format;
            return this;
        }

        public Options withRotationPolicy(FileRotationPolicy policy){
            this.rotationPolicy = policy;
            return this;
        }

        public Options withFileNameFormat(FileNameFormat format){
            this.fileNameFormat = format;
            return this;
        }

        public Options withFsUrl(String url){
            this.fsUrl = url;
            return this;
        }

        public Options withPath(String path){
            this.path = path;
            return this;
        }

    }

    public static final Logger LOG = LoggerFactory.getLogger(HdfsState.class);

    private FileSystem fs;
    private FSDataOutputStream out;
    private int rotation = 0;
    private long offset = 0;

    private Options options;

    HdfsState(Options options){
        this.options = options;
    }

    void prepare(Map conf, IMetricsContext metrics, int partitionIndex, int numPartitions){
        LOG.info("Preparing HDFS Bolt...");
        if(this.options.format == null) throw new IllegalStateException("RecordFormat must be specified.");;
        if(this.options.rotationPolicy == null) throw new IllegalStateException("RotationPolicy must be specified.");


        if(this.options.fsUrl ==  null || this.options.path == null){
            throw new IllegalStateException("File system URL and base path must be specified.");
        }
        this.options.fileNameFormat.prepare(conf, partitionIndex, numPartitions);

        try{
            Configuration hdfsConfig = new Configuration();
            this.fs = FileSystem.get(URI.create(this.options.fsUrl), hdfsConfig);
            out = this.fs.create(new Path(this.options.path, this.options.fileNameFormat.getName(this.rotation, System.currentTimeMillis())));
        } catch (Exception e){
            throw new RuntimeException("Error preparing HdfsBolt: " + e.getMessage(), e);
        }
    }



    @Override
    public void beginCommit(Long txId) {
    }

    @Override
    public void commit(Long txId) {
        try {
            this.out.hsync();
        } catch (IOException e) {
            LOG.warn("Commit failed.", e);
        }
    }

    public void updateState(List<TridentTuple> tuples, TridentCollector tridentCollector){
        try{
            for(TridentTuple tuple : tuples){

                byte[] bytes = this.options.format.format(tuple);
                out.write(bytes);
                this.offset += bytes.length;

                if(this.options.rotationPolicy.mark(tuple, this.offset)){
                    rotateOutputFile();
                    this.options.rotationPolicy.reset();
                    this.offset = 0;
                }
            }
        } catch (IOException e){
            LOG.warn("Failing batch due to IOException.", e);
            throw new FailedException(e);
        }
    }

    private void rotateOutputFile() throws IOException {
        LOG.info("Rotating output file...");
        long start = System.currentTimeMillis();
        this.out.hsync();
        this.out.close();
        this.rotation++;
        this.out = this.fs.create(
                new Path(
                        this.options.path,
                        this.options.fileNameFormat.getName(this.rotation, System.currentTimeMillis())
                )
        );

        long time = System.currentTimeMillis() - start;
        LOG.info("File rotation took {} ms.", time);
    }
}

@@ -292,16 +292,77 @@
 s.Path;%0A
+import org.apache.storm.hdfs.common.rotation.RotationAction;%0A
 import o
@@ -405,32 +405,32 @@
 FileNameFormat;%0A
-
 import org.apach
@@ -798,16 +798,44 @@
 et.URI;%0A
+import java.util.ArrayList;%0A
 import j
@@ -1166,16 +1166,109 @@
 eFormat;
+%0A        private ArrayList%3CRotationAction%3E rotationActions = new ArrayList%3CRotationAction%3E();
 %0A%0A      
@@ -1871,16 +1871,16 @@
  path)%7B%0A
-
         
@@ -1933,24 +1933,171 @@
 %0A        %7D%0A%0A
+        public Options addRotationAction(RotationAction action)%7B%0A            this.rotationActions.add(action);%0A            return this;%0A        %7D%0A%0A
     %7D%0A%0A    p
@@ -4639,50 +4639,19 @@
     
-this.out = this.fs.create(%0A               
+Path path =
  new
@@ -4657,24 +4657,16 @@
 w Path(%0A
-        
         
@@ -4704,32 +4704,24 @@
             
-  
-      
 this.options
@@ -4795,24 +4795,27 @@
         
+);%0A
         
 )%0A      
@@ -4806,29 +4806,170 @@
         
-)%0A        );%0A
+this.out = this.fs.create(path);%0A%0A        for(RotationAction action : this.options.rotationActions)%7B%0A            action.execute(this.fs, path);%0A%0A        %7D
 %0A       

add RotationActions to State implementation.