package com.atlassian.bitbucket.mesh.git;

import com.atlassian.bitbucket.mesh.AbstractGrpcServiceTest;
import com.atlassian.bitbucket.mesh.MockResponseObserver;
import com.atlassian.bitbucket.mesh.io.MoreFiles;
import com.atlassian.bitbucket.mesh.repository.Repository;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import com.atlassian.bitbucket.mesh.test.HasGitError;
import com.atlassian.bitbucket.mesh.test.annotation.Repo;
import com.google.common.base.Charsets;
import com.google.common.base.Splitter;
import com.google.protobuf.ByteString;
import io.grpc.Status;
import io.grpc.StatusRuntimeException;
import io.grpc.stub.StreamObserver;
import org.apache.commons.codec.binary.Hex;
import org.apache.commons.codec.digest.DigestUtils;
import org.apache.commons.lang3.RandomStringUtils;
import org.apache.commons.lang3.StringUtils;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.Timeout;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.MessageDigest;
import java.util.Iterator;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

public class GrpcLfsServiceTest extends AbstractGrpcServiceTest {

    private static final int DEFAULT_TEST_OBJECT_LENGTH = 16;
    private static final String OID_ZERO = StringUtils.repeat("0", 64);

    @Rule
    public final Timeout timeout = new Timeout(1, TimeUnit.MINUTES);

    @Repo
    private Repository repository;

    @Test
    public void testCanStore() {
        RpcLfsCanStoreResponse response = blockingStub().canStore(RpcLfsCanStoreRequest.newBuilder()
                .setRepository(repository.getId())
                .setSize(1)
                .build());
        assertTrue(response.getCanStore());
    }

    @Test
    public void testCanStoreReturnsFalseOnInsufficientSpace() {
        RpcLfsCanStoreResponse response = blockingStub().canStore(RpcLfsCanStoreRequest.newBuilder()
                .setRepository(repository.getId())
                .setSize(Long.MAX_VALUE / 2)
                .build());
        assertFalse(response.getCanStore());
    }

    @Test
    public void testObjectExistsNotPresent() {
        RpcLfsObjectExistsResponse response = blockingStub().objectExists(RpcLfsObjectExistsRequest.newBuilder()
                .setRepository(repository.getId())
                .setOid(OID_ZERO)
                .build());
        assertFalse(response.getExists());
    }

    @Test
    public void testObjectExistsPresent() throws IOException {
        String oid = createTestObject(DEFAULT_TEST_OBJECT_LENGTH);
        RpcLfsObjectExistsResponse response = blockingStub().objectExists(RpcLfsObjectExistsRequest.newBuilder()
                .setRepository(repository.getId())
                .setOid(oid)
                .build());
        assertTrue(response.getExists());
    }

    @Test
    public void testReadObject() throws IOException {
        int maxFragmentSize = 32 * 1024; // See LfsReadObjectHandler.IO_BLOCK_SIZE
        int size = maxFragmentSize + 1; // Two fragments

        String content = RandomStringUtils.randomAlphanumeric(size);
        String oid = createTestObject(content);

        Iterator<RpcStreamChunk> iterator = readObject(oid);
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();

        // First fragment
        assertTrue(iterator.hasNext());
        RpcStreamChunk chunk = iterator.next();
        outputStream.write(chunk.getData().toByteArray());
        assertFalse(chunk.getClosed());

        // Second fragment
        assertTrue(iterator.hasNext());
        chunk = iterator.next();
        outputStream.write(chunk.getData().toByteArray());
        assertTrue(chunk.getClosed());
        assertFalse(iterator.hasNext());

        // Validate combined message
        assertEquals(generateOid(outputStream.toByteArray()), oid);
    }

    @Test
    public void testReadObjectThrowsObjectNotFoundException() {
        Iterator<RpcStreamChunk> iterator = readObject(OID_ZERO);

        StatusRuntimeException thrown = assertThrows(StatusRuntimeException.class, iterator::next);
        assertSame(Status.Code.NOT_FOUND, thrown.getStatus().getCode());
        assertThat(thrown, HasGitError.noSuchObject(OID_ZERO));
    }

    @Test
    public void testReadObjectInvalidRangeEnd() throws IOException {
        Iterator<RpcStreamChunk> iterator = readObject(createTestObject(DEFAULT_TEST_OBJECT_LENGTH), 0, -1);

        StatusRuntimeException thrown = assertThrows(StatusRuntimeException.class, iterator::next);
        assertSame(Status.Code.OUT_OF_RANGE, thrown.getStatus().getCode());
        assertThat(thrown, HasGitError.lfsByteRangeNotSatisfiable(0, -1));
    }

