package org.rakam.clickhouse.analysis;

import com.facebook.presto.sql.RakamSqlFormatter;
import com.facebook.presto.sql.tree.Expression;
import com.google.common.collect.ImmutableList;
import com.google.inject.Inject;
import org.rakam.analysis.RetentionQueryExecutor;
import org.rakam.analysis.metadata.Metastore;
import org.rakam.collection.SchemaField;
import org.rakam.report.DelegateQueryExecution;
import org.rakam.report.QueryExecution;
import org.rakam.report.QueryExecutor;
import org.rakam.report.QueryResult;
import org.rakam.util.ValidationUtil;

import java.time.LocalDate;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;

import static com.facebook.presto.sql.RakamExpressionFormatter.formatIdentifier;
import static java.time.format.DateTimeFormatter.ISO_DATE;
import static org.rakam.collection.FieldType.INTEGER;
import static org.rakam.util.ValidationUtil.checkCollection;

public class ClickHouseRetentionQueryExecutor
        implements RetentionQueryExecutor
{
    private final QueryExecutor executor;
    private final Metastore metastore;

    @Inject
    public ClickHouseRetentionQueryExecutor(QueryExecutor executor, Metastore metastore)
    {
        this.executor = executor;
        this.metastore = metastore;
    }

    private static String formatExpression(Expression value) {
        return RakamSqlFormatter.formatExpression(value,
                name -> name.getParts().stream().map(e -> formatIdentifier(e, '`')).collect(Collectors.joining(".")),
                name -> name.getParts().stream()
                        .map(e -> formatIdentifier(e, '`')).collect(Collectors.joining(".")), '`');
    }

    @Override
    public QueryExecution query(String project, Optional<RetentionAction> firstAction, Optional<RetentionAction> returningAction, DateUnit dateUnit, Optional<String> dimension, Optional<Integer> period, LocalDate startDate, LocalDate endDate)
    {
        int startEpoch = (int) startDate.toEpochDay();
        int endEpoch = (int) endDate.toEpochDay();

        String firstActionQuery = firstAction.map(action -> String.format("SELECT `$date`, _user, _time %s FROM %s.%s %s",
                dimension.map(e -> "," + e).orElse(""), project, ValidationUtil.checkCollection(action.collection(), '`'),
                action.filter().map(f -> "WHERE " + formatExpression(f)).orElse("")))
                .orElseGet(() -> metastore.getCollectionNames(project).stream()
                        .map(collection -> String.format("SELECT `$date`, _user, _time %s FROM %s.%s",
                                dimension.map(e -> "," + e).orElse(""), project, ValidationUtil.checkCollection(collection, '`')))
                        .collect(Collectors.joining(" UNION ALL ")));

        String query = IntStream.range(startEpoch, endEpoch).boxed().flatMap(epoch -> {
            LocalDate beginDate = LocalDate.ofEpochDay(epoch);
            String date = startDate.format(ISO_DATE);
            String endDateStr = period.map(e -> beginDate.plus(e, dateUnit.getTemporalUnit()))
                    .map(e -> e.isAfter(endDate) ? endDate : e)
                    .orElse(LocalDate.ofEpochDay(endEpoch)).format(ISO_DATE);

            return Stream.of(String.format("select toDate('%s') as date, CAST(-1 AS Int64) as lead, uniq(_user) users from" +
                            " (%s) where `$date` between toDate('%s')  and toDate('%s') " +
                            " group by `$date`", date,
                    firstActionQuery,
                    date, endDateStr),
                    String.format("select toDate('%s') as date, (`$date` - toDate('%s')) - 1 as lead, sum(_user IN (select _user from (%s) WHERE `$date` = toDate('%s'))) as users from" +
                                    " (%s) where `$date` between toDate('%s') and toDate('%s') " +
                                    " group by `$date` order by `$date`", date, date,
                            firstActionQuery,
                            date,
                            firstActionQuery,
                            date, endDateStr));
        }).collect(Collectors.joining(" UNION ALL \n"));

        return new DelegateQueryExecution(executor.executeRawQuery(query), (result) -> {
            if (result.isFailed()) {
                return result;
            }

            Long uniqUser = new Long(-1);

            result.getResult().stream()
                    .filter(objects -> objects.get(1).equals(uniqUser))
                    .forEach(objects -> objects.set(1, null));

            return result;
        });
    }
}

@@ -907,24 +907,63 @@
 Identifier;%0A
+import static java.lang.String.format;%0A
 import stati
@@ -1123,16 +1123,78 @@
 lection;
+%0Aimport static org.rakam.util.ValidationUtil.checkTableColumn;
 %0A%0Apublic
@@ -1603,16 +1603,20 @@
 n value)
+%0A   
  %7B%0A     
@@ -2363,39 +2363,32 @@
 n.map(action -%3E 
-String.
 format(%22SELECT %60
@@ -2759,23 +2759,16 @@
 tion -%3E 
-String.
 format(%22
@@ -3486,23 +3486,16 @@
 ream.of(
-String.
 format(%22
@@ -3501,36 +3501,18 @@
 %22select 
-toDate('%25s') as date
+%25s
 , CAST(-
@@ -3696,23 +3696,131 @@
  %60$date%60
-%22, date
+ %25s%22,%0A                    dimension.map(e -%3E checkTableColumn(e, '%60')).orElse(format(%22toDate('%25s') as date%22, date))
 ,%0A      
@@ -3887,16 +3887,78 @@
 dDateStr
+, dimension.map(e -%3E %22,%22 + checkCollection(e, '%60')).orElse(%22%22)
 ),%0A     
@@ -3972,23 +3972,16 @@
         
-String.
 format(%22
@@ -3987,36 +3987,18 @@
 %22select 
-toDate('%25s') as date
+%25s
 , (%60$dat
@@ -4267,16 +4267,19 @@
  %60$date%60
+ %25s
  order b
@@ -4284,23 +4284,159 @@
  by 
-%60$
+%25s%22,%0A                            dimension.map(e -%3E checkTableColumn(e, '%60')).orElse(format(%22toDate('%25s') as 
 date
-%60
 %22, date
-,
+)),%0A                           
  dat
@@ -4608,16 +4608,198 @@
 dDateStr
+,%0A                            dimension.map(e -%3E %22, %22 + checkCollection(e, '%60')).orElse(%22%22),%0A                            dimension.map(e -%3E checkCollection(e, '%60')).orElse(%22%60$date%60%22)
 ));%0A    

fix bug in click house retention query