package com.atlassian.stash.internal.hook;

import com.atlassian.bitbucket.hook.*;
import com.atlassian.bitbucket.util.IoUtils;
import com.atlassian.bitbucket.repository.Repository;
import com.atlassian.bitbucket.repository.RepositorySupplier;
import com.atlassian.bitbucket.scm.ScmService;
import com.atlassian.bitbucket.server.ApplicationPropertiesService;
import com.atlassian.bitbucket.util.FileUtils;
import com.atlassian.bitbucket.util.Timer;
import com.atlassian.bitbucket.util.TimerUtils;
import com.atlassian.plugin.spring.AvailableToPlugins;
import com.atlassian.sal.api.executor.ThreadLocalContextManager;
import com.atlassian.security.random.SecureTokenGenerator;
import com.atlassian.stash.internal.ApplicationConstants;
import com.atlassian.stash.internal.scm.InternalScmService;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Strings;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.io.Files;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.ClassPathResource;
import org.springframework.stereotype.Component;

import javax.annotation.Nonnull;
import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.net.InetAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.nio.BufferOverflowException;
import java.util.List;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutorService;
import java.util.function.Function;

import static com.google.common.base.Preconditions.checkNotNull;

@Component("hookService")
@AvailableToPlugins(HookService.class)
public class DefaultHookService implements HookService {

    private static final String CALLBACK_SCRIPT = "hook-callback.pl";
    private static final String COORDINATOR_SCRIPT = "hook-coordinator.sh";
    /**
     * Uses the default backlog for accepting new connections. (Currently 50, but may vary by JDK)
     */
    private static final int DEFAULT_BACKLOG = 0;

    private static final Logger log = LoggerFactory.getLogger(DefaultHookService.class);

    private final ExecutorService executor;
    private final File hookScriptDirectory;
    private final RepositorySupplier repositorySupplier;
    private final ScmService scmService;
    private final SecureTokenGenerator tokenGenerator;
    private final ConcurrentMap<String, HookState> stateForScmRequest = Maps.newConcurrentMap();
    private final ThreadLocalContextManager<Object> stateManager;

    private final long hookBufferCapacity;
    private final String hookAddress;
    private final int hookPort;

    private ServerSocket serverSocket;
    private volatile boolean shutdown = false;

    @Autowired
    public DefaultHookService(ApplicationPropertiesService applicationPropertiesService, ExecutorService executor,
                              RepositorySupplier repositorySupplier, InternalScmService scmService,
                              ThreadLocalContextManager<Object> stateManager, SecureTokenGenerator tokenGenerator,
                              @Value(ApplicationConstants.PROP_HOOK_CALLBACK_BUFFER_CAPACITY) long hookBufferCapacity,
                              @Value(ApplicationConstants.PROP_HOOK_CALLBACK_SOCKET_ADDRESS) String hookAddress,
                              @Value(ApplicationConstants.PROP_HOOK_CALLBACK_SOCKET_PORT) int hookPort) throws IOException {
        this.executor = executor;
        this.repositorySupplier = repositorySupplier;
        this.scmService = scmService;
        this.stateManager = stateManager;
        this.tokenGenerator = tokenGenerator;

        this.hookBufferCapacity = Math.max(hookBufferCapacity, 32768); //At least a 32k buffer
        this.hookAddress = Strings.emptyToNull(hookAddress);
        this.hookPort = Math.max(hookPort, 0); //Don't allow negative ports

        hookScriptDirectory = new File(applicationPropertiesService.getBinDir(), "git-hooks");
    }

    @Override
    public <T> T doWithHookRequest(int repositoryId, @Nonnull Function<HookRequestHandle, T> callback) {
        checkNotNull(callback, "callback");
        try (HookRequestHandle handle = registerRequest(repositoryId)) {
            return callback.apply(handle);
        }
    }

    @Nonnull
    @Override
    public HookRequestHandle registerRequest(int repositoryId) {
        String scmRequestId = tokenGenerator.generateToken();
        DefaultRequestHandle handle = new DefaultRequestHandle(scmRequestId);
        stateForScmRequest.put(scmRequestId, new HookState(handle, repositoryId, stateManager.getThreadLocalContext()));
        return handle;
    }

    @PostConstruct
    public void startup() throws IOException {
        installHookScripts();
        startHookCallbackListener();
    }

    @PreDestroy
    public void shutdown() {
        shutdown = true;
        try {
            if (serverSocket != null) {
                serverSocket.close();
            }
        } catch (IOException e) {
            log.warn("Could not close socket for hook callbacks", e);
        }
    }

