package ru.yandex.webmaster3.worker.searchurl;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.NavigableMap;
import java.util.Optional;
import java.util.TreeMap;
import java.util.UUID;
import java.util.stream.Collectors;

import com.datastax.driver.core.utils.UUIDs;
import lombok.Setter;
import org.joda.time.Instant;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import ru.yandex.webmaster3.core.WebmasterException;
import ru.yandex.webmaster3.core.http.WebmasterErrorResponse;
import ru.yandex.webmaster3.core.worker.task.PeriodicTaskState;
import ru.yandex.webmaster3.core.worker.task.PeriodicTaskType;
import ru.yandex.webmaster3.core.worker.task.TaskResult;
import ru.yandex.webmaster3.storage.clickhouse.replication.MdbClickhouseReplicationManager;
import ru.yandex.webmaster3.storage.clickhouse.replication.data.ClickhouseReplicationCommand;
import ru.yandex.webmaster3.storage.clickhouse.replication.data.ClickhouseReplicationPriority;
import ru.yandex.webmaster3.storage.clickhouse.replication.data.ClickhouseReplicationTaskGroup;
import ru.yandex.webmaster3.storage.searchurl.offline.dao.SearchBaseImportTablesRepository;
import ru.yandex.webmaster3.storage.searchurl.offline.data.SearchBaseImportInfo;
import ru.yandex.webmaster3.storage.searchurl.offline.data.SearchBaseImportStageEnum;
import ru.yandex.webmaster3.storage.searchurl.offline.data.SearchBaseImportTaskType;
import ru.yandex.webmaster3.storage.util.clickhouse2.ClickhouseException;
import ru.yandex.webmaster3.storage.util.clickhouse2.ClickhouseHost;
import ru.yandex.webmaster3.storage.util.clickhouse2.ClickhouseServer;
import ru.yandex.webmaster3.storage.util.yt.YtException;
import ru.yandex.webmaster3.storage.util.yt.YtMapReduceCommand;
import ru.yandex.webmaster3.storage.util.yt.YtNode;
import ru.yandex.webmaster3.storage.util.yt.YtNodeAttributes;
import ru.yandex.webmaster3.storage.util.yt.YtOperationId;
import ru.yandex.webmaster3.storage.util.yt.YtPath;
import ru.yandex.webmaster3.storage.util.yt.YtService;
import ru.yandex.webmaster3.storage.util.yt.YtUtils;
import ru.yandex.webmaster3.storage.ytimport.ImportPriority;
import ru.yandex.webmaster3.storage.ytimport.MdbYtClickhouseImportManager;
import ru.yandex.webmaster3.storage.ytimport.YtClickhouseImportCommand;
import ru.yandex.webmaster3.storage.ytimport.YtClickhouseImportState;
import ru.yandex.webmaster3.storage.ytimport.YtClickhouseImportStateEnum;
import ru.yandex.webmaster3.storage.ytimport.YtClickhouseTableRelation;
import ru.yandex.webmaster3.worker.PeriodicTask;
import ru.yandex.webmaster3.worker.TaskSchedule;

/**
 * @author avhaliullin
 */
public abstract class AbstractSearchBaseImportTask extends PeriodicTask<PeriodicTaskState> {
    private static final Logger log = LoggerFactory.getLogger(AbstractSearchBaseImportTask.class);

    private static final String ATTR_SOURCE_TABLE_BASE_DATE = "source_base_timestamp";
    private static final String ACCEPTANCE_SUFFIX = ".acceptance";

    private static final int PARTS_COUNT = 6;

    private final SearchBaseImportTaskType taskType;
    private final String dbName;
    private final String insertSpec;

    @Setter
    private YtService ytService;
    @Setter
    private YtPath sourceTablePath;
    @Setter
    private YtPath ytPrepareExecutable;
    @Setter
    private YtPath ytWorkDir;
    @Setter
    private ClickhouseServer clickhouseServer;
    @Setter
    private SearchBaseImportTablesRepository searchBaseImportTablesCDao;
    @Setter
    private MdbYtClickhouseImportManager ytClickhouseImportManager;
    @Setter
    private MdbClickhouseReplicationManager clickhouseReplicationManager;
    @Setter
    private PeriodicTaskType type;

