import org.apache.commons.codec.binary.Base64;
import org.apache.http.HttpHeaders;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPut;
import org.apache.http.entity.FileEntity;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.DefaultRedirectStrategy;
import org.apache.http.impl.client.HttpClientBuilder;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;

import java.io.File;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.UUID;


/**
 * This example mainly demonstrates how to use stream load to import data
 * Including file type (CSV) and data in JSON format
 *
 */
public class DorisStreamLoader {
    // FE IP Address
    private final static String HOST = "10.220.146.10";
    // FE port
    private final static int PORT = 8030;
    // db name
    private final static String DATABASE = "test_2";
    // table name
    private final static String TABLE = "doris_test_sink";
    //user name
    private final static String USER = "root";
    //user password
    private final static String PASSWD = "";
    //The path of the local file to be imported
    private final static String LOAD_FILE_NAME = "c:/es/1.csv";

    //http path of stream load task submission
    private final static String loadUrl = String.format("http://%s:%s/api/%s/%s/_stream_load",
            HOST, PORT, DATABASE, TABLE);

    //构建HTTP客户端
    private final static HttpClientBuilder httpClientBuilder = HttpClients
            .custom()
            .setRedirectStrategy(new DefaultRedirectStrategy() {
                @Override
                protected boolean isRedirectable(String method) {
                    // If the connection target is FE, you need to deal with 307 redirect。
                    return true;
                }
            });

    /**
     * 文件数据导入
     * @param file
     * @throws Exception
     */
    public void load(File file) throws Exception {
        try (CloseableHttpClient client = httpClientBuilder.build()) {
            HttpPut put = new HttpPut(loadUrl);
            put.removeHeaders(HttpHeaders.CONTENT_LENGTH);
            put.removeHeaders(HttpHeaders.TRANSFER_ENCODING);
            put.setHeader(HttpHeaders.EXPECT, "100-continue");
            put.setHeader(HttpHeaders.AUTHORIZATION, basicAuthHeader(USER, PASSWD));

            // You can set stream load related properties in the Header, here we set label and column_separator.
            put.setHeader("label", UUID.randomUUID().toString());
            put.setHeader("column_separator", ",");

            // Set up the import file. Here you can also use StringEntity to transfer arbitrary data.
            FileEntity entity = new FileEntity(file);
            put.setEntity(entity);

            try (CloseableHttpResponse response = client.execute(put)) {
                String loadResult = "";
                if (response.getEntity() != null) {
                    loadResult = EntityUtils.toString(response.getEntity());
                }

                final int statusCode = response.getStatusLine().getStatusCode();
                if (statusCode != 200) {
                    throw new IOException(String.format("Stream load failed. status: %s load result: %s", statusCode, loadResult));
                }
                System.out.println("Get load result: " + loadResult);
            }
        }
    }

    /**
     * JSON格式的数据导入
     * @param jsonData
     * @throws Exception
     */
    public void loadJson(String jsonData) throws Exception {
        try (CloseableHttpClient client = httpClientBuilder.build()) {
            HttpPut put = new HttpPut(loadUrl);
            put.removeHeaders(HttpHeaders.CONTENT_LENGTH);
            put.removeHeaders(HttpHeaders.TRANSFER_ENCODING);
            put.setHeader(HttpHeaders.EXPECT, "100-continue");
            put.setHeader(HttpHeaders.AUTHORIZATION, basicAuthHeader(USER, PASSWD));

            // You can set stream load related properties in the Header, here we set label and column_separator.
            put.setHeader("label", UUID.randomUUID().toString());
            put.setHeader("column_separator", ",");
            put.setHeader("format", "json");

            // Set up the import file. Here you can also use StringEntity to transfer arbitrary data.
            StringEntity entity = new StringEntity(jsonData);
            put.setEntity(entity);

            try (CloseableHttpResponse response = client.execute(put)) {
                String loadResult = "";
                if (response.getEntity() != null) {
                    loadResult = EntityUtils.toString(response.getEntity());
                }

                final int statusCode = response.getStatusLine().getStatusCode();
                if (statusCode != 200) {
                    throw new IOException(String.format("Stream load failed. status: %s load result: %s", statusCode, loadResult));
                }
                System.out.println("Get load result: " + loadResult);
            }
        }
    }

    /**
     * 封装认证信息
     * @param username
     * @param password
     * @return
     */
    private String basicAuthHeader(String username, String password) {
        final String tobeEncode = username + ":" + password;
        byte[] encoded = Base64.encodeBase64(tobeEncode.getBytes(StandardCharsets.UTF_8));
        return "Basic " + new String(encoded);
    }