    private void installHookScript(String scriptName, File targetDir) throws IOException {
        File scriptFile = new File(targetDir, scriptName);

        ClassPathResource script = new ClassPathResource("/hooks/" + scriptName, DefaultHookService.class);
        try (InputStream scriptStream = script.getInputStream()) {
            IoUtils.copy(scriptStream, scriptFile);
        }

        if (!scriptFile.setExecutable(true)) {
            throw new IOException(scriptFile.getAbsolutePath() + " could not be set executable.");
        }
    }

    private void installHookScripts() throws IOException {
        if (hookScriptDirectory.exists() && !hookScriptDirectory.isDirectory()) {
            log.warn("{} must be a directory. Attempting to move the file", hookScriptDirectory.getAbsolutePath());
            try {
                File moved = new File(hookScriptDirectory.getParent(), hookScriptDirectory.getName() + ".moved");
                Files.move(hookScriptDirectory, moved);
                log.warn("The file at {} has been renamed to {}",
                        hookScriptDirectory.getAbsolutePath(), moved.getName());
            } catch (IOException e) {
                throw new IOException(hookScriptDirectory.getAbsolutePath() + " exists and is not a directory, " +
                        "preventing installation of required hook support scripts.", e);
            }
        }

        try {
            FileUtils.mkdir(hookScriptDirectory);
        } catch (IllegalStateException e) {
            throw new IOException(hookScriptDirectory.getAbsolutePath() +
                    " could not be created. Hook support scripts cannot be installed.", e);
        }

        // ensure that the latest version of the hook scripts is installed
        try {
            installHookScript(CALLBACK_SCRIPT, hookScriptDirectory);
            installHookScript(COORDINATOR_SCRIPT, hookScriptDirectory);
        } catch (IOException e) {
            throw new IOException("Hook support scripts could not be written to " +
                    hookScriptDirectory.getAbsolutePath(), e);
        }
    }

    private void startHookCallbackListener() throws IOException {
        //hookPort may be 0, in which case an open port will automatically be chosen and used.
        //hookAddress may be null, in which case the _loopback_ (localhost) is used. This is generally the
        //best configuration as it ensures the callback socket is only available to hooks on-box.
        serverSocket = new ServerSocket(hookPort, DEFAULT_BACKLOG, InetAddress.getByName(hookAddress));

        Thread socketListenerThread = new Thread("hook-callback-listener") {

            @Override
            public void run() {
                while (!shutdown) {
                    final Socket socket;
                    try {
                        socket = serverSocket.accept();
                    } catch (Exception e) {
                        if (!shutdown) {
                            log.warn("A hook connection could not be established; accept failed", e);
                        }
                        continue;
                    }

                    executor.submit(new Runnable() {

                        @Override
                        public void run() {
                            try (Socket ignored = socket; //To close the socket automatically
                                 SocketTransferInput input = new SocketTransferInput(socket.getInputStream());
                                 SocketTransferOutput output = new SocketTransferOutput(socket.getOutputStream())
                            ) {
                                handleRawRequest(input, output);
                            } catch (Exception e) {
                                if (!shutdown) {
                                    log.warn("Hook socket I/O failed before the repository/hook could be identified", e);
                                }
                            }
                        }
                    });
                }
            }
        };
        socketListenerThread.setDaemon(true);
        socketListenerThread.start();

        log.info("Hook callback socket listening on {}:{}", serverSocket.getInetAddress().getHostAddress(), serverSocket.getLocalPort());
    }

    @VisibleForTesting
    void handleRawRequest(SocketTransferInput input, SocketTransferOutput output) throws Exception {
        HookState hookState = null;
        DefaultHookRequest request = null;
        DefaultHookResponse response = new DefaultHookResponse();
        Timer timer = null;
        try {
            // Create request and fetch hook state
            hookState = getHookState(input);
            boolean acceptRequest = false;
            try {
                request = createHookRequest(input);

                timer = TimerUtils.start("DefaultHookService hook callback " + request.getHookName() +
                        " for repository " + hookState.getRepositoryId());

                // delegate out to the handler for this callback
                acceptRequest = handleRequest(request, response, hookState);
            } catch (BufferOverflowException e) {
                response.err().println("This push is too large to process.");
            } finally {
                DefaultRequestHandle handle = hookState.getHandle();
                handle.setAccepted(acceptRequest);
                handle.setCalled(true);
            }

            // write response to output
            writeResponse(output, response, acceptRequest);

        } catch (Exception e) {
            if (!shutdown) {
                // Handle it if we can to provide more information
                if (hookState != null) {
                    if (request != null) {
                        log.error(String.format("An exception occurred communicating with '%s' hook in repository %s", request.getHookName(), hookState.getRepositoryId()), e);
                    } else {
                        log.error(String.format("An exception occurred communicating with an unknown hook in repository %s", hookState.getRepositoryId()), e);
                    }
                } else {
                    throw e;
                }
            }
        } finally {
            if (timer != null) {
                timer.close();
            }
        }
    }

