/**
 * Copyright (c) 2025 Dominic Masters
 * 
 * This software is released under the MIT License.
 * https://opensource.org/licenses/MIT
 */

#include "client/client.h"
#include "assert/assert.h"
#include "console/console.h"

errorret_t networkedClientConnect(
  client_t *client,
  const networkedclientconnect_t connInfo
) {
  int32_t ret;
  packet_t packet;
  errorret_t err;
  char_t *ip = "127.0.0.1";

  assertNotNull(client, "Client is NULL");
  assertTrue(client->type == CLIENT_TYPE_NETWORKED, "Client is not networked");
  assertIsMainThread("Client connect must be on main thread");

  client->state = CLIENT_STATE_CONNECTING;
  consolePrint("Connecting to server %s:%d", ip, connInfo.port);

  // Create a socket
  client->networked.socket = socket(AF_INET, SOCK_STREAM, 0);
  if(client->networked.socket < 0) {
    return error(
      "Failed to create socket %s",
      errno > 0 ? strerror(errno) : "Unknown error"
    );
  }

  // Set ip address and port
  client->networked.address.sin_family = AF_INET;
  client->networked.address.sin_port = htons(connInfo.port);
  client->networked.address.sin_addr.s_addr = inet_addr(ip);

  ret = inet_pton(AF_INET, ip, &client->networked.address.sin_addr);
  if(ret <= 0) {
    close(client->networked.socket);
    return error(
      "Invalid or bad IP address %s: %s",
      ip,
      errno > 0 ? strerror(errno) : "Unknown error"
    );
  }

  // Connect to the server
  ret = connect(
    client->networked.socket,
    (struct sockaddr *)&client->networked.address,
    sizeof(client->networked.address)
  );
  if(ret < 0) {
    close(client->networked.socket);
    switch(errno) {
      case ECONNREFUSED:
        return error("Failed to connect: Connection refused");
      case ETIMEDOUT:
        return error("Failed to connect: Connection timed out");
      case ENETUNREACH:
        return error("Failed to connect: Network unreachable");
      default:
        return error("Failed to connect: Unknown error");
    }
  }

  // Initialize mutex and condition variable
  pthread_mutex_init(&client->networked.lock, NULL);
  pthread_cond_init(&client->networked.cond, NULL);

  // Initialize read and write locks
  pthread_mutex_init(&client->networked.readLock, NULL);
  pthread_mutex_init(&client->networked.writeLock, NULL);

  // Send the version
  {
    const char_t *message = "DUSK|"DUSK_VERSION;
    ssize_t sent = send(
      client->networked.socket,
      message,
      strlen(message),
      0
    );
  }

  // We should now receive a welcome packet
  err = networkedClientReadPacket(client, &packet);
  if(err) return err;
  switch(packet.type) {
    case PACKET_TYPE_DISCONNECT:
      err = packetDisconnectClient(&packet);
      if(err) return err;
      break;
    
    case PACKET_TYPE_WELCOME:
      err = packetWelcomeClient(&packet);
      if(err) return err;
      break;
    
    default:
      return error("Server did not send welcome message.");
  }

  // Connection was established, hand off to read thread
  ret = pthread_create(
    &client->networked.readThread,
    NULL,
    networkedClientReadThread,
    client
  );

  if(ret != 0) {
    close(client->networked.socket);
    pthread_mutex_destroy(&client->networked.readLock);
    pthread_mutex_destroy(&client->networked.writeLock);
    return error(
      "Failed to create client read thread %s",
      errno > 0 ? strerror(errno) : "Unknown error"
    );
  }

  // Wait for the read thread to signal that it has started
  pthread_mutex_lock(&client->networked.lock);
  while (client->state == CLIENT_STATE_CONNECTING) {
    pthread_cond_wait(&client->networked.cond, &client->networked.lock);
  }
  pthread_mutex_unlock(&client->networked.lock);

  // Start the write thread after the handshake
  ret = pthread_create(
    &client->networked.writeThread,
    NULL,
    networkedClientWriteThread, // Renamed from networkedClientRightThread
    client
  );

  if(ret != 0) {
    close(client->networked.socket);
    pthread_mutex_destroy(&client->networked.readLock);
    pthread_mutex_destroy(&client->networked.writeLock);
    return error(
      "Failed to create client write thread %s",
      errno > 0 ? strerror(errno) : "Unknown error"
    );
  }

  return ERROR_OK;
}