    @Test
    public void testReadObjectInvalidRangeStart() throws IOException {
        Iterator<RpcStreamChunk> iterator = readObject(createTestObject(DEFAULT_TEST_OBJECT_LENGTH), -1, 10);

        StatusRuntimeException thrown = assertThrows(StatusRuntimeException.class, iterator::next);
        assertSame(Status.Code.OUT_OF_RANGE, thrown.getStatus().getCode());
        assertThat(thrown, HasGitError.lfsByteRangeNotSatisfiable(-1, 10));
    }

    @Test
    public void testReadObjectInvalidRangeStartBeyondEnd() throws IOException {
        Iterator<RpcStreamChunk> iterator = readObject(createTestObject(DEFAULT_TEST_OBJECT_LENGTH), 5, 0);

        StatusRuntimeException thrown = assertThrows(StatusRuntimeException.class, iterator::next);
        assertSame(Status.Code.OUT_OF_RANGE, thrown.getStatus().getCode());
        assertThat(thrown, HasGitError.lfsByteRangeNotSatisfiable(5, 0));
    }

    @Test
    public void testReadObjectRangeRequest() throws IOException {
        String oid = createTestObject("Hello World!");

        // Start substring
        Iterator<RpcStreamChunk> iterator = readObject(oid, 0, 2);
        RpcStreamChunk chunk = iterator.next();
        assertEquals("Hel", chunk.getData().toString(StandardCharsets.UTF_8.name()));
        assertTrue(chunk.getClosed());
        assertFalse(iterator.hasNext());

        // Mid substring
        iterator = readObject(oid, 3, 7);
        chunk = iterator.next();
        assertEquals("lo Wo", chunk.getData().toString(StandardCharsets.UTF_8.name()));
        assertTrue(chunk.getClosed());
        assertFalse(iterator.hasNext());

        // End substring
        iterator = readObject(oid, 8, 11);
        chunk = iterator.next();
        assertEquals("rld!", chunk.getData().toString(StandardCharsets.UTF_8.name()));
        assertTrue(chunk.getClosed());
        assertFalse(iterator.hasNext());
    }

    // IETF RFC 7233 says:
    // "If the last-byte-pos value is absent, or if the value is greater than or equal to the current length of the
    // representation data, the byte range is interpreted as the remainder of the representation"
    @Test
    public void testReadObjectRangeRequestEndBeyondFileExtent() throws IOException {
        String oid = createTestObject(DEFAULT_TEST_OBJECT_LENGTH);

        Iterator<RpcStreamChunk> iterator = readObject(oid, 0, DEFAULT_TEST_OBJECT_LENGTH * 2);
        RpcStreamChunk chunk = iterator.next();
        assertEquals(DEFAULT_TEST_OBJECT_LENGTH, chunk.getData().size());
        assertEquals(oid, generateOid(chunk.getData().toByteArray()));
        assertFalse(iterator.hasNext());
    }

    @Test
    public void testReadObjectZeroLength() throws IOException {
        String oid = createTestObject(0);
        Iterator<RpcStreamChunk> iterator = readObject(oid);
        assertEquals(0, iterator.next().getData().size());
    }

    @Test
    public void testWriteObject() throws InterruptedException, TimeoutException {
        String content = RandomStringUtils.randomAlphanumeric(DEFAULT_TEST_OBJECT_LENGTH);
        String oid = generateOid(content.getBytes(StandardCharsets.UTF_8));

        MockResponseObserver<RpcLfsWriteObjectResponse> responseObserver = writeObject(oid, content);
        responseObserver.await();

        assertEquals(1, responseObserver.getFragments().size());
        assertTrue(objectExistsOnDisk(oid));
    }

    @Test
    public void testWriteObjectFailsInReadOnlyMode() {
        try {
            repository.setReadOnly(true);

            String content = RandomStringUtils.randomAlphanumeric(DEFAULT_TEST_OBJECT_LENGTH);
            String oid = generateOid(content.getBytes(StandardCharsets.UTF_8));
            StatusRuntimeException thrown = assertThrows(StatusRuntimeException.class,
                    () -> writeObject(oid, content).await());

            RpcError error = RpcErrorUtils.fromThrowable(thrown);
            assertNotNull(error);
            assertTrue(error.hasRepositoryReadOnly());
            assertFalse(objectExistsOnDisk(oid));
        } finally {
            repository.setReadOnly(false);
        }
    }

    @Test
    public void testWriteObjectThatAlreadyExists() throws InterruptedException, TimeoutException, IOException {
        String content = RandomStringUtils.randomAlphanumeric(DEFAULT_TEST_OBJECT_LENGTH);
        String oid = createTestObject(content);

        MockResponseObserver<RpcLfsWriteObjectResponse> responseObserver = writeObject(oid, content);
        responseObserver.await();

        // Uploading an object that already exists should be transparent to the caller, it should look like a
        // normal upload.
        assertEquals(1, responseObserver.getFragments().size());
        assertTrue(objectExistsOnDisk(oid));
    }