    @VisibleForTesting
    HookState getHookState(SocketTransferInput input) throws IOException {
        // first part = request id
        String scmRequestId = input.readRequestId();
        log.trace("Received hook callback with request id {}", scmRequestId);

        return checkNotNull(stateForScmRequest.get(scmRequestId), "hookState");
    }

    @VisibleForTesting
    DefaultHookRequest createHookRequest(SocketTransferInput input) throws IOException {
        // second part = hook-identifier
        String hookIdentifier = input.readHookType();
        log.trace("Received callback for hookType {}", hookIdentifier);

        // remainder = hook arguments and input
        final List<String> arguments = Lists.newArrayList();
        StringBuilder stdIn = new StringBuilder();
        SocketTransferInput.Chunk chunk;
        long bufferSize = 0;
        while ((chunk = input.readChunk()).getType() != SocketTransferInput.ChunkType.END) {
            if (hookBufferCapacity > 0 && bufferSize > hookBufferCapacity) {
                throw new BufferOverflowException();
            }
            switch (chunk.getType()) {
                case ARG:
                    arguments.add(chunk.getValue());
                    break;
                case STDIN:
                    stdIn.append(chunk.getValue());
                    break;
                default:
                    throw new IllegalStateException("Invalid input from hook request");
            }
            bufferSize += ((long) chunk.getValue().length()) * 2; // chars -> bytes
        }

        return new DefaultHookRequest(hookIdentifier, arguments, stdIn);
    }

    @VisibleForTesting
    void writeResponse(SocketTransferOutput output, DefaultHookResponse response, boolean acceptRequest) throws IOException {
        String out = response.getOutput();
        if (out.length() > 0) {
            output.writeStdOut(out);
        }

        String err = response.getError();
        if (err.length() > 0) {
            output.writeStdErr(err);
        }

        output.writeExitCode((short) (acceptRequest ? 0 : 1));
        output.flush();
    }

    @VisibleForTesting
    boolean handleRequest(HookRequest request, HookResponse response, HookState hookState) throws IOException {
        stateManager.setThreadLocalContext(hookState.getThreadLocalState());
        try {
            return doHandleRequest(request, response, hookState);
        } finally {
            stateManager.clearThreadLocalContext();
        }
    }

    @VisibleForTesting
    boolean doHandleRequest(HookRequest request, HookResponse response, HookState hookState) throws IOException {
        Repository repository = repositorySupplier.getById(hookState.getRepositoryId());
        if (repository == null) {
            log.warn("Failed to find repository with id {} for hook callback", hookState.getRepositoryId());
            return false;
        }
        HookHandler handler = scmService.getHookHandlerFactory(repository).create(request);
        return handler == null || handler.handle(request, response);
    }

    @VisibleForTesting
    static class HookState {
        private final DefaultRequestHandle handle;
        private final int repositoryId;
        private final Object threadLocalState;

        private HookState(DefaultRequestHandle handle, int repositoryId, Object threadLocalState) {
            this.handle = handle;
            this.repositoryId = repositoryId;
            this.threadLocalState = threadLocalState;
        }

        public DefaultRequestHandle getHandle() {
            return handle;
        }

        public int getRepositoryId() {
            return repositoryId;
        }

        public Object getThreadLocalState() {
            return threadLocalState;
        }
    }

    private class DefaultRequestHandle implements HookRequestHandle {

        private final String requestId;

        private volatile boolean accepted;
        private volatile boolean called;

        private DefaultRequestHandle(String requestId) {
            this.requestId = requestId;
        }

        @Override
        public void close() {
            stateForScmRequest.remove(requestId);
        }

        @Nonnull
        @Override
        public File getCallbackScript() {
            return new File(hookScriptDirectory, CALLBACK_SCRIPT);
        }

        @Nonnull
        @Override
        public File getCoordinatorScript() {
            return new File(hookScriptDirectory, COORDINATOR_SCRIPT);
        }

        @Nonnull
        @Override
        public String getHostAddress() {
            return serverSocket.getInetAddress().getHostAddress();
        }

        @Override
        public int getPort() {
            return serverSocket.getLocalPort();
        }

        @Nonnull
        @Override
        public String getRequestId() {
            return requestId;
        }

        @Override
        public boolean isAccepted() {
            return accepted;
        }

        @Override
        public boolean isCalled() {
            return called;
        }

        public void setAccepted(boolean accepted) {
            this.accepted = accepted;
        }

        public void setCalled(boolean called) {
            this.called = called;
        }
    }
}
