package ru.dragonestia.kubik.client.proxy;

import com.google.gson.Gson;
import okhttp3.*;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import ru.dragonestia.kubik.ServerEntry;
import ru.dragonestia.kubik.client.Kubik;
import ru.dragonestia.kubik.dto.KubikEventUpdate;

import java.util.HashSet;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.logging.Logger;

public class ProxyControl {

    private static final Gson gson = new Gson();

    private final Kubik kubik;
    private final ScheduledExecutorService executor;
    private final ProxyServerRegistry registry;

    private final Object lock = new Object();
    private final Map<String, ServerEntry> knownServers = new ConcurrentHashMap<>();
    private final Logger logger = Logger.getLogger(ProxyControl.class.getName());

    private WebSocket ws;
    private ScheduledFuture<?> future;

    public ProxyControl(Kubik kubik, ScheduledExecutorService executor, ProxyServerRegistry registry) {
        this.kubik = kubik;
        this.executor = executor;
        this.registry = registry;

        createConnection();
    }

    private synchronized void handle(KubikEventUpdate update) {
        for (var toDelete: update.unregistered()) {
            var server = knownServers.remove(toDelete);
            if (server != null) registry.unregister(toDelete);
        }

        if (update.initial()) {
            var notAdded = new HashSet<>(knownServers.keySet());
            for (var toAdd: update.registered()) {
                var newServer = new ServerEntry(toAdd.id(), toAdd.address(), toAdd.port());
                var prevServer = knownServers.put(toAdd.id(), newServer);
                if (!newServer.equals(prevServer)) registry.register(toAdd.id(), toAdd.address(), toAdd.port());
                notAdded.remove(toAdd.id());
            }
            for (var toDelete: notAdded) {
                knownServers.remove(toDelete);
                registry.unregister(toDelete);
            }
            return;
        }

        for (var toAdd: update.registered()) {
            var newServer = new ServerEntry(toAdd.id(), toAdd.address(), toAdd.port());
            var prevServer = knownServers.put(toAdd.id(), newServer);
            if (!newServer.equals(prevServer)) registry.register(toAdd.id(), toAdd.address(), toAdd.port());
        }
    }

    public void dispose() {
        synchronized (lock) {
            if (future != null) {
                future.cancel(false);
                future = null;
            }
            if (ws != null) {
                ws.cancel();
                ws = null;
            }
        }
    }

    private void createConnection() {
        synchronized (lock) {
            if (future != null) future = null;

            var request = new Request.Builder()
                    .url(kubik.getWsUrl() + "/api/v1/server/updates/ws")
                    .build();

            ws = kubik.getOkHttpClient().newWebSocket(request, new WebSocketListener() {

                @Override
                public void onOpen(@NotNull WebSocket webSocket, @NotNull Response response) {
                    //System.out.println("Connected to discovery updates: " + kubik.getWsUrl());
                }

                @Override
                public void onMessage(@NotNull WebSocket webSocket, @NotNull String text) {
                    handle(gson.fromJson(text, KubikEventUpdate.class));
                }

                @Override
                public void onFailure(@NotNull WebSocket webSocket, @NotNull Throwable t, @Nullable Response response) {
                    synchronized (lock) {
                        ws.cancel();
                        ws = null;

                        future = executor.schedule(() -> createConnection(), kubik.getConfig().getReconnectDelay(), TimeUnit.SECONDS);
                    }
                }

                @Override
                public void onClosing(@NotNull WebSocket webSocket, int code, @NotNull String reason) {
                    synchronized (lock) {
                        ws.cancel();
                        ws = null;

                        future = executor.schedule(() -> createConnection(), kubik.getConfig().getReconnectDelay(), TimeUnit.SECONDS);
                    }
                }
            });
        }
    }
}