void networkedClientDisconnect(client_t *client) {
  assertNotNull(client, "Client is NULL");
  assertIsMainThread("Client disconnect must be on main thread");
  assertTrue(client->type == CLIENT_TYPE_NETWORKED, "Client is not networked");
  assertTrue(client->state == CLIENT_STATE_CONNECTED, "Client not connected");

  pthread_mutex_lock(&client->networked.lock);
  client->state = CLIENT_STATE_DISCONNECTING;

  struct timespec timeout;
  clock_gettime(CLOCK_REALTIME, &timeout);
  timeout.tv_sec += 1; // Wait for up to 1 second

  while (client->state == CLIENT_STATE_DISCONNECTING) {
    if(pthread_cond_timedwait(
      &client->networked.cond,
      &client->networked.lock,
      &timeout
    ) == ETIMEDOUT) {
      consolePrint("Client disconnect timed out, force closing");
      break;
    }
  }
  pthread_mutex_unlock(&client->networked.lock);

  client->state = CLIENT_STATE_DISCONNECTED;

  // Wait for the threads to finish
  if (client->networked.readThread) {
    pthread_join(client->networked.readThread, NULL);
    client->networked.readThread = 0;
  }

  if (client->networked.writeThread) {
    pthread_join(client->networked.writeThread, NULL);
    client->networked.writeThread = 0;
  }

  // Destroy read and write locks
  pthread_mutex_destroy(&client->networked.readLock);
  pthread_mutex_destroy(&client->networked.writeLock);

  // Close the socket
  if (client->networked.socket) {
    shutdown(client->networked.socket, SHUT_RDWR);
    close(client->networked.socket);
    client->networked.socket = 0;
  }

  // Destroy mutex and condition variable
  pthread_mutex_destroy(&client->networked.lock);
  pthread_cond_destroy(&client->networked.cond);
}

errorret_t networkedClientWrite(
  const client_t *client,
  const uint8_t *data,
  const size_t len
) {
  assertNotNull(client, "Client is NULL");
  assertNotNull(data, "Data is NULL");
  assertTrue(len > 0, "Data length is 0");
  assertTrue(client->type == CLIENT_TYPE_NETWORKED, "Client is not networked");

  if(client->state == CLIENT_STATE_DISCONNECTED) {
    return error("Client is disconnected");
  }

  ssize_t sent = send(client->networked.socket, data, len, 0);
  if(sent < 0) return error("Failed to send data");
  return ERROR_OK;
}

errorret_t networkedClientWritePacket(
  const client_t *client,
  const packet_t *packet
) {
  assertNotNull(packet, "Packet is NULL");
  assertTrue(packet->type != PACKET_TYPE_INVALID, "Packet type is INVALID");
  assertTrue(packet->length > 0, "Packet length is 0");
  assertTrue(
    packet->length <= sizeof(packetdata_t),
    "Packet length is too large (1)"
  );

  size_t fullSize = sizeof(packet_t) - sizeof(packet->data) + packet->length;
  assertTrue(fullSize <= sizeof(packet_t), "Packet size is too large (2)");
  return networkedClientWrite(client, (const uint8_t *)packet, fullSize);
}

