/**
 * TCP Client for Pico W
 * Simple TCP client for sending sensor data messages to a backend server
 */

#include "tcp_client.h"
#include "lwip/pbuf.h"
#include "lwip/tcp.h"
#include "pico/cyw43_arch.h"
#include "pico/stdlib.h"
#include "pico/time.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define DEBUG_printf printf
#define MAX_RETRY_COUNT 3

typedef struct {
  struct tcp_pcb *tcp_pcb;
  ip_addr_t remote_addr;
  uint16_t remote_port;
  const char *message;
  int message_len;
  int sent_len;
  bool complete;
  bool success;
  bool connected;
  uint32_t timeout_ms;
  absolute_time_t timeout_time;
} tcp_client_state_t;

static err_t tcp_client_close(tcp_client_state_t *state) {
  err_t err = ERR_OK;
  if (state->tcp_pcb != NULL) {
    tcp_arg(state->tcp_pcb, NULL);
    tcp_poll(state->tcp_pcb, NULL, 0);
    tcp_sent(state->tcp_pcb, NULL);
    tcp_recv(state->tcp_pcb, NULL);
    tcp_err(state->tcp_pcb, NULL);
    err = tcp_close(state->tcp_pcb);
    if (err != ERR_OK) {
      DEBUG_printf("tcp_client: close failed %d, calling abort\n", err);
      tcp_abort(state->tcp_pcb);
      err = ERR_ABRT;
    }
    state->tcp_pcb = NULL;
  }
  return err;
}

static err_t tcp_client_result(tcp_client_state_t *state, int status) {
  if (status == 0) {
    DEBUG_printf("tcp_client: send success\n");
    state->success = true;
  } else {
    DEBUG_printf("tcp_client: send failed %d\n", status);
    state->success = false;
  }
  state->complete = true;
  return tcp_client_close(state);
}

static err_t tcp_client_sent(void *arg, struct tcp_pcb *tpcb, u16_t len) {
  tcp_client_state_t *state = (tcp_client_state_t *)arg;
  DEBUG_printf("tcp_client: sent %u bytes\n", len);
  state->sent_len += len;

  if (state->sent_len >= state->message_len) {
    DEBUG_printf("tcp_client: message sent completely\n");
    return tcp_client_result(state, 0);
  }

  return ERR_OK;
}

static err_t tcp_client_connected(void *arg, struct tcp_pcb *tpcb, err_t err) {
  tcp_client_state_t *state = (tcp_client_state_t *)arg;
  if (err != ERR_OK) {
    DEBUG_printf("tcp_client: connect failed %d\n", err);
    return tcp_client_result(state, err);
  }

  state->connected = true;
  DEBUG_printf("tcp_client: connected, sending message\n");

  // Send the message
  cyw43_arch_lwip_begin();
  err_t write_err =
      tcp_write(tpcb, state->message, state->message_len, TCP_WRITE_FLAG_COPY);
  if (write_err != ERR_OK) {
    DEBUG_printf("tcp_client: failed to write data %d\n", write_err);
    cyw43_arch_lwip_end();
    return tcp_client_result(state, -1);
  }

  // Flush the data
  err_t output_err = tcp_output(tpcb);
  cyw43_arch_lwip_end();

  if (output_err != ERR_OK) {
    DEBUG_printf("tcp_client: failed to output data %d\n", output_err);
    return tcp_client_result(state, -1);
  }

  return ERR_OK;
}

static err_t tcp_client_poll(void *arg, struct tcp_pcb *tpcb) {
  tcp_client_state_t *state = (tcp_client_state_t *)arg;
  DEBUG_printf("tcp_client: poll timeout\n");
  return tcp_client_result(state, -1);
}

static void tcp_client_err(void *arg, err_t err) {
  tcp_client_state_t *state = (tcp_client_state_t *)arg;
  if (err != ERR_ABRT) {
    DEBUG_printf("tcp_client: error %d\n", err);
    state->success = false;
    state->complete = true;
  }
}