    protected AbstractSearchBaseImportTask(SearchBaseImportTaskType taskType, String dbName, String insertSpec) {
        this.taskType = taskType;
        this.dbName = dbName;
        this.insertSpec = insertSpec;
    }

    protected abstract String distributedTableName(Instant searchBaseDate);

    protected abstract String partTableName(Instant searchBaseDate, int part);

    protected abstract String partTableCreateSpec(Instant searchBaseDate, int part);

    protected abstract String shardTableCreateQuery(SearchBaseImportInfo importInfo);

    protected abstract String distributedTableCreateQuery(SearchBaseImportInfo importInfo);

    protected abstract String symlinkTableDropQuery(SearchBaseImportInfo importInfo);

    protected abstract String symlinkTableCreateQuery(SearchBaseImportInfo importInfo);

    @Override
    public Result run(UUID runId) throws Exception {
        SearchBaseImportInfo importInfo = searchBaseImportTablesCDao.getLatestBaseInfo(taskType);
        if (importInfo == null || !importInfo.getStage().isToBeProcessed()) {
            NavigableMap<Instant, YtPath> latestBases = discoverLatestBases();
            for (Map.Entry<Instant, YtPath> entry : latestBases.entrySet()) {
                Instant newBaseDate = entry.getKey();
                YtPath sourceTable = entry.getValue();
                if (importInfo == null || newBaseDate.isAfter(importInfo.getSearchBaseDate())) {
                    String dc = pickAliveDC();
                    int shardsCount = getHostsForDC(dc).size();
                    List<UUID> importIds = new ArrayList<>();
                    for (int i = 0; i < PARTS_COUNT; i++) {
                        importIds.add(UUIDs.timeBased());
                    }
                    importInfo = new SearchBaseImportInfo(
                            taskType,
                            newBaseDate,
                            dbName,
                            SearchBaseImportStageEnum.DISCOVERED,
                            distributedTableName(newBaseDate),
                            shardsCount,
                            PARTS_COUNT,
                            dc,
                            importIds,
                            false,
                            sourceTable,
                            null,
                            null
                    );
                    log.info("Found new base {} in {}", newBaseDate, sourceTable);
                    syncState(importInfo);
                    break;
                }
            }
        }
        if (importInfo == null || !importInfo.getStage().isToBeProcessed()) {
            log.info("Nothing new found");
            return new Result(TaskResult.SUCCESS);
        }

        if (importInfo.getSourceTable() == null) {
            importInfo = importInfo.withSourceTable(getAcceptanceSourceTable());
            syncState(importInfo);
        }

        while (importInfo.getStage().isToBeProcessed()) {
            switch (importInfo.getStage()) {
                case DISCOVERED:
                    importInfo = prepareTables(importInfo);
                    break;
                case IMPORTING:
                    importInfo = importData(importInfo);
                    break;
                case IMPORTED:
                    importInfo = replicateTables(importInfo);
                    break;
                case REPLICATED:
                    importInfo = createViewTables(importInfo);
                    break;
                default:
                    throw new RuntimeException("Unknown stage " + importInfo.getStage());
            }
        }
        return new Result(TaskResult.SUCCESS);
    }

    private SearchBaseImportInfo createViewTables(SearchBaseImportInfo importInfo) throws ClickhouseException {
        clickhouseServer.executeOnAllHosts(shardTableCreateQuery(importInfo));
        clickhouseServer.executeOnAllHosts(distributedTableCreateQuery(importInfo));
        clickhouseServer.executeOnAllHosts(symlinkTableDropQuery(importInfo));
        clickhouseServer.executeOnAllHosts(symlinkTableCreateQuery(importInfo));
        return syncState(importInfo.withStage(SearchBaseImportStageEnum.READY));
    }

