#include <stdio.h>
#include <string.h>
#include <strings.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <netdb.h>
#include <netinet/in.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>

#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_log.h"
#include "esp_tls.h"
#include "esp_crt_bundle.h"

#include "main.h"
#include "at_uart.h"
#include "at_cip.h"
#include "wifi.h"

static const char *TAG = "at_cip";

static int  cip_sock = -1;
static int  cip_type = CIP_TYPE_NONE;
static char cip_host[64];
static int  cip_port = 0;
static esp_tls_t *cip_tls = NULL;

void at_cip_init()
{
  cip_sock = -1;
  cip_type = CIP_TYPE_NONE;
  cip_host[0] = 0;
  cip_port = 0;
  cip_tls = NULL;
}

bool at_cip_is_connected()
{
  return cip_sock >= 0 || cip_tls != NULL;
}

void at_cip_status_line(char *buf, size_t buflen)
{
  if (!at_cip_is_connected())
  {
    snprintf(buf, buflen, "STATUS:0");
    return;
  }

  const char *type_str = cip_type == CIP_TYPE_UDP ? "UDP" : cip_type == CIP_TYPE_SSL ? "SSL" : "TCP";
  snprintf(buf, buflen, "STATUS:3\r\n+CIPSTATUS:0,\"%s\",\"%s\",%d,0,0", type_str, cip_host, cip_port);
}

// Resolves and connects a raw socket via the plain BSD gethostbyname()/connect()
// path. Shared by plain TCP and (as of the SSL support below) also by the TLS
// path: esp_tls's own built-in connect step goes through getaddrinfo() instead,
// which crashes (LoadProhibited, NULL FreeRTOS queue handle deep inside lwIP's
// tcpip_task message-passing - confirmed via backtrace against esp-spi.elf,
// root cause not fully understood) the first time it's exercised from this
// task. gethostbyname() uses a different, older lwIP DNS code path that has
// always worked fine here, so SSL connects through it too and only hands
// esp_tls an already-open socket, skipping esp_tls's DNS/connect step entirely.
static int at_cip_connect_raw_tcp(int sock_type, const char *host, int port, bool set_rcv_timeout)
{
  // gethostbyname() (and, for the SSL caller, esp_tls's own getaddrinfo()
  // fallback below) both post through tcpip_task's mailbox - if this is the
  // very first network command since boot (no AT+CWJAP/AT+CWLAP yet), that
  // mailbox doesn't exist yet and the post crashes with a NULL-queue-handle
  // LoadProhibited (EXCVADDR=0x54, confirmed via console backtrace: reproduced
  // with a plain AT+CIPSTART="TCP" to a literal IP right after a fresh boot,
  // fixed by ensuring the WiFi stack is up first, same lazy-init used by
  // AT+CWLAP - safe to call again if already initialized).
  if (!net_mtx)
    initialize_wifi();

  struct hostent *server = gethostbyname(host);
  if (!server)
  {
    ESP_LOGW(TAG, "DNS lookup failed for '%s'", host);
    return -1;
  }

  int sock = socket(AF_INET, sock_type, 0);
  if (sock < 0)
  {
    ESP_LOGW(TAG, "socket() failed");
    return -1;
  }

  struct sockaddr_in addr;
  memset(&addr, 0, sizeof(addr));
  addr.sin_family = AF_INET;
  addr.sin_port = htons(port);
  memcpy(&addr.sin_addr.s_addr, server->h_addr, server->h_length);

  // A short receive timeout lets at_cip_rx_task poll periodically instead of
  // blocking forever, so it notices at_cip_close() promptly. Skipped here for
  // the SSL caller: applying it before the TLS handshake made mbedtls's own
  // blocking reads see spurious timeouts mid-handshake (a partial/timed-out
  // read on a socket mbedtls expects to block on desyncs its record parser -
  // confirmed by this exact change turning consistent MBEDTLS_ERR_SSL_INVALID_RECORD
  // handshake failures into a clean handshake). The SSL path sets this itself
  // after the handshake completes, once the socket moves to steady-state reads.
  if (set_rcv_timeout)
  {
    struct timeval tv = { 0, 200 * 1000 };
    setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
  }

  if (sock_type == SOCK_DGRAM)
    return sock; // UDP: no connect() handshake to wait on

  // Bound the connect() itself with an explicit timeout via a non-blocking
  // connect + select(): an unreachable-but-not-immediately-rejected address
  // (e.g. a black-holed IP) can otherwise leave the OS/lwip TCP SYN retry
  // logic blocking this call for a minute or more, freezing the whole AT
  // command interface for that entire time.
  int flags = fcntl(sock, F_GETFL, 0);
  fcntl(sock, F_SETFL, flags | O_NONBLOCK);

  int rc = connect(sock, (struct sockaddr*)&addr, sizeof(addr));

  if (rc < 0 && errno != EINPROGRESS)
  {
    ESP_LOGW(TAG, "connect() to %s:%d failed immediately", host, port);
    close(sock);
    return -1;
  }

  if (rc < 0) // EINPROGRESS - wait for it to complete
  {
    fd_set wfds;
    FD_ZERO(&wfds);
    FD_SET(sock, &wfds);
    struct timeval conn_tv = { 10, 0 }; // 10s connect timeout

    int sel = select(sock + 1, NULL, &wfds, NULL, &conn_tv);

    if (sel <= 0)
    {
      ESP_LOGW(TAG, "connect() to %s:%d timed out", host, port);
      close(sock);
      return -1;
    }

    int so_err = 0;
    socklen_t so_len = sizeof(so_err);
    getsockopt(sock, SOL_SOCKET, SO_ERROR, &so_err, &so_len);

    if (so_err != 0)
    {
      ESP_LOGW(TAG, "connect() to %s:%d failed: errno %d", host, port, so_err);
      close(sock);
      return -1;
    }
  }

  fcntl(sock, F_SETFL, flags); // back to blocking for the (rare) case send() would need it
  return sock;
}

