package ru.yandex.webmaster3.worker.http.abt;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CompletableFuture;

import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.google.common.collect.Range;
import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.Setter;
import lombok.ToString;
import lombok.extern.slf4j.Slf4j;
import org.joda.time.Instant;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import ru.yandex.webmaster3.core.data.WebmasterHostId;
import ru.yandex.webmaster3.core.util.IdUtils;
import ru.yandex.webmaster3.storage.abt.dao.AbtExperimentsHistoryYDao;
import ru.yandex.webmaster3.storage.abt.dao.AbtHostExperimentYDao;
import ru.yandex.webmaster3.storage.abt.dao.AbtUserExperimentYDao;
import ru.yandex.webmaster3.storage.abt.model.ExperimentInfo;
import ru.yandex.webmaster3.storage.abt.model.ExperimentScope;
import ru.yandex.webmaster3.storage.download.DownloadStatus;
import ru.yandex.webmaster3.storage.util.yt.AsyncTableReader;
import ru.yandex.webmaster3.storage.util.yt.YtException;
import ru.yandex.webmaster3.storage.util.yt.YtPath;
import ru.yandex.webmaster3.storage.util.yt.YtService;
import ru.yandex.webmaster3.storage.util.yt.YtTable;
import ru.yandex.webmaster3.storage.util.yt.YtTableReadDriver;

@Slf4j
@Service
@AllArgsConstructor(onConstructor_ = @Autowired)
public class YtImportExperimentService {
    public static final int BATCH_SIZE = 4096;

    private final YtService ytService;
    private final AbtExperimentsHistoryYDao abtExperimentsHistoryYDao;
    private final AbtHostExperimentYDao abtHostExperimentYDao;
    private final AbtUserExperimentYDao abtUserExperimentYDao;

    public void run(YtImportExperimentService.Data data) {
        CompletableFuture.runAsync(()->startTask(data));
    }

    public void startTask(YtImportExperimentService.Data data){
        Instant time = Instant.now();
        String description = "Started loading exp: " + data.getExperiment()
                + " with action: " + data.getTaskType()
                + " from table: " + data.getYtCluster() + ":" + data.getYtTable()
                + " column: " + data.getYtField();
        abtExperimentsHistoryYDao.insert(
                data.getExperiment(),
                data.getGroup(),
                data.getScope().name(),
                time,
                0L,
                description,
                DownloadStatus.IN_PROGRESS
        );
        log.info(description);

        try {
            RespData resp = process(data);

            String dsc = "Finished loading exp: " + data.getExperiment()
                    + " with action: " + data.getTaskType()
                    + " from table: " + data.getYtCluster() + ":" + data.getYtTable()
                    + " column: " + data.getYtField();
            log.info(dsc);
            abtExperimentsHistoryYDao.update(
                    data.getExperiment(),
                    data.getGroup(),
                    data.getScope().name(),
                    time,
                    resp.getRowCount(),
                    description,
                    DownloadStatus.DONE
            );
        } catch (Exception e) {
            String dsc = "Error while loading exp: " + data.getExperiment()
                    + " with action: " + data.getTaskType()
                    + " from table: " + data.getYtCluster() + ":" + data.getYtTable()
                    + " column: " + data.getYtField();
            log.error(dsc);
            abtExperimentsHistoryYDao.update(
                    data.getExperiment(),
                    data.getGroup(),
                    data.getScope().name(),
                    time,
                    0L,
                    dsc,
                    DownloadStatus.INTERNAL_ERROR
            );

            throw e;
        }
    }

    private RespData process(YtImportExperimentService.Data data) {
        YtPath tablePath = YtPath.create(data.ytCluster, data.ytTable);
        ExperimentInfo experimentInfo = new ExperimentInfo(data.experiment,data.group);
        TaskType taskType = data.getTaskType();
        ExperimentScope scope = data.getScope();
        RespData resp = new RespData();

        ytService.inTransaction(tablePath).execute(ytCypressService -> {
            YtTable table = (YtTable) ytCypressService.getNode(tablePath);
            resp.setRowCount(table.getRowCount());

            int rowsCount = 0;
            Map<WebmasterHostId, ExperimentInfo> hostsBatch = new HashMap<>();
            Map<Long, ExperimentInfo> usersBatch = new HashMap<>();
            AsyncTableReader<ObjectNode> tableReader = new AsyncTableReader<>(
                    ytCypressService,
                    tablePath,
                    Range.all(),
                    YtTableReadDriver.createYSONDriver(ObjectNode.class)
            ).splitInParts(10000L).withRetry(5);
            try (AsyncTableReader.TableIterator<ObjectNode> iter = tableReader.read()) {
                while (iter.hasNext()) {
                    ObjectNode row = iter.next();

                    rowsCount++;
                    if (rowsCount % 100_000 == 0) {
                        log.info("Processed rows: {}", rowsCount);
                    }

                    switch (scope) {
                        case USER:
                            long userId = row.get(data.ytField).asLong();
                            usersBatch.put(userId, experimentInfo);
                            if (usersBatch.size() % BATCH_SIZE == 0) {
                                flushUsersBatch(taskType, usersBatch);
                            }
                            break;

                        case HOST:
                        case DOMAIN:
                            WebmasterHostId hostId = IdUtils.urlToHostId(row.get(data.ytField).asText());
                            hostsBatch.put(hostId, experimentInfo);
                            if (hostsBatch.size() % BATCH_SIZE == 0) {
                                flushHostsBatch(taskType, hostsBatch);
                            }
                            break;

                        default:
                            throw new RuntimeException("Unsupported experiment scope: " + scope);
                    }

                }

                if (!hostsBatch.isEmpty()) {
                    flushHostsBatch(taskType, hostsBatch);
                }

                if (!usersBatch.isEmpty()) {
                    flushUsersBatch(taskType, usersBatch);
                }

                return true;
            } catch (Exception e) {
                log.error("Exception while processing table", e);
                throw new YtException(e);
            }
        });

        return resp;
    }

    private void flushUsersBatch(TaskType taskType, Map<Long, ExperimentInfo> batch) {
        if (taskType == TaskType.ADD){
            abtUserExperimentYDao.insert(batch);
        } else {
            abtUserExperimentYDao.delete(batch);
        }

        batch.clear();
    }

    private void flushHostsBatch(TaskType taskType, Map<WebmasterHostId, ExperimentInfo> batch) {
        if (taskType == TaskType.ADD){
            abtHostExperimentYDao.insert(batch);
        } else {
            abtHostExperimentYDao.delete(batch);
        }

        batch.clear();
    }

    enum TaskType {
        ADD, REMOVE
    }

    @Getter
    @ToString
    public static class Data {
        final String experiment;
        final String group;
        final TaskType taskType;
        final String ytCluster;
        final String ytTable;
        final String ytField;
        final ExperimentScope scope;


        public Data(
                @JsonProperty("experiment") String experiment, @JsonProperty("group") String group,
                @JsonProperty("ytCluster") String ytCluster, @JsonProperty("ytTable") String ytTable,
                @JsonProperty("ytField") String ytField, @JsonProperty("scope")ExperimentScope scope,
                @JsonProperty("ytField") TaskType taskType
        ) {
            this.experiment = experiment;
            this.group = group;
            this.ytCluster = ytCluster;
            this.ytTable = ytTable;
            this.ytField = ytField;
            this.scope = scope;
            this.taskType = taskType;
        }
    }

    @Getter
    @Setter
    public static class RespData {
       private Long rowCount;
    }
}