    private SearchBaseImportInfo replicateTables(SearchBaseImportInfo importInfo) throws ClickhouseException {
        UUID replicationTask = importInfo.getReplicationManagerTask();
        if (replicationTask == null) {
            replicationTask = UUIDs.timeBased();
            importInfo = importInfo.withReplicationManagerTask(replicationTask);
            syncState(importInfo);
            List<ClickhouseReplicationCommand.TableInfo> tables = new ArrayList<>();
            for (int shard = 0; shard < importInfo.getShardsCount(); shard++) {
                for (int part = 0; part < PARTS_COUNT; part++) {
                    tables.add(new ClickhouseReplicationCommand.TableInfo(
                            partTableName(importInfo.getSearchBaseDate(), part),
                            partTableCreateSpec(importInfo.getSearchBaseDate(), part),
                            shard
                    ));
                }
            }
            ClickhouseReplicationCommand replicationCommand = new ClickhouseReplicationCommand(
                    replicationTask,
                    importInfo.getDbName(),
                    ClickhouseReplicationPriority.OFFLINE,
                    tables
            );
            clickhouseReplicationManager.enqueueReplication(replicationCommand);
            log.info("Started replication task {}", replicationTask);
        }
        ClickhouseReplicationTaskGroup taskState = clickhouseReplicationManager.waitForCompletionAndGetState(replicationTask);
        if (taskState == null) {
            String message = "Replication task " + replicationTask + " was lost";
            log.error(message);
            throw new RuntimeException(message);
        }
        return syncState(importInfo.withStage(SearchBaseImportStageEnum.REPLICATED));
    }

    private SearchBaseImportInfo importData(SearchBaseImportInfo importInfo)  {
        UUID importManagerTask = importInfo.getImportManagerTask();
        if (importManagerTask == null) {
            importManagerTask = UUIDs.timeBased();
            importInfo = importInfo.withImportManagerTask(importManagerTask);
            syncState(importInfo);
            List<YtClickhouseTableRelation> tables = new ArrayList<>();
            for (int shard = 0; shard < importInfo.getShardsCount(); shard++) {
                for (int part = 0; part < PARTS_COUNT; part++) {
                    tables.add(new YtClickhouseTableRelation(
                            preparedTablePath(importInfo.getSearchBaseDate(), shard, part),
                            shard,
                            partTableName(importInfo.getSearchBaseDate(), part),
                            partTableCreateSpec(importInfo.getSearchBaseDate(), part)
                    ));
                }
            }
            YtClickhouseImportCommand importCommand = new YtClickhouseImportCommand(
                    importManagerTask,
                    tables,
                    importInfo.getDbName(),
                    insertSpec,
                    ImportPriority.OFFLINE
            );
            ytClickhouseImportManager.startImport(importCommand);
            log.info("Started import task {}", importManagerTask);
        }
        YtClickhouseImportState importState = ytClickhouseImportManager.waitForCompletionAndGetState(importManagerTask);

        if (importState == null || importState.getState() == YtClickhouseImportStateEnum.FAILED) {
            String message = importState == null
                    ? "Import task " + importManagerTask + " was lost"
                    : "Import task " + importManagerTask + " was failed";
            log.error(message);
            throw new RuntimeException(message);
        }
        return syncState(importInfo.withStage(SearchBaseImportStageEnum.IMPORTED));
    }

    /**
     * Удаляет таблицы с промежуточными данными (prepared)
     */
    public void dropPreparedTables(SearchBaseImportInfo importInfo) throws WebmasterException {
        try {
            ytService.withoutTransaction(cypressService -> {
                // смотрим, где лежит наше добро
                YtPath path = preparedTableBasePath(importInfo.getSearchBaseDate());
                if (cypressService.exists(path)) {
                    cypressService.remove(path, true);
                }
                return true;
            });
        } catch (Exception e) {
            throw new WebmasterException("Failed to drop prepared tables",
                    new WebmasterErrorResponse.YTServiceErrorResponse(getClass(), e), e);
        }
    }