static int at_cip_start_ssl(const char *host, int port)
{
  int sock = at_cip_connect_raw_tcp(SOCK_STREAM, host, port, false);
  if (sock < 0)
    return -1;

  // Kept local until the handshake fully succeeds below, deliberately not
  // written to the global cip_tls yet: at_cip_is_connected() (checked by
  // at_cip_rx_task's poll loop, running concurrently in its own task) goes by
  // cip_tls != NULL, and that task calling esp_tls_conn_read() on a context
  // that's mid-handshake raced with esp_tls_conn_new_sync() below over the
  // same mbedtls ssl_context - confirmed via mbedtls debug log showing
  // interleaved flush-output calls and a corrupted ssl->out_left, manifesting
  // as MBEDTLS_ERR_SSL_INTERNAL_ERROR then MBEDTLS_ERR_SSL_INVALID_RECORD.
  // The plain-TCP path never had this problem because it keeps its socket in
  // a local variable the same way, only publishing to cip_sock at the very end.
  esp_tls_t *tls = esp_tls_init();
  if (!tls)
  {
    ESP_LOGW(TAG, "esp_tls_init failed");
    close(sock);
    return -1;
  }

  esp_tls_set_conn_sockfd(tls, sock);
  esp_tls_set_conn_state(tls, ESP_TLS_CONNECTING); // skip esp_tls's own DNS/connect step - see at_cip_connect_raw_tcp()

  esp_tls_cfg_t cfg = {};
  cfg.crt_bundle_attach = esp_crt_bundle_attach;
  cfg.timeout_ms = 10000; // matches the connect timeout already spent above

  int rc = esp_tls_conn_new_sync(host, strlen(host), port, &cfg, tls);
  if (rc != 1)
  {
    ESP_LOGW(TAG, "TLS handshake with %s:%d failed", host, port);
    esp_tls_conn_destroy(tls); // also closes sock, it now owns the fd
    return -1;
  }

  // Now that the handshake is done, throttle reads the same way the plain-TCP
  // path does, so at_cip_rx_task's poll loop notices at_cip_close() promptly.
  struct timeval tv = { 0, 200 * 1000 };
  setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));

  cip_tls = tls; // now safe to publish: handshake is done, at_cip_rx_task may use it
  cip_type = CIP_TYPE_SSL;
  strlcpy(cip_host, host, sizeof(cip_host));
  cip_port = port;

  ESP_LOGI(TAG, "connected SSL %s:%d", host, port);
  return 0;
}

int at_cip_start(const char *type, const char *host, int port)
{
  if (at_cip_is_connected())
    return -2; // already connected

  if (!strcasecmp(type, "SSL"))
    return at_cip_start_ssl(host, port);

  int sock_type = 0;

  if (!strcasecmp(type, "TCP"))
    sock_type = SOCK_STREAM;
  else if (!strcasecmp(type, "UDP"))
    sock_type = SOCK_DGRAM;
  else
  {
    ESP_LOGW(TAG, "unknown type '%s'", type);
    return -1;
  }

  int sock = at_cip_connect_raw_tcp(sock_type, host, port, true);
  if (sock < 0)
    return -1;

  cip_sock = sock;
  cip_type = (sock_type == SOCK_DGRAM) ? CIP_TYPE_UDP : CIP_TYPE_TCP;
  strlcpy(cip_host, host, sizeof(cip_host));
  cip_port = port;

  ESP_LOGI(TAG, "connected %s %s:%d", cip_type == CIP_TYPE_UDP ? "UDP" : "TCP", host, port);
  return 0;
}

int at_cip_send(const uint8_t *data, size_t len)
{
  if (cip_tls) return (int)esp_tls_conn_write(cip_tls, data, len);
  if (cip_sock < 0) return -1;
  return send(cip_sock, data, len, 0);
}

void at_cip_close()
{
  if (!at_cip_is_connected()) return;

  if (cip_tls)
  {
    esp_tls_conn_destroy(cip_tls);
    cip_tls = NULL;
  }
  else
  {
    close(cip_sock);
    cip_sock = -1;
  }

  cip_type = CIP_TYPE_NONE;
  cip_host[0] = 0;
  cip_port = 0;
}

void at_cip_rx_task(void *arg)
{
  uint8_t buf[1024];

  while (1)
  {
    if (!at_cip_is_connected())
    {
      vTaskDelay(pdMS_TO_TICKS(100));
      continue;
    }

    esp_tls_t *tls = cip_tls; // snapshot; at_cip_close() may change these concurrently
    int sock = cip_sock;
    int n = tls ? (int)esp_tls_conn_read(tls, buf, sizeof(buf)) : recv(sock, buf, sizeof(buf), 0);

    if (n > 0)
    {
      char hdr[24];
      snprintf(hdr, sizeof(hdr), "\r\n+IPD,%d:", n);
      at_send_str(hdr);
      at_send_raw(buf, n);
    }
    else if (n == 0)
    {
      // Peer closed (TCP/SSL only - UDP recv() doesn't return 0 on a connected socket).
      at_cip_close();
      at_send_str("\r\nCLOSED\r\n");
    }
    // n < 0: either the SO_RCVTIMEO timeout expired (EAGAIN/EWOULDBLOCK, expected -
    // esp_tls_conn_read() surfaces the same underlying-socket timeout as a generic
    // negative return too, not distinguished from a real error here, matching how
    // the plain-TCP path already treats every negative read as "retry") or the
    // connection was closed out from under us by at_cip_close() (also fine,
    // at_cip_is_connected() will be false at the top next time).
  }
}