/*
 * Aiven Kafka GCS Connector
 * Copyright (c) 2019 Aiven Oy
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

package io.aiven.kafka.connect.gcs.testutils;

import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Base64;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;
import java.util.zip.GZIPInputStream;

import io.aiven.kafka.connect.gcs.config.CompressionType;

import com.github.luben.zstd.ZstdInputStream;
import com.google.cloud.storage.Blob;
import com.google.cloud.storage.BlobInfo;
import com.google.cloud.storage.Storage;
import org.xerial.snappy.SnappyInputStream;

public final class BucketAccessor {
    private final Storage storage;
    private final String bucketName;
    private final boolean cache;

    private List<String> blobNamesCache = null;
    private final Map<String, String> stringContentCache = new HashMap<>();
    private final Map<String, List<String>> linesCache = new HashMap<>();
    private final Map<String, List<List<String>>> decodedLinesCache = new HashMap<>();

    public BucketAccessor(final Storage storage,
                          final String bucketName,
                          final boolean cache) {
        Objects.requireNonNull(storage, "storage cannot be null");
        Objects.requireNonNull(bucketName, "bucketName cannot be null");

        this.storage = storage;
        this.bucketName = bucketName;
        this.cache = cache;
    }

    public BucketAccessor(final Storage storage,
                          final String bucketName) {
        this(storage, bucketName, false);
    }

    public final void ensureWorking() {
        if (storage.get(bucketName) == null) {
            throw new RuntimeException("Cannot access GCS bucket \"" + bucketName + "\"");
        }
    }

    public final List<String> getBlobNames() {
        if (cache) {
            if (blobNamesCache == null) {
                blobNamesCache = getBlobNames0();
            }
            return blobNamesCache;
        } else {
            return getBlobNames0();
        }
    }

    private List<String> getBlobNames0() {
        return StreamSupport.stream(storage.list(bucketName).iterateAll().spliterator(), false)
            .map(BlobInfo::getName)
            .sorted()
            .collect(Collectors.toList());
    }

    /**
     * Get blob names with the prefix.
     *
     * <p>Doesn't support caching.
     */
    public final List<String> getBlobNames(final String prefix) {
        Objects.requireNonNull(prefix, "prefix cannot be null");

        final Storage.BlobListOption blobListOption = Storage.BlobListOption.prefix(prefix);
        return StreamSupport.stream(storage.list(bucketName, blobListOption).iterateAll().spliterator(), false)
            .map(BlobInfo::getName)
            .sorted()
            .collect(Collectors.toList());
    }

    public final void clear(final String prefix) {
        Objects.requireNonNull(prefix, "prefix cannot be null");

        final Storage.BlobListOption blobListOption = Storage.BlobListOption.prefix(prefix);
        for (final Blob blob : storage.get(bucketName).list(blobListOption).iterateAll()) {
            assert blob.delete();
        }

        if (cache) {
            blobNamesCache = null;
            stringContentCache.clear();
            linesCache.clear();
            decodedLinesCache.clear();
        }
    }

    public final String readStringContent(final String blobName,
                                          final String compression) {
        Objects.requireNonNull(blobName, "blobName cannot be null");
        if (cache) {
            return stringContentCache.computeIfAbsent(
                blobName, k -> readStringContent0(blobName, compression));
        } else {
            return readStringContent0(blobName, compression);
        }
    }

    private String readStringContent0(final String blobName,
                                      final String compression) {
        final byte[] blobBytes = storage.readAllBytes(bucketName, blobName);
        try (final ByteArrayInputStream bais = new ByteArrayInputStream(blobBytes);
             final InputStream decompressedStream = getDecompressedStream(bais, compression);
             final InputStreamReader reader = new InputStreamReader(decompressedStream, StandardCharsets.UTF_8);
             final BufferedReader bufferedReader = new BufferedReader(reader)) {
            return bufferedReader.readLine();
        } catch (final IOException e) {
            throw new RuntimeException(e);
        }
    }

    public final List<String> readLines(final String blobName, final String compression) {
        Objects.requireNonNull(blobName, "blobName cannot be null");
        if (cache) {
            return linesCache.computeIfAbsent(blobName, k -> readLines0(blobName, compression));
        } else {
            return readLines0(blobName, compression);
        }
    }

    private List<String> readLines0(final String blobName, final String compression) {
        Objects.requireNonNull(blobName, "blobName cannot be null");
        final byte[] blobBytes = storage.readAllBytes(bucketName, blobName);
        try (final ByteArrayInputStream bais = new ByteArrayInputStream(blobBytes);
             final InputStream decompressedStream = getDecompressedStream(bais, compression);
             final InputStreamReader reader = new InputStreamReader(decompressedStream, StandardCharsets.UTF_8);
             final BufferedReader bufferedReader = new BufferedReader(reader)) {

            return bufferedReader.lines().collect(Collectors.toList());
        } catch (final IOException e) {
            throw new RuntimeException(e);
        }
    }

    private InputStream getDecompressedStream(final InputStream inputStream, final String compression)
            throws IOException {
        Objects.requireNonNull(inputStream, "inputStream cannot be null");
        Objects.requireNonNull(compression, "compression cannot be null");

        final CompressionType compressionType = CompressionType.forName(compression);
        switch (compressionType) {
            case ZSTD:
                return new ZstdInputStream(inputStream);
            case GZIP:
                return new GZIPInputStream(inputStream);
            case SNAPPY:
                return new SnappyInputStream(inputStream);
            default:
                return inputStream;
        }
    }

    public final List<List<String>> readAndDecodeLines(final String blobName,
                                                       final String compression,
                                                       final int... fieldsToDecode) {
        Objects.requireNonNull(blobName, "blobName cannot be null");
        Objects.requireNonNull(fieldsToDecode, "fieldsToDecode cannot be null");

        if (cache) {
            return decodedLinesCache.computeIfAbsent(
                blobName, k -> readAndDecodeLines0(blobName, compression, fieldsToDecode));
        } else {
            return readAndDecodeLines0(blobName, compression, fieldsToDecode);
        }
    }

    private List<List<String>> readAndDecodeLines0(final String blobName,
                                                   final String compression,
                                                   final int[] fieldsToDecode) {
        return readLines(blobName, compression).stream()
            .map(l -> l.split(","))
            .map(fields -> decodeRequiredFields(fields, fieldsToDecode))
            .collect(Collectors.toList());
    }

    private List<String> decodeRequiredFields(final String[] originalFields, final int[] fieldsToDecode) {
        Objects.requireNonNull(originalFields, "originalFields cannot be null");
        Objects.requireNonNull(fieldsToDecode, "fieldsToDecode cannot be null");

        final List<String> result = Arrays.asList(originalFields);
        for (final int fieldIdx : fieldsToDecode) {
            result.set(fieldIdx, b64Decode(result.get(fieldIdx)));
        }
        return result;
    }

    private String b64Decode(final String value) {
        Objects.requireNonNull(value, "value cannot be null");

        return new String(Base64.getDecoder().decode(value), StandardCharsets.UTF_8);
    }
}