package ru.yandex.webmaster3.storage.util.yt;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.URI;
import java.util.Iterator;
import java.util.Optional;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

import com.google.common.base.Stopwatch;
import com.google.common.io.CountingInputStream;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.http.Header;
import org.apache.http.HttpEntity;
import org.apache.http.HttpHeaders;
import org.apache.http.HttpRequest;
import org.apache.http.StatusLine;
import org.apache.http.client.config.RequestConfig;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Required;
import org.springframework.web.util.UriComponentsBuilder;

import ru.yandex.webmaster3.core.http.HttpConstants;
import ru.yandex.webmaster3.core.util.ProgressUtils;
import ru.yandex.webmaster3.storage.util.LenvalStreamParser;
import ru.yandex.webmaster3.storage.util.ProgressLogInputStream;

/**
 * @author aherman
 */
public class YtTableReaderService {
    private static final Logger log = LoggerFactory.getLogger(YtTableReaderService.class);

    private static final String FORMAT_YAMR_LENVAL = "application/x-yamr-lenval";
    private static final String ENCODING_SNAPPY = "y-snappy";

    private static final String REQUEST_READ = "/api/v3/read_table";

    private static final String PARAMETER_PATH = "path";
    private static final String PARAMETER_TRANSACTION_ID = "transaction_id";

    private static final String HEADER_AUTHORIZATION = "Authorization";

    private static final String HEADER_X_YT_PROXY = "X-YT-Proxy";
    private static final String HEADER_X_YT_REQUEST_ID = "X-YT-Request-Id";
    private static final String HEADER_X_YT_SOCKET_ID = "X-YT-Socket-Id";

    private static final String HEADER_X_YT_RESPONSE_PARAMETERS = "X-YT-Response-Parameters";

    // Trailing
    private static final String HEADER_X_YT_ERROR = "X-YT-Error";
    private static final String HEADER_X_YT_RESPONSE_CODE = "X-YT-Response-Code";
    private static final String HEADER_X_YT_RESPONSE_MESSAGE = "X-YT-Response-Message";

    private static final int RETRY_COUNT = 3;

    private URI ytProxyUri;
    private int socketTimeout = 5 * 60 * 1000;
    private int connectTimeout = HttpConstants.DEFAULT_CONNECT_TIMEOUT;
    private String oAuthToken;
    private File localCacheFolder;

    public void init() {
        if (!localCacheFolder.exists() && !localCacheFolder.mkdirs()) {
            throw new RuntimeException("Unable to create YT cache folder");
        }
        File[] files = localCacheFolder.listFiles();
        for (File file : files) {
            log.info("Clean old cache file: {}", file.getAbsolutePath());
            file.delete();
        }
    }