errorret_t networkedClientReadPacket(
  const client_t *client,
  packet_t *packet
) {
  uint8_t buffer[sizeof(packet_t)];

  assertNotNull(client, "Client is NULL");
  assertNotNull(packet, "Packet is NULL");
  assertTrue(client->type == CLIENT_TYPE_NETWORKED, "Client is not networked");

  if(client->state == CLIENT_STATE_DISCONNECTED) {
    return error("Client is disconnected");
  }

  // Read the packet header
  ssize_t read = recv(
    client->networked.socket,
    buffer,
    sizeof(packettype_t),
    0
  );
  if(read != sizeof(packettype_t)) {
    return error(
      "Failed to read packet header %s",
      errno > 0 ? strerror(errno) : "Unknown error"
    );
  }

  packet->type = *(packettype_t *)buffer;

  // Read the packet length
  read = recv(
    client->networked.socket,
    buffer,
    sizeof(uint32_t),
    0
  );
  if(read != sizeof(uint32_t)) {
    return error(
      "Failed to read packet length %s",
      errno > 0 ? strerror(errno) : "Unknown error"
    );
  }
  if(read > sizeof(packetdata_t)) {
    return error("Packet length is too large");
  }
  packet->length = *(uint32_t *)buffer;

  // Now, read the packet data
  read = recv(
    client->networked.socket,
    (uint8_t *)&packet->data,
    packet->length,
    0
  );
  if(read != packet->length) {
    return error(
      "Failed to read packet data %s",
      errno > 0 ? strerror(errno) : "Unknown error"
    );
  }
  return ERROR_OK;
}

void * networkedClientReadThread(void *arg) {
  assertNotNull(arg, "Client thread argument is NULL");
  assertNotMainThread("Client thread must not be on main thread");
  
  packet_t packet;
  client_t *client = (client_t *)arg;
  assertTrue(
    client->type == CLIENT_TYPE_NETWORKED,
    "Client thread argument is not networked"
  );
  assertTrue(
    client->state == CLIENT_STATE_CONNECTING,
    "Client thread argument is not connecting"
  );

  // Notify the main thread that we are connected
  pthread_mutex_lock(&client->networked.lock);
  client->state = CLIENT_STATE_CONNECTED;
  pthread_cond_signal(&client->networked.cond);
  pthread_mutex_unlock(&client->networked.lock);

  while(client->state == CLIENT_STATE_CONNECTED) {
    pthread_mutex_lock(&client->networked.readLock);
    
    // Read a packet from the server
    errorret_t err = networkedClientReadPacket(client, &packet);
    if(err) {
      consolePrint("Failed to read packet %s", errorString());
      errorFlush();
      break;
    }

    packetQueuePushIn(
      &client->packetQueue,
      &packet
    );

    pthread_mutex_unlock(&client->networked.readLock);
  }

  client->state = CLIENT_STATE_DISCONNECTED;

  pthread_mutex_lock(&client->networked.lock);
  client->state = CLIENT_STATE_DISCONNECTED;
  pthread_cond_signal(&client->networked.cond); // Notify the main thread
  pthread_mutex_unlock(&client->networked.lock);
}

void * networkedClientWriteThread(void *arg) {
  packet_t packet;
  int32_t ret;
  errorret_t err;

  assertNotNull(arg, "Write thread argument is NULL");
  assertNotMainThread("Write thread must not be on main thread");

  client_t *client = (client_t *)arg;
  assertTrue(
    client->type == CLIENT_TYPE_NETWORKED,
    "Write thread argument is not networked"
  );

  while(client->state == CLIENT_STATE_CONNECTED) {
    pthread_mutex_lock(&client->networked.writeLock);

    ret = packetQueuePopOut(
      &client->packetQueue,
      &packet
    );
    
    if(ret == 0) {
      pthread_mutex_unlock(&client->networked.writeLock);
      continue;
    }

    if(ret < 0) {
      consolePrint("Failed to pop packet from queue %s", errorString());
      errorFlush();
      client->state = CLIENT_STATE_DISCONNECTING;
      pthread_mutex_unlock(&client->networked.writeLock);
      break;
    }

    err = networkedClientWritePacket(client, &packet);
    if(err) {
      consolePrint("Failed to write packet %s", errorString());
      errorFlush();
      client->state = CLIENT_STATE_DISCONNECTING;
      pthread_mutex_unlock(&client->networked.writeLock);
      break;
    }

    pthread_mutex_unlock(&client->networked.writeLock); // Unlock after writing
  }

  return NULL;
}