package de.dhbwstuttgart.server;

import com.fasterxml.jackson.core.JsonProcessingException;
import de.dhbwstuttgart.server.packet.ErrorPacket;
import de.dhbwstuttgart.server.packet.IClientToServerPacket;
import de.dhbwstuttgart.server.packet.IPacket;
import de.dhbwstuttgart.server.packet.MessagePacket;
import de.dhbwstuttgart.server.packet.PacketContainer;
import de.dhbwstuttgart.util.Logger;
import java.net.InetSocketAddress;
import java.util.Collections;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.SynchronousQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import org.java_websocket.WebSocket;
import org.java_websocket.drafts.Draft;
import org.java_websocket.drafts.Draft_6455;
import org.java_websocket.extensions.permessage_deflate.PerMessageDeflateExtension;
import org.java_websocket.handshake.ClientHandshake;
import org.java_websocket.server.WebSocketServer;

public class SocketServer extends WebSocketServer {

  public static Logger logger = new Logger("SocketServer");
  public static final int maxTasksPerSession = 100;
  private static boolean serverRunning = false;

  /**
   * Increase this every time a breaking change to the server communication is done.
   * This will prevent errors when the server version and client version do not match.
   */
  public static final String packetProtocolVersion = "1";

  // create an executor for tasks that will always keep at least one task around
  private final ThreadPoolExecutor taskExecutor = new ThreadPoolExecutor(1, Integer.MAX_VALUE,60L, TimeUnit.SECONDS, new SynchronousQueue<>());
  // create an executor for scheduling timeouts
  private final ScheduledExecutorService timeoutExecutor = Executors.newSingleThreadScheduledExecutor();

  public SocketServer(int port) {
    super(new InetSocketAddress(port));
    this.setConnectionLostTimeout(30);

    serverRunning = true;
    // add a shutdown hook to close all connections when the process ends or is stopped by a SIGINT signal
    Runtime.getRuntime().addShutdownHook(new Thread(this::onShutdown));
  }

  public static boolean isServerRunning() {
    return serverRunning;
  }

  private void onShutdown() {
    serverRunning = false;
    try {
      for (var webSocket : this.getConnections()) {
        this.sendError(webSocket, "Sorry, i am shutting down. You are now on your own, good Luck!", true);
        webSocket.close();
      }
      this.stop();
      taskExecutor.shutdown();
      timeoutExecutor.shutdown();
    } catch (InterruptedException exception) {
      // we are shutting down anyway
    }
  }

  @Override
  public void onOpen(WebSocket webSocket, ClientHandshake clientHandshake) {
    String ppv = clientHandshake.getFieldValue("packetProtocolVersion");
    if (!ppv.equals(packetProtocolVersion)) {
      this.sendError(webSocket,
        "Mismatch in packet protocol version! Client (you): \"" + ppv + "\" and Server (me): \"" + packetProtocolVersion + "\"",
        true
      );
      webSocket.close(1);
      return;
    }

    SocketData socketData = new SocketData(webSocket);
    logger.info("New connection: " + socketData.id + " (with ppv " + ppv + ")");

    try {
      sendMessage(webSocket, "Welcome to the server!");

      // wait 10 seconds for the client to send a task and close the connection if nothing has been received until then
      final int secondsUntilTimeout = 10;
      timeoutExecutor.schedule(() -> {
          if (webSocket.<SocketData>getAttachment().totalTasks.get() > 0 || !webSocket.isOpen()) {
            return;
          }
          sendMessage(webSocket, "No task received after " + secondsUntilTimeout + " seconds. Closing connection...");
          webSocket.close();
        },
        secondsUntilTimeout,
        TimeUnit.SECONDS
      );

      // and finally, when your program wants to exit
    } catch (Exception e) {
      logger.exception(e);
      webSocket.close(1, e.getMessage());
    }
  }