    public void readTableCached(String tablePath, Pair<Long, Long> range, long subrangeSize,
            TableProcessor tableProcessor) throws Exception
    {
        RequestConfig requestConfig = RequestConfig.custom()
                .setConnectTimeout(connectTimeout)
                .setSocketTimeout(socketTimeout)
                .setContentCompressionEnabled(false)
                .build();

        ExecutorService cachers = Executors.newSingleThreadExecutor(
                new ThreadFactoryBuilder()
                        .setDaemon(true)
                        .setNameFormat("yt-cacher-%d")
                        .build()
        );
        BlockingQueue<CachedTablePart> cacheQueue = new ArrayBlockingQueue<CachedTablePart>(3);
        try (CloseableHttpClient httpClient = HttpClients.custom()
                .setDefaultRequestConfig(requestConfig)
                .build())
        {
            tableProcessor.setTableSize(range.getRight() - range.getLeft());
            cachers.execute(new Runnable() {
                @Override
                public void run() {
                    log.info("Start cacher");
                    Iterator<Pair<Long, Long>> subrangeIt = generateSubranges(range, subrangeSize);

                    String errorMessage = null;
                    YtStatus errorStatus = YtStatus.UNKNOWN;

                    try {
                        while (subrangeIt.hasNext()) {
                            Pair<Long, Long> subrange = subrangeIt.next();
                            File cacheFile = getCacheFile(tablePath, subrange);
                            if (cacheFile.exists()) {
                                log.warn("Cache file exists, delete: {}", cacheFile.getAbsolutePath());
                                cacheFile.delete();
                            }
                            HttpGet httpRequest = createReadRequest(tablePath, null, Optional.of(subrange));
                            log.info("YT request: {}", httpRequest.getURI());
                            int iteration = 0;

                            Stopwatch downloadTime = Stopwatch.createStarted();
                            boolean compressed = false;
                            long size = 0;
                            while (iteration < RETRY_COUNT) {
                                try (CloseableHttpResponse httpResponse = httpClient.execute(httpRequest)) {
                                    YtStatus ytStatus = getYtStatus(httpResponse);
                                    Header[] allHeaders = httpResponse.getAllHeaders();
                                    Header contentEncoding = httpResponse.getFirstHeader(HttpHeaders.CONTENT_ENCODING);
                                    if (contentEncoding != null) {
                                        compressed = ENCODING_SNAPPY.equalsIgnoreCase(contentEncoding.getValue());
                                    }
                                    String responseParameters = "";
                                    logYtResponseParameters(ytStatus, allHeaders, responseParameters);
                                    HttpEntity entity = httpResponse.getEntity();
                                    errorStatus = ytStatus;
                                    if (!isSuccess(ytStatus)) {
                                        log.error("Unable to download table part: {} {}-{} {}", tablePath,
                                                subrange.getLeft(),
                                                subrange.getRight(), ytStatus);
                                        iteration++;
                                        continue;
                                    } else if (entity == null) {
                                        log.error("Empty response body: {} {}", tablePath, ytStatus);
                                        iteration++;
                                        continue;
                                    } else {
                                        log.info("Cache table subrange: {} {}-{} {}", tablePath, subrange.getLeft(),
                                                subrange.getRight(), cacheFile.getAbsolutePath());
                                        try (BufferedOutputStream bos = new BufferedOutputStream(
                                                new FileOutputStream(cacheFile)))
                                        {
                                            size = IOUtils.copyLarge(
                                                    new ProgressLogInputStream(
                                                            new CountingInputStream(entity.getContent()),
                                                            "Cache " + getTableName(tablePath,
                                                                    Optional.of(subrange))),
                                                    bos
                                            );
                                        }
                                        break;
                                    }
                                } catch (IOException e) {
                                    log.error("IO error when caching table part: {} {}-{}", tablePath,
                                            subrange.getLeft(),
                                            subrange.getRight(), e);
                                }
                                log.error("Retry table download: {} {}-{}", tablePath, subrange.getLeft(),
                                        subrange.getRight());
                                iteration++;
                                Thread.sleep(TimeUnit.SECONDS.toMillis(30));
                            }
                            if (iteration == RETRY_COUNT) {
                                log.error("Retries exceeded, stop");
                                errorMessage = "Retry exceeded";
                                break;
                            }
                            downloadTime.stop();
                            log.info("Subrange downloaded: {} {}-{} {} {}", tablePath,
                                    subrange.getLeft(), subrange.getRight(),
                                    size, ProgressUtils.formatTime(downloadTime.elapsed(TimeUnit.MILLISECONDS))
                            );
                            cacheQueue.put(new CachedTablePart(cacheFile, subrange, errorStatus, compressed));
                        }
                    } catch (Exception e) {
                        errorMessage = e.getMessage();
                        log.error("Error in cacher", e);
                    }
                    try {
                        if (errorMessage != null) {
                            cacheQueue.put(new CachedTablePart(null, true, errorMessage, errorStatus));
                        } else {
                            cacheQueue.put(new CachedTablePart());
                        }
                    } catch (InterruptedException e) {
                        log.warn("Interrupted");
                    }
                    log.info("Stop cacher");
                }
            });
            while (true) {
                CachedTablePart tablePart = cacheQueue.take();
                if (tablePart.end) {
                    if (tablePart.error) {
                        throw new YtException(tablePart.errorMessage);
                    }
                    break;
                }
                log.info("Start subrange processing: {} {}-{}",
                        tablePath,
                        tablePart.subrange.getLeft(), tablePart.subrange.getRight()
                );
                Stopwatch processTime = Stopwatch.createStarted();
                try (InputStream is = new BufferedInputStream(new FileInputStream(tablePart.cacheFile))) {
                    InputStream uncompressedIs = is;
                    if (tablePart.compressed) {
                        uncompressedIs = new YtSnappyCompressedStream(is);
                    }
                    CountingInputStream cis = new CountingInputStream(uncompressedIs);
                    String tableName = getTableName(tablePath, Optional.of(tablePart.subrange));
                    ProgressLogInputStream pis = new ProgressLogInputStream(cis, "Read " + tableName);
                    LenvalStreamParser streamParser = new LenvalStreamParser(pis);
                    streamParser.process(tableProcessor);
                } catch (IOException e) {
                    log.error("Unable to read cached table", e);
                } finally {
                    if (tablePart.cacheFile.exists()) {
                        tablePart.cacheFile.delete();
                    }
                }
                log.info("Subrange processed: {} {}-{} {}",
                        tablePath,
                        tablePart.subrange.getLeft(), tablePart.subrange.getRight(),
                        ProgressUtils.formatTime(processTime.elapsed(TimeUnit.MILLISECONDS))
                );
            }

            tableProcessor.finish();
        }
    }

    private static class CachedTablePart {
        private final File cacheFile;
        private final Pair<Long, Long> subrange;
        private final boolean end;
        private final boolean error;
        private final String errorMessage;
        private final YtStatus status;
        private final boolean compressed;

        public CachedTablePart() {
            this.cacheFile = null;
            this.subrange = null;
            this.end = true;
            this.compressed = false;
            this.status = null;
            this.error = false;
            this.errorMessage = null;
        }