    public static void main(String[] args) throws Exception {
        DorisStreamLoader loader = new DorisStreamLoader();
        //file load
        //File file = new File(LOAD_FILE_NAME);
        //loader.load(file);
        //json load
        String jsonData = "{\"id\":556393582,\"number\":\"123344\",\"price\":\"23.5\",\"skuname\":\"test\",\"skudesc\":\"zhangfeng_test,test\"}";
        loader.loadJson(jsonData);

    }
}

public class DorisSink extends RichSinkFunction<String> {
 
     private static final Logger log = LoggerFactory.getLogger(DorisSink.class);
 
     private final static List<String> DORIS_SUCCESS_STATUS = new ArrayList<>(Arrays.asList("Success", "Publish Timeout"));
 
     private DorisStreamLoad dorisStreamLoad;
 
     private String columns;
 
     private String jsonFormat;
 
     public DorisSink(DorisStreamLoad dorisStreamLoad, String columns, String jsonFormat) {
         this.dorisStreamLoad = dorisStreamLoad;
         this.columns = columns;
         this.jsonFormat = jsonFormat;
    }
 
     @Override
     public void open(Configuration parameters) throws Exception {
         super.open(parameters);
    }
 
 
     /**
      * 判断StreamLoad是否成功
      * @param respContent streamload返回的响应信息（JSON格式）
      * @return
      */
     public static Boolean checkStreamLoadStatus(RespContent respContent) {
         if (DORIS_SUCCESS_STATUS.contains(respContent.getStatus())
                 && respContent.getNumberTotalRows() == respContent.getNumberLoadedRows()) {
             return true;
        } else {
             return false;
        }
    }
 
     @Override
     public void invoke(String value, Context context) throws Exception {
         DorisStreamLoad.LoadResponse loadResponse = dorisStreamLoad.loadBatch(value, columns, jsonFormat);
         if (loadResponse != null && loadResponse.status == 200) {
             RespContent respContent = JSON.parseObject(loadResponse.respContent, RespContent.class);
             if (!checkStreamLoadStatus(respContent)) {
                 log.error("Stream Load fail{}:", loadResponse);
            }
        } else {
             log.error("Stream Load Request failed:{}", loadResponse);
        }
    }
 }

public class DorisStreamLoad implements Serializable {
 
     private static final Logger log = LoggerFactory.getLogger(DorisStreamLoad.class);
 //连接地址，这里使用的是连接FE
     private static String loadUrlPattern = "http://%s/api/%s/%s/_stream_load?";
     //fe ip地址
     private String hostPort;
     //数据库
     private String db;
     //要导入的数据表名
     private String tbl;
     //用户名
     private String user;
     //密码
     private String passwd;
     private String loadUrlStr;
     private String authEncoding;
 
 
     public DorisStreamLoad(String hostPort, String db, String tbl, String user, String passwd) {
         this.hostPort = hostPort;
         this.db = db;
         this.tbl = tbl;
         this.user = user;
         this.passwd = passwd;
         this.loadUrlStr = String.format(loadUrlPattern, hostPort, db, tbl);
         this.authEncoding = Base64.getEncoder().encodeToString(String.format("%s:%s", user, passwd).getBytes(StandardCharsets.UTF_8));
    }
 //获取http连接信息
     private HttpURLConnection getConnection(String urlStr, String label, String columns, String jsonformat) throws IOException {
         URL url = new URL(urlStr);
         HttpURLConnection conn = (HttpURLConnection) url.openConnection();
         conn.setInstanceFollowRedirects(false);
         conn.setRequestMethod("PUT");
         conn.setRequestProperty("Authorization", "Basic " + authEncoding);
         conn.addRequestProperty("Expect", "100-continue");
         conn.addRequestProperty("Content-Type", "text/plain; charset=UTF-8");
         conn.addRequestProperty("label", label);
         conn.addRequestProperty("max_filter_ratio", "0");
         conn.addRequestProperty("strict_mode", "true");
         conn.addRequestProperty("columns", columns);
         conn.addRequestProperty("format", "json");
         conn.addRequestProperty("jsonpaths", jsonformat);
         conn.addRequestProperty("strip_outer_array", "true");
         conn.setDoOutput(true);
         conn.setDoInput(true);
 
         return conn;
    }
 
     public static class LoadResponse {
         public int status;
         public String respMsg;
         public String respContent;
 
         public LoadResponse(int status, String respMsg, String respContent) {
             this.status = status;
             this.respMsg = respMsg;
             this.respContent = respContent;
        }
 