    @Test
    public void testWriteObjectWithInvalidChecksum() {
        String content = RandomStringUtils.randomAlphanumeric(DEFAULT_TEST_OBJECT_LENGTH);
        String checksum = generateOid(content.getBytes(StandardCharsets.UTF_8));
        String bad_oid = OID_ZERO;

        StatusRuntimeException thrown = assertThrows(StatusRuntimeException.class,
                () -> writeObject(bad_oid, content).await());

        assertThat(thrown, HasGitError.lfsChecksumValidationFailed(bad_oid, checksum));
        // The object should not persisted as either the client provided OID or the calculated OID (i.e the checksum)
        assertFalse(objectExistsOnDisk(bad_oid));
        assertFalse(objectExistsOnDisk(checksum));
    }

    @Test
    public void testWriteObjectZeroLength() throws InterruptedException, TimeoutException {
        String content = "";
        String oid = generateOid(content.getBytes(StandardCharsets.UTF_8));

        MockResponseObserver<RpcLfsWriteObjectResponse> responseObserver = writeObject(oid, content);
        // Wait for request to complete
        responseObserver.await();

        assertEquals(1, responseObserver.getFragments().size());
        assertTrue(objectExistsOnDisk(oid));
    }

    private static RpcLfsByteRange.Builder toByteRange(long start, long end) {
        return RpcLfsByteRange.newBuilder()
                .setStart(start)
                .setEnd(end);
    }

    public LfsServiceGrpc.LfsServiceBlockingStub blockingStub() {
        return LfsServiceGrpc.newBlockingStub(managedChannel);
    }

    private String createTestObject(int length) throws IOException {
        String content = RandomStringUtils.randomAlphanumeric(length);
        return createTestObject(content);
    }

    private String createTestObject(String content) throws IOException {
        String oid = generateOid(content.getBytes(StandardCharsets.UTF_8));

        // Write the contents directly to the datastore filesystem (stored in location <sha256[0:1]>/sha256[2:63])
        Path dir = MoreFiles.mkdir(repository.getLfsDir().resolve(oid.substring(0, 2)));
        Path file = dir.resolve(oid.substring(2));
        Files.write(file, content.getBytes(StandardCharsets.UTF_8));

        return oid;
    }

    private String generateOid(byte[] content) {
        MessageDigest md = DigestUtils.getSha256Digest();
        md.update(content);
        return Hex.encodeHexString(md.digest());
    }

    private boolean objectExistsOnDisk(String oid) {
        return Files.isRegularFile(repository.getLfsDir().resolve(oid.substring(0, 2)).resolve(oid.substring(2)));
    }

    private Iterator<RpcStreamChunk> readObject(String oid) {
        return blockingStub().readObject(RpcLfsReadObjectRequest.newBuilder()
                .setRepository(repository.getId())
                .setOid(oid)
                .build());
    }

    private Iterator<RpcStreamChunk> readObject(String oid, long start, long end) {
        return blockingStub().readObject(RpcLfsReadObjectRequest.newBuilder()
                .setRepository(repository.getId())
                .setOid(oid)
                .setRange(toByteRange(start, end))
                .build());
    }

    private LfsServiceGrpc.LfsServiceStub stub() {
        return LfsServiceGrpc.newStub(managedChannel);
    }

    private MockResponseObserver<RpcLfsWriteObjectResponse> writeObject(String oid, String content) {
        MockResponseObserver<RpcLfsWriteObjectResponse> responseObserver = new MockResponseObserver<>();
        StreamObserver<RpcLfsWriteObjectFragment> observer = stub().writeObject(responseObserver);

        // Send request (OID and repository)
        observer.onNext(RpcLfsWriteObjectFragment.newBuilder()
                .setRequest(RpcLfsWriteObjectRequest.newBuilder()
                        .setRepository(repository.getId())
                        .setOid(oid))
                .build());

        // Split content to ensure testing with multiple fragments
        Iterator<String> contentFragments = Splitter.fixedLength(7).split(content).iterator();
        while (contentFragments.hasNext()) {
            observer.onNext(RpcLfsWriteObjectFragment.newBuilder()
                    .setContent(RpcStreamChunk.newBuilder()
                            .setData(ByteString.copyFrom(contentFragments.next(), Charsets.UTF_8))
                            .setClosed(!contentFragments.hasNext()))
                    .build());
        }

        observer.onCompleted();
        return responseObserver;
    }
}
