/*
 * Copyright 2017 ThoughtWorks, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.thoughtworks.go.agent;

import com.thoughtworks.go.agent.service.AgentUpgradeService;
import com.thoughtworks.go.agent.service.SslInfrastructureService;
import com.thoughtworks.go.buildsession.ArtifactsRepository;
import com.thoughtworks.go.buildsession.BuildSession;
import com.thoughtworks.go.buildsession.BuildVariables;
import com.thoughtworks.go.config.AgentRegistry;
import com.thoughtworks.go.domain.BuildSettings;
import com.thoughtworks.go.plugin.access.packagematerial.PackageRepositoryExtension;
import com.thoughtworks.go.plugin.access.pluggabletask.TaskExtension;
import com.thoughtworks.go.plugin.access.scm.SCMExtension;
import com.thoughtworks.go.plugin.infra.PluginManager;
import com.thoughtworks.go.publishers.GoArtifactsManipulator;
import com.thoughtworks.go.remote.AgentIdentifier;
import com.thoughtworks.go.remote.AgentInstruction;
import com.thoughtworks.go.remote.BuildRepositoryRemote;
import com.thoughtworks.go.remote.work.ConsoleOutputTransmitter;
import com.thoughtworks.go.remote.work.RemoteConsoleAppender;
import com.thoughtworks.go.remote.work.Work;
import com.thoughtworks.go.server.service.AgentBuildingInfo;
import com.thoughtworks.go.util.*;
import com.thoughtworks.go.util.command.TaggedStreamConsumer;
import com.thoughtworks.go.websocket.Action;
import com.thoughtworks.go.websocket.Message;
import com.thoughtworks.go.websocket.MessageEncoding;
import org.eclipse.jetty.websocket.api.Session;
import org.eclipse.jetty.websocket.api.annotations.*;
import org.eclipse.jetty.websocket.api.extensions.Frame;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.InputStream;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

@WebSocket
public class AgentWebSocketClientController extends AgentController {
    private static final Logger LOG = LoggerFactory.getLogger(AgentWebSocketClientController.class);
    private final SslInfrastructureService sslInfrastructureService;
    private final GoArtifactsManipulator manipulator;
    private HttpService httpService;
    private WebSocketClientHandler webSocketClientHandler;
    private WebSocketSessionHandler webSocketSessionHandler;
    private AtomicReference<BuildSession> buildSession = new AtomicReference<>();
    private JobRunner runner;
    private PackageRepositoryExtension packageRepositoryExtension;
    private SCMExtension scmExtension;
    private TaskExtension taskExtension;
    private BuildRepositoryRemote server;

    public AgentWebSocketClientController(BuildRepositoryRemote server, GoArtifactsManipulator manipulator,
                                          SslInfrastructureService sslInfrastructureService, AgentRegistry agentRegistry,
                                          AgentUpgradeService agentUpgradeService, SubprocessLogger subprocessLogger,
                                          SystemEnvironment systemEnvironment, PluginManager pluginManager,
                                          PackageRepositoryExtension packageRepositoryExtension, SCMExtension scmExtension,
                                          TaskExtension taskExtension, HttpService httpService,
                                          WebSocketClientHandler webSocketClientHandler, WebSocketSessionHandler webSocketSessionHandler) {
        super(sslInfrastructureService, systemEnvironment, agentRegistry, pluginManager, subprocessLogger, agentUpgradeService);
        this.server = server;
        this.manipulator = manipulator;
        this.packageRepositoryExtension = packageRepositoryExtension;
        this.scmExtension = scmExtension;
        this.taskExtension = taskExtension;
        this.sslInfrastructureService = sslInfrastructureService;
        this.httpService = httpService;
        this.webSocketClientHandler = webSocketClientHandler;
        this.webSocketSessionHandler = webSocketSessionHandler;
    }

    @Override
    public void ping() {
        // Do nothing
    }

    @Override
    public void execute() {
        // Do nothing
    }

    @Override
    public void work() throws Exception {
        if (sslInfrastructureService.isRegistered()) {
            if (webSocketSessionHandler.isNotRunning()) {
                webSocketSessionHandler.clearCallBacks();
                webSocketSessionHandler.setSession(webSocketClientHandler.connect(this));
            }
            updateServerAgentRuntimeInfo();
        }
    }


    void process(Message message) throws InterruptedException {
        switch (message.getAction()) {
            case cancelBuild:
                cancelJobIfThereIsOneRunning();
                cancelBuild();
                break;
            case setCookie:
                String cookie = MessageEncoding.decodeData(message.getData(), String.class);
                getAgentRuntimeInfo().setCookie(cookie);
                LOG.info("Got cookie: {}", cookie);
                break;
            case assignWork:
                cancelJobIfThereIsOneRunning();
                Work work = MessageEncoding.decodeWork(message.getData());
                LOG.debug("Got work from server: [{}]", work.description());
                getAgentRuntimeInfo().idle();
                runner = new JobRunner();
                try {
                    runner.run(work, agentIdentifier(),
                            new BuildRepositoryRemoteAdapter(runner, webSocketSessionHandler),
                            manipulator, getAgentRuntimeInfo(),
                            packageRepositoryExtension, scmExtension,
                            taskExtension);
                } finally {
                    getAgentRuntimeInfo().idle();
                    updateServerAgentRuntimeInfo();
                }
                break;
            case build:
                cancelBuild();
                BuildSettings buildSettings = MessageEncoding.decodeData(message.getData(), BuildSettings.class);
                runBuild(buildSettings);
                break;
            case reregister:
                LOG.warn("Reregister: invalidate current agent certificate fingerprint {} and stop websocket webSocketClient.", getAgentRegistry().uuid());
                webSocketSessionHandler.stop();
                sslInfrastructureService.invalidateAgentCertificate();
                break;
            case acknowledge:
                webSocketSessionHandler.acknowledge(message);
                break;
            default:
                throw new RuntimeException("Unknown action: " + message.getAction());

        }
    }

    private void runBuild(BuildSettings buildSettings) {
        URLService urlService = new URLService();
        TaggedStreamConsumer buildConsole;

        if (getSystemEnvironment().isConsoleLogsThroughWebsocketEnabled()) {
            buildConsole = new ConsoleOutputWebsocketTransmitter(webSocketSessionHandler, buildSettings.getBuildId());
        } else {
            buildConsole = new ConsoleOutputTransmitter(
                new RemoteConsoleAppender(
                    urlService.prefixPartialUrl(buildSettings.getConsoleUrl()),
                    httpService)
            );
        }

        ArtifactsRepository artifactsRepository = new UrlBasedArtifactsRepository(
                httpService,
                urlService.prefixPartialUrl(buildSettings.getArtifactUploadBaseUrl()),
                urlService.prefixPartialUrl(buildSettings.getPropertyBaseUrl()),
                new ZipUtil());

        DefaultBuildStateReporter buildStateReporter = new DefaultBuildStateReporter(webSocketSessionHandler, getAgentRuntimeInfo());

        TimeProvider clock = new TimeProvider();
        BuildVariables buildVariables = new BuildVariables(getAgentRuntimeInfo(), clock);
        BuildSession build = new BuildSession(
                buildSettings.getBuildId(),
                buildStateReporter,
                buildConsole,
                buildVariables,
                artifactsRepository,
                httpService, clock, new File("."));

        this.buildSession.set(build);

        build.setEnv("GO_SERVER_URL", getSystemEnvironment().getServiceUrl());
        getAgentRuntimeInfo().idle();
        try {
            getAgentRuntimeInfo().busy(new AgentBuildingInfo(buildSettings.getBuildLocatorForDisplay(), buildSettings.getBuildLocator()));
            build.build(buildSettings.getBuildCommand());
        } finally {
            try {
                buildConsole.stop();
            } finally {
                getAgentRuntimeInfo().idle();
            }
        }
        this.buildSession.set(null);
    }

    private void cancelBuild() throws InterruptedException {
        BuildSession build = this.buildSession.get();
        if (build == null) {
            return;
        }
        getAgentRuntimeInfo().cancel();
        if (!build.cancel(30, TimeUnit.SECONDS)) {
            LOG.error("Waited 30 seconds for canceling job finish, but the job is still running. Maybe canceling job does not work as expected, here is buildSession details: " + buildSession.get());
        }
    }

    private void cancelJobIfThereIsOneRunning() throws InterruptedException {
        if (runner == null || !runner.isRunning()) {
            return;
        }
        LOG.info("Cancel running job");
        runner.handleInstruction(new AgentInstruction(true), getAgentRuntimeInfo());
        runner.waitUntilDone(30);
        if (runner.isRunning()) {
            LOG.error("Waited 30 seconds for canceling job finish, but the job is still running. Maybe canceling job does not work as expected, here is running job details: " + runner);
        }
    }

    private void updateServerAgentRuntimeInfo() {
        AgentIdentifier agent = agentIdentifier();
        LOG.trace("{} is pinging server [{}]", agent, server);
        getAgentRuntimeInfo().refreshUsableSpace();
        webSocketSessionHandler.sendAndWaitForAcknowledgement(new Message(Action.ping, MessageEncoding.encodeData(getAgentRuntimeInfo())));
        LOG.trace("{} pinged server [{}]", agent, server);
    }

    private Executor executor = Executors.newFixedThreadPool(5);

    @OnWebSocketConnect
    public void onConnect(Session session) {
        LOG.info(session + " connected.");
    }

    @OnWebSocketMessage
    public void onMessage(InputStream raw) {
        final Message msg = MessageEncoding.decodeMessage(raw);
        LOG.debug("{} message: {}", webSocketSessionHandler.getSessionName(), msg);
        executor.execute(new Runnable() {
            @Override
            public void run() {
                try {
                    LOG.debug("Processing message[" + msg + "].");
                    process(msg);
                } catch (InterruptedException e) {
                    LOG.error("Process message[" + msg + "] is interruptted.", e);
                } catch (RuntimeException e) {
                    LOG.error("Unexpected error while processing message[" + msg + "]: " + e.getMessage(), e);
                } finally {
                    LOG.debug("Finished trying to process message[" + msg + "].");
                }
            }
        });
    }

    @OnWebSocketClose
    public void onClose(int closeCode, String closeReason) {
        LOG.debug("{} closed. code: {}, reason: {}", webSocketSessionHandler.getSessionName(), closeCode, closeReason);
    }

    @OnWebSocketError
    public void onError(Throwable error) {
        LOG.error(webSocketSessionHandler.getSessionName() + " error", error);
    }

    @OnWebSocketFrame
    public void onFrame(Frame frame) {
        LOG.debug("{} receive frame: {}", webSocketSessionHandler.getSessionName(), frame.getPayloadLength());
    }

}