  @Override
  public void onClose(WebSocket webSocket, int code, String reason, boolean remote) {
    SocketData socketData = webSocket.getAttachment();
    logger.info("Connection closed: " + socketData.id);
    logger.info(
      "Disconnected client " + socketData.id + " " +
        "with code " + code + " " +
        (reason.isEmpty() ? "" : "and reason " + reason + " ") +
        "(closed by client: " + remote + ")"
    );

  }

  @Override
  public void onMessage(WebSocket webSocket, String s) {
    // logger.info("Received: " + s.substring(0, 50));
    IPacket reconstructedPacket = PacketContainer.deserialize(s);
    try {
      this.onPacketReceived(webSocket, reconstructedPacket);
    } catch (JsonProcessingException e) {
      logger.exception(e);
      this.log(webSocket, "Error on processing incoming package: " + e.getMessage());
    }
  }

  @Override
  public void onError(WebSocket webSocket, Exception e) {
    if (webSocket != null) {
      log(webSocket, e.getMessage());
      webSocket.close();
    }
    logger.exception(e);
  }

  @Override
  public void onStart() {
    logger.success("Websocket server started on port " + this.getPort());
  }

  /**
   * A shorthand method for sending informational messages to the client
   */
  public void sendMessage(WebSocket webSocket, String text) {
    try {
      MessagePacket message = MessagePacket.create(text);
      webSocket.send(PacketContainer.serialize(message));
    } catch (Exception e) {
      System.err.println("Failed to send message: " + text);
      logger.exception(e);
    }
  }

  /**
   * A shorthand method for sending error messages to the client
   */
  public void sendError(WebSocket webSocket, String text, boolean isFatal) {
    try {
      ErrorPacket error = ErrorPacket.create(text, isFatal);
      webSocket.send(PacketContainer.serialize(error));
    } catch (Exception e) {
      logger.exception(e);
      log(webSocket, "Failed to send error: " + text);
    }
  }

  /**
   * The server-side implementation on how to handle certain packets when received
   */
  private void onPacketReceived(WebSocket webSocket, IPacket packet) throws JsonProcessingException {
    SocketData socketData = webSocket.getAttachment();

    // limit the number of tasks per connection
    if (socketData.totalTasks.get() >= maxTasksPerSession) {
      sendError(webSocket, "Exceeded the maximum amount of " + maxTasksPerSession + " tasks per session", true);
      webSocket.close();
      return;
    }

    // only allow packets that are meant to be handled by the server
    if (!(packet instanceof IClientToServerPacket<?> clientToServerPacket)) {
      sendMessage(webSocket, "The package of type " + packet.getClass().getName() + " is not handled by the server!");
      return;
    }

    // update the socket data
    socketData.unhandledTasks.incrementAndGet();
    socketData.totalTasks.incrementAndGet();

    // add the packet to the queue so it can be started by the worker
    CompletableFuture.runAsync(() -> {
      clientToServerPacket.onHandle(webSocket, this);
      int remainingUnhandledTasks = socketData.unhandledTasks.decrementAndGet();

      if (socketData.closeIfNoTasksLeft) {
        // if the websocket has 0 unhandled Tasks, close the connection
        if (remainingUnhandledTasks <= 0) {
          sendMessage(webSocket, "All requested tasks finished! Closing connection...");
          webSocket.close();
        }
      }
    }, taskExecutor);
  }

  public void log(WebSocket webSocket, String msg) {
    String socketId = (webSocket == null) ? "???" : webSocket.<SocketData>getAttachment().id;
    logger.info("[" + socketId + "] " + msg);
  }

  /**
   * The data that is associated server-side with any connected client.
   * This makes it possible to store information that can be mapped to any existing connection.
   */
  public static class SocketData {

    public final String id;
    public final AtomicInteger unhandledTasks = new AtomicInteger(0);
    public final AtomicInteger totalTasks = new AtomicInteger(0);
    public boolean closeIfNoTasksLeft = false;

    public SocketData(WebSocket webSocket) {
      this.id = UUID.randomUUID().toString();
      webSocket.setAttachment(this);
    }
  }
}