         @Override
         public String toString() {
             StringBuilder sb = new StringBuilder();
             sb.append("status: ").append(status);
             sb.append(", resp msg: ").append(respMsg);
             sb.append(", resp content: ").append(respContent);
             return sb.toString();
        }
    }
 //执行数据导入
     public LoadResponse loadBatch(String data, String columns, String jsonformat) {
         Calendar calendar = Calendar.getInstance();
         //导入的lable，全局唯一
         String label = String.format("flink_import_%s%02d%02d_%02d%02d%02d_%s",
                 calendar.get(Calendar.YEAR), calendar.get(Calendar.MONTH) + 1, calendar.get(Calendar.DAY_OF_MONTH),
                 calendar.get(Calendar.HOUR_OF_DAY), calendar.get(Calendar.MINUTE), calendar.get(Calendar.SECOND),
                 UUID.randomUUID().toString().replaceAll("-", ""));
 
         HttpURLConnection feConn = null;
         HttpURLConnection beConn = null;
         try {
             // build request and send to fe
             feConn = getConnection(loadUrlStr, label, columns, jsonformat);
             int status = feConn.getResponseCode();
             // fe send back http response code TEMPORARY_REDIRECT 307 and new be location
             if (status != 307) {
                 throw new Exception("status is not TEMPORARY_REDIRECT 307, status: " + status);
            }
             String location = feConn.getHeaderField("Location");
             if (location == null) {
                 throw new Exception("redirect location is null");
            }
             // build request and send to new be location
             beConn = getConnection(location, label, columns, jsonformat);
             // send data to be
             BufferedOutputStream bos = new BufferedOutputStream(beConn.getOutputStream());
             bos.write(data.getBytes());
             bos.close();
 
             // get respond
             status = beConn.getResponseCode();
             String respMsg = beConn.getResponseMessage();
             InputStream stream = (InputStream) beConn.getContent();
             BufferedReader br = new BufferedReader(new InputStreamReader(stream));
             StringBuilder response = new StringBuilder();
             String line;
             while ((line = br.readLine()) != null) {
                 response.append(line);
            }
             return new LoadResponse(status, respMsg, response.toString());
 
        } catch (Exception e) {
             e.printStackTrace();
             String err = "failed to load audit via AuditLoader plugin with label: " + label;
             log.warn(err, e);
             return new LoadResponse(-1, e.getMessage(), err);
        } finally {
             if (feConn != null) {
                 feConn.disconnect();
            }
             if (beConn != null) {
                 beConn.disconnect();
            }
        }
    }
 
 }

 public class FlinkKafka2Doris {
     //kafka address
     private static final String bootstrapServer = "xxx:9092,xxx:9092,xxx:9092";
     //kafka groupName
     private static final String groupName = "test_flink_doris_group";
     //kafka topicName
     private static final String topicName = "test_flink_doris";
     //doris ip port
     private static final String hostPort = "xxx:8030";
     //doris dbName
     private static final String dbName = "db1";
     //doris tbName
     private static final String tbName = "tb1";
     //doris userName
     private static final String userName = "root";
     //doris password
     private static final String password = "";
     //doris columns
     private static final String columns = "name,age,price,sale";
     //json format
     private static final String jsonFormat = "[\"$.name\",\"$.age\",\"$.price\",\"$.sale\"]";
 
     public static void main(String[] args) throws Exception {
 
         Properties props = new Properties();
         props.put("bootstrap.servers", bootstrapServer);
         props.put("group.id", groupName);
         props.put("key.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
         props.put("value.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
         props.put("auto.offset.reset", "earliest");
         props.put("max.poll.records", "10000");
 
         StreamExecutionEnvironment blinkStreamEnv = StreamExecutionEnvironment.getExecutionEnvironment();
         blinkStreamEnv.enableCheckpointing(10000);         blinkStreamEnv.getCheckpointConfig().enableExternalizedCheckpoints(CheckpointConfig.ExternalizedCheckpointCleanup.RETAIN_ON_CANCELLATION);
         FlinkKafkaConsumer<String> flinkKafkaConsumer = new FlinkKafkaConsumer<>(topicName,
                 new SimpleStringSchema(),
                 props); 
         DataStreamSource<String> dataStreamSource = blinkStreamEnv.addSource(flinkKafkaConsumer);
         DorisStreamLoad dorisStreamLoad = new DorisStreamLoad(hostPort, dbName, tbName, userName, password); 
         dataStreamSource.addSink(new DorisSink(dorisStreamLoad,columns,jsonFormat));
         blinkStreamEnv.execute("flink kafka to doris");
    }
 }