static err_t tcp_client_recv(void *arg, struct tcp_pcb *tpcb, struct pbuf *p,
                             err_t err) {
  tcp_client_state_t *state = (tcp_client_state_t *)arg;

  if (!p) {
    DEBUG_printf("tcp_client: connection closed by server\n");
    return tcp_client_result(state, 0);
  }

  cyw43_arch_lwip_check();

  if (p->tot_len > 0) {
    DEBUG_printf("tcp_client: received %d bytes (ignoring)\n", p->tot_len);
    tcp_recved(tpcb, p->tot_len);
  }
  pbuf_free(p);

  return ERR_OK;
}

static bool tcp_client_open(tcp_client_state_t *state) {
  DEBUG_printf("tcp_client: connecting to %s port %u\n",
               ip4addr_ntoa(&state->remote_addr), state->remote_port);

  state->tcp_pcb = tcp_new_ip_type(IP_GET_TYPE(&state->remote_addr));
  if (!state->tcp_pcb) {
    DEBUG_printf("tcp_client: failed to create pcb\n");
    return false;
  }

  tcp_arg(state->tcp_pcb, state);
  tcp_poll(state->tcp_pcb, tcp_client_poll, 10);
  tcp_sent(state->tcp_pcb, tcp_client_sent);
  tcp_recv(state->tcp_pcb, tcp_client_recv);
  tcp_err(state->tcp_pcb, tcp_client_err);

  state->sent_len = 0;
  state->timeout_time = make_timeout_time_ms(state->timeout_ms);

  cyw43_arch_lwip_begin();
  err_t err = tcp_connect(state->tcp_pcb, &state->remote_addr,
                          state->remote_port, tcp_client_connected);
  cyw43_arch_lwip_end();

  if (err != ERR_OK) {
    DEBUG_printf("tcp_client: tcp_connect failed %d\n", err);
    return false;
  }

  return true;
}

bool tcp_client_init(tcp_client_config *config, const char *server_ip,
                     uint16_t server_port, uint32_t timeout_ms) {
  if (!config || !server_ip) {
    return false;
  }

  strncpy(config->server_ip, server_ip, sizeof(config->server_ip) - 1);
  config->server_ip[sizeof(config->server_ip) - 1] = '\0';
  config->server_port = server_port;
  config->timeout_ms = timeout_ms;
  config->initialized = true;
  config->internal_state = NULL;

  DEBUG_printf("tcp_client: initialized for %s:%u\n", config->server_ip,
               config->server_port);
  return true;
}

bool tcp_client_send_message(tcp_client_config *config, const char *message) {
  if (!config || !config->initialized || !message) {
    DEBUG_printf("tcp_client: invalid config or message\n");
    return false;
  }

  tcp_client_state_t *state = calloc(1, sizeof(tcp_client_state_t));
  if (!state) {
    DEBUG_printf("tcp_client: failed to allocate state\n");
    return false;
  }

  // Convert IP address
  if (!ip4addr_aton(config->server_ip, &state->remote_addr)) {
    DEBUG_printf("tcp_client: invalid IP address %s\n", config->server_ip);
    free(state);
    return false;
  }

  state->remote_port = config->server_port;
  state->message = message;
  state->message_len = strlen(message);
  state->timeout_ms = config->timeout_ms;
  state->complete = false;
  state->success = false;
  state->connected = false;

  DEBUG_printf("tcp_client: sending message (%d bytes): %s\n",
               state->message_len, message);

  if (!tcp_client_open(state)) {
    DEBUG_printf("tcp_client: failed to open connection\n");
    free(state);
    return false;
  }

  // Wait for completion or timeout
  while (!state->complete) {
    cyw43_arch_poll();
    cyw43_arch_wait_for_work_until(make_timeout_time_ms(1000));
    if (time_reached(state->timeout_time)) {
      DEBUG_printf("tcp_client: operation timed out\n");
      tcp_client_close(state);
      state->complete = true;
      state->success = false;
      break;
    }
  }

  bool success = state->success;
  free(state);

  return success;
}

void tcp_client_cleanup(tcp_client_config *config) {
  if (config) {
    config->initialized = false;
    config->internal_state = NULL;
  }
}