        public CachedTablePart(File cacheFile, Pair<Long, Long> subrange, YtStatus status, boolean compressed) {
            this.cacheFile = cacheFile;
            this.subrange = subrange;
            this.end = false;
            this.status = status;
            this.compressed = compressed;
            this.error = false;
            this.errorMessage = null;
        }

        public CachedTablePart(Pair<Long, Long> subrange, boolean error, String errorMessage, YtStatus status) {
            this.cacheFile = null;
            this.subrange = subrange;
            this.end = true;
            this.compressed = false;
            this.status = status;
            this.error = error;
            this.errorMessage = errorMessage;
        }
    }

    private File getCacheFile(String tablePath, Pair<Long, Long> subrange) {
        String cacheFileName = getTableName(tablePath) + "_" + subrange.getLeft() + "_" + subrange.getRight();
        return new File(localCacheFolder, cacheFileName);
    }

    private YtStatus getYtStatus(CloseableHttpResponse httpResponse) {
        StatusLine statusLine = httpResponse.getStatusLine();
        return YtStatus.R.fromValueOrUnknown(statusLine.getStatusCode());
    }

    private static String getTableName(String tablePath, Optional<Pair<Long, Long>> range) {
        String name = getTableName(tablePath);
        if (range.isPresent()) {
            Pair<Long, Long> pair = range.get();
            name += " " + pair.getLeft() + ":" + pair.getRight();
        }
        return name;
    }

    @NotNull
    private static String getTableName(String tablePath) {
        int slash = tablePath.lastIndexOf('/');
        String name = tablePath;
        if (slash >= 0) {
            name = tablePath.substring(slash);
        }
        return name;
    }

    private String logYtResponseParameters(YtStatus ytStatus, Header[] allHeaders, String responseParameters) {
        StringBuilder sb = new StringBuilder(32);
        for (Header header : allHeaders) {
            if (HEADER_X_YT_RESPONSE_PARAMETERS.equals(header.getName())) {
                responseParameters = header.getValue();
            }
            if (header.getName().startsWith("X-YT")
                    || HttpHeaders.CONTENT_ENCODING.equalsIgnoreCase(header.getName())
                    || HttpHeaders.CONTENT_TYPE.equalsIgnoreCase(header.getName()))
            {
                sb.append(' ').append(header.getName()).append('=').append(header.getValue());
            }
        }
        log.info("YT response: {} {}", ytStatus, sb);
        return responseParameters;
    }

    protected HttpGet createReadRequest(String tablePath, String transactionId, Optional<Pair<Long, Long>> range) {
        if (range.isPresent()) {
            tablePath = tablePath + "[#" + range.get().getLeft() + ":#" + range.get().getRight() + "]";
        }
        UriComponentsBuilder builder = UriComponentsBuilder.fromUri(ytProxyUri)
                .path(REQUEST_READ)
                .queryParam(PARAMETER_PATH, tablePath);
        if (!StringUtils.isEmpty(transactionId)) {
            builder.queryParam(PARAMETER_TRANSACTION_ID, transactionId);
        }
        URI requestUri = builder.build().toUri();
        HttpGet request = new HttpGet(requestUri);
        request.setHeader(HttpHeaders.ACCEPT, FORMAT_YAMR_LENVAL);
        request.setHeader(HttpHeaders.ACCEPT_ENCODING, ENCODING_SNAPPY);
        setAuthToken(request);
        return request;
    }

    private static Iterator<Pair<Long, Long>> generateSubranges(Pair<Long, Long> range, long subrangeSize) {
        return new Iterator<Pair<Long, Long>>() {
            long nextStart = range.getLeft();
            @Override
            public boolean hasNext() {
                return nextStart < range.getRight();
            }

            @Override
            public Pair<Long, Long> next() {
                long nextEnd = Math.min(nextStart + subrangeSize, range.getRight());
                Pair<Long, Long> result = Pair.of(nextStart, nextEnd);
                nextStart = nextEnd;
                return result;
            }
        };
    }

    private void setAuthToken(HttpRequest request) {
        request.setHeader(HEADER_AUTHORIZATION, "OAuth " + oAuthToken);
    }

    protected static boolean isSuccess(YtStatus status) {
        return status == YtStatus.YT_200_OK || status == YtStatus.YT_202_SUCCESS;
    }

    public interface TableProcessor extends LenvalStreamParser.EntryProcessor {
        public void setTableSize(long linesCount);
        public void finish() throws Exception;
    }

    @Required
    public void setYtProxyUri(URI ytProxyUri) {
        this.ytProxyUri = ytProxyUri;
    }

    public void setSocketTimeout(int socketTimeout) {
        this.socketTimeout = socketTimeout;
    }

    public void setConnectTimeout(int connectTimeout) {
        this.connectTimeout = connectTimeout;
    }

    @Required
    public void setoAuthToken(String oAuthToken) {
        this.oAuthToken = oAuthToken;
    }

    @Required
    public void setLocalCacheFolder(File localCacheFolder) {
        this.localCacheFolder = localCacheFolder;
    }
}