    private SearchBaseImportInfo prepareTables(SearchBaseImportInfo importInfo) throws YtException {
        List<YtPath> outTables = preparedTables(importInfo);
        return ytService.inTransaction(importInfo.getSourceTable()).query(cypressService -> {
            YtUtils.recreateTables(cypressService, outTables, new YtNodeAttributes().setCompressionCodec("none"));
            YtMapReduceCommand mrCommand = YtUtils.newPrepareTablesForImportBuilder()
                    .addInputTable(importInfo.getSourceTable())
                    .setOutputTables(outTables)
                    .setTask(importInfo.getTaskType().getPrepareTaskName())
                    .setBinary(ytPrepareExecutable)
                    .setMapperMemoryLimit(1073741824L)
                    .setReducerMemoryLimit(1073741824L)
                    .buildV2();
            YtOperationId operationId = cypressService.mapReduce(mrCommand);
            if (!cypressService.waitFor(operationId)) {
                throw new RuntimeException("YT task failed " + operationId);
            }
            return syncState(importInfo.withStage(SearchBaseImportStageEnum.IMPORTING));
        });
    }

    private NavigableMap<Instant, YtPath> discoverLatestBases() throws YtException, InterruptedException {
        return ytService.withoutTransactionQuery(cypressService -> {
            List<YtPath> candidates = new ArrayList<>();
            candidates.add(getProductionSourceTable());
            candidates.add(getAcceptanceSourceTable());
            TreeMap<Instant, YtPath> result = new TreeMap<>();
            for (YtPath table : candidates) {
                YtNode sourceNode = cypressService.getNode(table);
                result.put(new Instant(sourceNode.getNodeMeta().get(ATTR_SOURCE_TABLE_BASE_DATE).asLong() * 1000L), table);
            }
            return result;
        });
    }

    private SearchBaseImportInfo syncState(SearchBaseImportInfo importInfo)  {
        log.info("Storing stage {} for import {} from base {}", importInfo.getStage(), importInfo.getTaskType(), importInfo.getSearchBaseDate());
        searchBaseImportTablesCDao.insertStage(importInfo);
        importInfo = searchBaseImportTablesCDao.getBaseInfo(importInfo.getTaskType(), importInfo.getSearchBaseDate());
        if (importInfo.isPaused()) {
            throw new WebmasterException("Task was paused: " + importInfo.getTaskType() + " for base " + importInfo.getSearchBaseDate(),
                    new WebmasterErrorResponse.InternalUnknownErrorResponse(getClass(), null));
        }
        return importInfo;
    }

    private String pickAliveDC() {
        Optional<String> dcO = clickhouseServer.pickAliveDC();
        return dcO.orElseThrow(() -> new RuntimeException("Not found alive clickhouse DC"));
    }

    private List<ClickhouseHost> getHostsForDC(String dc) {
        return clickhouseServer.getHosts()
                .stream().filter(host -> host.getDcName().equals(dc))
                .collect(Collectors.toList());

    }

    private List<YtPath> preparedTables(SearchBaseImportInfo importInfo) {
        List<YtPath> result = new ArrayList<>();
        for (int shard = 0; shard < importInfo.getShardsCount(); shard++) {
            for (int part = 0; part < importInfo.getPartsCount(); part++) {
                result.add(preparedTablePath(importInfo.getSearchBaseDate(), shard, part));
            }
        }
        return result;
    }

    private YtPath preparedTablePath(Instant baseDate, int shard, int part) {
        return YtPath.path(preparedTableBasePath(baseDate), String.format("shard%02d-part%02d", shard, part));
    }

    private YtPath preparedTableBasePath(Instant baseDate) {
        return YtPath.path(ytWorkDir, String.valueOf(baseDate.getMillis() / 1000L));
    }

    private YtPath getProductionSourceTable() {
        return sourceTablePath;
    }

    private YtPath getAcceptanceSourceTable() {
        return YtPath.path(sourceTablePath.getParent(), sourceTablePath.getName() + ACCEPTANCE_SUFFIX);
    }

    @Override
    public TaskSchedule getSchedule() {
        return TaskSchedule.startByCron("0 * * * * *");
    }

    @Override
    public PeriodicTaskType getType() {
        return type;
    }
}
