#include <stdio.h>
#include <string.h>

#include "pico/error.h"
#include "pico/cyw43_arch.h"
#include "pico/stdio_uart.h"
#include "pico/stdlib.h"
#include "tusb.h"

#include "lwip/apps/mqtt.h"
#include "lwip/dns.h"
#include "lwip/ip4_addr.h"
#include "lwip/netif.h"
#include "lwip/timeouts.h"

#include "modbus_rtu.h"

#include "wifi_config.h"
#include "mqtt_config.h"

// ---- User settings (mirror your ESPHome config) ----
#define MODBUS_SLAVE_ID       1
#define MODBUS_BAUDRATE       9600
// Device expects a correct CRC in the query, but may respond with an incorrect CRC.
// So: TX includes CRC, RX ignores CRC validation.
#define MODBUS_DISABLE_CRC    1
#define POLL_INTERVAL_MS      5000
#define RESPONSE_TIMEOUT_MS   800

// We read holding registers 0..19 (20 regs)
#define READ_START_ADDR       0
#define READ_QUANTITY         20

// We'll use UART0 for logs because RP2040 USB is used for host.
static void log_uart_init(void) {
  // Note: these are GPIO numbers (GP16/GP17), not physical header pin numbers.
  stdio_uart_init_full(uart0, 115200, 16, 17);
  sleep_ms(50);
}

// ---- Debug helpers ----
#ifndef MODBUS_DEBUG
#define MODBUS_DEBUG 1
#endif

static void hexdump_line(const uint8_t *buf, size_t len, size_t max_len) {
  const size_t n = (len < max_len) ? len : max_len;
  for (size_t i = 0; i < n; i++) {
    printf("%02X", buf[i]);
    if (i + 1 != n) printf(" ");
  }
  if (len > n) printf(" ...");
}

static uint32_t ms_now(void) {
  return to_ms_since_boot(get_absolute_time());
}

// ---- MQTT publisher (lwIP) ----
static mqtt_client_t *g_mqtt = NULL;
static volatile bool g_mqtt_connected = false;
static ip_addr_t g_mqtt_broker_ip;
static volatile bool g_mqtt_dns_done = false;
static volatile bool g_mqtt_dns_ok = false;

static const char *lwip_err_name(err_t err) {
  switch (err) {
    case ERR_OK: return "ERR_OK";
    case ERR_MEM: return "ERR_MEM";
    case ERR_BUF: return "ERR_BUF";
    case ERR_TIMEOUT: return "ERR_TIMEOUT";
    case ERR_RTE: return "ERR_RTE";
    case ERR_INPROGRESS: return "ERR_INPROGRESS";
    case ERR_VAL: return "ERR_VAL";
    case ERR_WOULDBLOCK: return "ERR_WOULDBLOCK";
    case ERR_USE: return "ERR_USE";
    case ERR_ALREADY: return "ERR_ALREADY";
    case ERR_ISCONN: return "ERR_ISCONN";
    case ERR_CONN: return "ERR_CONN";
    case ERR_IF: return "ERR_IF";
    case ERR_ABRT: return "ERR_ABRT";
    case ERR_RST: return "ERR_RST";
    case ERR_CLSD: return "ERR_CLSD";
    case ERR_ARG: return "ERR_ARG";
    default: return "ERR_?";
  }
}

static const char *mqtt_status_name(mqtt_connection_status_t st) {
  switch (st) {
    case MQTT_CONNECT_ACCEPTED: return "ACCEPTED";
    case MQTT_CONNECT_REFUSED_PROTOCOL_VERSION: return "REFUSED_PROTOCOL";
    case MQTT_CONNECT_REFUSED_IDENTIFIER: return "REFUSED_ID";
    case MQTT_CONNECT_REFUSED_SERVER: return "REFUSED_SERVER";
    case MQTT_CONNECT_REFUSED_USERNAME_PASS: return "REFUSED_USERPASS";
    case MQTT_CONNECT_REFUSED_NOT_AUTHORIZED_: return "REFUSED_NOAUTH";
    case MQTT_CONNECT_DISCONNECTED: return "DISCONNECTED";
    case MQTT_CONNECT_TIMEOUT: return "TIMEOUT";
    default: return "UNKNOWN";
  }
}

static void mqtt_request_cb(void *arg, err_t err) {
  (void)arg;
  if (err != ERR_OK) {
    printf("MQTT: publish ack err=%d (%s)\n", (int)err, lwip_err_name(err));
  }
}

static void mqtt_connection_cb(mqtt_client_t *client, void *arg, mqtt_connection_status_t status) {
  (void)client;
  (void)arg;
  g_mqtt_connected = (status == MQTT_CONNECT_ACCEPTED);
  printf("MQTT: connection status=%d (%s)\n", (int)status, mqtt_status_name(status));
}

static void mqtt_dns_found_cb(const char *name, const ip_addr_t *ipaddr, void *arg) {
  (void)name;
  (void)arg;
  if (ipaddr) {
    g_mqtt_broker_ip = *ipaddr;
    g_mqtt_dns_ok = true;
  } else {
    g_mqtt_dns_ok = false;
  }
  g_mqtt_dns_done = true;
}

static bool mqtt_resolve_broker(uint32_t timeout_ms) {
  const char *host = MQTT_BROKER_HOST;

  if (ipaddr_aton(host, &g_mqtt_broker_ip)) {
    return true;
  }

  g_mqtt_dns_done = false;
  g_mqtt_dns_ok = false;
  ip_addr_set_any(IPADDR_TYPE_ANY, &g_mqtt_broker_ip);

  cyw43_arch_lwip_begin();
  ip_addr_t cached;
  err_t derr = dns_gethostbyname(host, &cached, mqtt_dns_found_cb, NULL);
  if (derr == ERR_OK) {
    g_mqtt_broker_ip = cached;
    g_mqtt_dns_ok = true;
    g_mqtt_dns_done = true;
  } else if (derr != ERR_INPROGRESS) {
    g_mqtt_dns_ok = false;
    g_mqtt_dns_done = true;
  }
  cyw43_arch_lwip_end();

  const uint32_t start = ms_now();
  while (!g_mqtt_dns_done && (ms_now() - start) < timeout_ms) {
    cyw43_arch_lwip_begin();
    sys_check_timeouts();
    cyw43_arch_lwip_end();
    sleep_ms(10);
  }

  return g_mqtt_dns_done && g_mqtt_dns_ok;
}

static void mqtt_init_and_connect(void) {
  if (!g_mqtt) {
    g_mqtt = mqtt_client_new();
    if (!g_mqtt) {
      printf("MQTT: mqtt_client_new failed\n");
      return;
    }
  }

  if (!mqtt_resolve_broker(3000)) {
    printf("MQTT: DNS failed for '%s'\n", MQTT_BROKER_HOST);
    return;
  }

  const struct mqtt_connect_client_info_t ci = {
      .client_id = MQTT_CLIENT_ID,
      .client_user = MQTT_USERNAME,
      .client_pass = MQTT_PASSWORD,
      .keep_alive = 60,
      .will_topic = NULL,
      .will_msg = NULL,
      .will_msg_len = 0,
      .will_qos = 0,
      .will_retain = 0,
  };

  cyw43_arch_lwip_begin();
  err_t err = mqtt_client_connect(g_mqtt, &g_mqtt_broker_ip, (u16_t)MQTT_BROKER_PORT,
                                 mqtt_connection_cb, NULL, &ci);
  cyw43_arch_lwip_end();

  if (err != ERR_OK) {
    printf("MQTT: connect err=%d\n", (int)err);
  }
}

static void mqtt_publish_readings(float active_power_w,
                                  float rms_current_a,
                                  float voltage_v,
                                  float frequency_hz,
                                  float power_factor,
                                  float annual_kwh,
                                  float active_kwh,
                                  float reactive_kwh,
                                  float load_time_h,
                                  uint16_t hours_day,
                                  uint16_t device_addr) {
  if (!g_mqtt_connected) {
    // Best-effort reconnect.
    mqtt_init_and_connect();
    if (!g_mqtt_connected) return;
  }

  static char payload[320];
  const uint32_t uptime_ms = ms_now();
  const int n = snprintf(payload, sizeof(payload),
                         "{\"uptime_ms\":%lu,\"active_power_w\":%.2f,\"rms_current_a\":%.3f,\"voltage_v\":%.1f,"
                         "\"frequency_hz\":%.2f,\"power_factor\":%.3f,\"annual_kwh\":%.3f,\"active_kwh\":%.3f,"
                         "\"reactive_kwh\":%.3f,\"load_time_h\":%.2f,\"hours_day\":%u,\"device_addr\":%u}",
                         (unsigned long)uptime_ms,
                         active_power_w,
                         rms_current_a,
                         voltage_v,
                         frequency_hz,
                         power_factor,
                         annual_kwh,
                         active_kwh,
                         reactive_kwh,
                         load_time_h,
                         (unsigned)hours_day,
                         (unsigned)device_addr);
  if (n <= 0 || (size_t)n >= sizeof(payload)) return;

  cyw43_arch_lwip_begin();
  err_t err = mqtt_publish(g_mqtt, MQTT_TOPIC_READINGS,
                           payload, (u16_t)n,
                           (u8_t)MQTT_QOS, (u8_t)MQTT_RETAIN,
                           mqtt_request_cb, NULL);
  cyw43_arch_lwip_end();

  if (err != ERR_OK) {
    printf("MQTT: publish call err=%d (%s) payload_len=%d topic='%s'\n",
           (int)err, lwip_err_name(err), n, MQTT_TOPIC_READINGS);
  }
}

static modbus_rtu_master_t g_mb;
static uint8_t g_cdc_idx = 0xFF;

static uint32_t g_usb_rx_bytes_total = 0;
static uint32_t g_usb_rx_last_print_ms = 0;

typedef enum {
  APP_WAIT_USB = 0,
  APP_IDLE,
  APP_WAIT_RESPONSE,
} app_state_t;

static app_state_t g_state = APP_WAIT_USB;
static uint32_t g_next_poll_ms = 0;

static uint32_t wifi_backoff_ms(uint32_t attempt_1_based) {
  // attempt=1 -> base, attempt=2 -> 2x, attempt=3 -> 4x, ... clamped.
  uint32_t delay = WIFI_RETRY_BASE_DELAY_MS;
  if (attempt_1_based > 1) {
    const uint32_t shift = attempt_1_based - 1;
    if (shift >= 31) {
      delay = WIFI_RETRY_MAX_DELAY_MS;
    } else {
      uint32_t mult = 1u << shift;
      if (mult == 0) mult = 0xFFFFFFFFu;
      delay = WIFI_RETRY_BASE_DELAY_MS * mult;
    }
  }
  if (delay > WIFI_RETRY_MAX_DELAY_MS) delay = WIFI_RETRY_MAX_DELAY_MS;
  // Add small jitter (0..255ms) to avoid sync collisions.
  delay += (time_us_32() & 0xFFu);
  return delay;
}

static bool wifi_wait_for_dhcp_ip(char *ip_buf, size_t ip_buf_len) {
  absolute_time_t deadline = make_timeout_time_ms(WIFI_DHCP_TIMEOUT_MS);
  while (absolute_time_diff_us(get_absolute_time(), deadline) > 0) {
    bool has_ip = false;
    if (ip_buf && ip_buf_len) ip_buf[0] = '\0';

    cyw43_arch_lwip_begin();
    struct netif *netif = &cyw43_state.netif[CYW43_ITF_STA];
    const ip4_addr_t *ip = netif_ip4_addr(netif);
    if (netif_is_up(netif) && ip && !ip4_addr_isany(ip)) {
      if (ip_buf && ip_buf_len) {
        snprintf(ip_buf, ip_buf_len, "%s", ip4addr_ntoa(ip));
      }
      has_ip = true;
    }
    cyw43_arch_lwip_end();

    if (has_ip) return true;
    sleep_ms(200);
  }
  return false;
}

static bool wifi_init_and_connect(void) {
  int init_err = cyw43_arch_init_with_country(WIFI_COUNTRY);
  if (init_err != 0) {
    printf("WiFi init failed: %d\n", init_err);
    return false;
  }

  cyw43_arch_enable_sta_mode();

  for (uint32_t attempt = 1; attempt <= WIFI_CONNECT_RETRIES; attempt++) {

    printf("WiFi: connecting to SSID '%s' (attempt %lu/%u)...\n",
           WIFI_SSID, (unsigned long)attempt, (unsigned)WIFI_CONNECT_RETRIES);
    int err = cyw43_arch_wifi_connect_timeout_ms(
        WIFI_SSID,
        WIFI_PASSWORD,
        WIFI_AUTH,
        WIFI_CONNECT_TIMEOUT_MS);
    if (err != 0) {
      if (err == PICO_ERROR_BADAUTH) {
        printf("WiFi connect failed (bad auth): %d\n", err);
        return false;
      }
      printf("WiFi connect failed: %d\n", err);
      const uint32_t delay = wifi_backoff_ms(attempt);
      printf("WiFi: retrying in %lu ms\n", (unsigned long)delay);
      sleep_ms(delay);
      continue;
    }

    char ip_buf[16] = {0};
    if (!wifi_wait_for_dhcp_ip(ip_buf, sizeof(ip_buf))) {
      printf("WiFi DHCP timeout (attempt %lu/%u)\n",
             (unsigned long)attempt, (unsigned)WIFI_CONNECT_RETRIES);
      const uint32_t delay = wifi_backoff_ms(attempt);
      printf("WiFi: retrying in %lu ms\n", (unsigned long)delay);
      sleep_ms(delay);
      continue;
    }

    printf("WiFi up. IP: %s\n", ip_buf);

    // Connect MQTT once at boot (best-effort).
    mqtt_init_and_connect();
    return true;
  }

  printf("WiFi: giving up after %u attempts\n", (unsigned)WIFI_CONNECT_RETRIES);
  return false;
}

static void start_poll(void) {
  uint8_t frame[16];
  // Always include CRC on TX (your device requires it).
  size_t n = modbus_rtu_build_read_holding(frame, sizeof(frame), MODBUS_SLAVE_ID, READ_START_ADDR, READ_QUANTITY, false);
  if (n == 0) return;

#if MODBUS_DEBUG
  printf("Modbus TX (%u bytes): ", (unsigned)n);
  hexdump_line(frame, n, 32);
  printf("\n");
#endif

  // Clear any stale RX bytes
  g_mb.rx_len = 0;

  tuh_cdc_write(g_cdc_idx, frame, (uint32_t)n);
  tuh_cdc_write_flush(g_cdc_idx);

  g_mb.awaiting_response = true;
  g_mb.request_started_ms = ms_now();
  g_state = APP_WAIT_RESPONSE;
}

static void print_decoded(const uint16_t regs[READ_QUANTITY]) {
  float active_power_w = modbus_float_dcba(regs[0], regs[1]);
  float rms_current_a  = modbus_float_dcba(regs[2], regs[3]);
  float voltage_v      = modbus_float_dcba(regs[4], regs[5]);
  float frequency_hz   = modbus_float_dcba(regs[6], regs[7]);
  float power_factor   = modbus_float_dcba(regs[8], regs[9]);
  float annual_kwh     = modbus_float_dcba(regs[10], regs[11]);
  float active_kwh     = modbus_float_dcba(regs[12], regs[13]);
  float reactive_kwh   = modbus_float_dcba(regs[14], regs[15]);
  float load_time_h    = modbus_float_dcba(regs[16], regs[17]) / 60.0f;
  uint16_t hours_day   = regs[18];
  uint16_t device_addr = regs[19];

  printf("Active Power: %.2f W\n", active_power_w);
  printf("RMS Current: %.3f A\n", rms_current_a);
  printf("Voltage: %.1f V\n", voltage_v);
  printf("Frequency: %.2f Hz\n", frequency_hz);
  printf("Power Factor: %.3f\n", power_factor);
  printf("Annual Power Consumption: %.3f kWh\n", annual_kwh);
  printf("Active Consumption: %.3f kWh\n", active_kwh);
  printf("Reactive Consumption: %.3f kWh\n", reactive_kwh);
  printf("Load Time: %.2f h\n", load_time_h);
  printf("Work Hours Per Day: %u h\n", hours_day);
  printf("Device Address: %u\n", device_addr);
  printf("----\n");

  // Publish once per refresh.
  mqtt_publish_readings(active_power_w,
                        rms_current_a,
                        voltage_v,
                        frequency_hz,
                        power_factor,
                        annual_kwh,
                        active_kwh,
                        reactive_kwh,
                        load_time_h,
                        hours_day,
                        device_addr);
}

static void app_task(void) {
  const uint32_t now = ms_now();

  if (g_state == APP_IDLE) {
    if ((int32_t)(now - g_next_poll_ms) >= 0) {
      start_poll();
      g_next_poll_ms = now + POLL_INTERVAL_MS;
    }
    return;
  }

  if (g_state == APP_WAIT_RESPONSE) {
    // Non-blocking receive: accumulate whatever is available.
    while (tuh_cdc_read_available(g_cdc_idx)) {
      uint8_t tmp[64];
      uint32_t n = tuh_cdc_read(g_cdc_idx, tmp, sizeof(tmp));
      if (n == 0) break;

#if MODBUS_DEBUG
      g_usb_rx_bytes_total += n;
      // Print at most every 100ms to avoid flooding UART.
      if ((now - g_usb_rx_last_print_ms) > 100) {
        printf("USB RX chunk (%lu bytes), total=%lu: ", (unsigned long)n, (unsigned long)g_usb_rx_bytes_total);
        hexdump_line(tmp, (size_t)n, 32);
        printf("\n");
        printf("Modbus RX buf len=%u: ", (unsigned)g_mb.rx_len);
        hexdump_line(g_mb.rx_buf, g_mb.rx_len, 32);
        printf("\n");
        g_usb_rx_last_print_ms = now;
      }
#endif

      modbus_rtu_feed(&g_mb, tmp, (size_t)n);
    }

    size_t data_off = 0;
    uint16_t word_count = 0;

    if (modbus_rtu_try_parse_read_holding_response(&g_mb, MODBUS_SLAVE_ID, READ_QUANTITY, &data_off, &word_count)) {
#if MODBUS_DEBUG
      printf("Modbus: parsed response (word_count=%u, data_off=%u, rx_len=%u)\n",
             (unsigned)word_count, (unsigned)data_off, (unsigned)g_mb.rx_len);
#endif
      uint16_t regs[READ_QUANTITY];
      memset(regs, 0, sizeof(regs));

      // Data is big-endian 16-bit registers.
      const uint8_t *p = &g_mb.rx_buf[data_off];
      for (uint16_t i = 0; i < word_count && i < READ_QUANTITY; i++) {
        regs[i] = modbus_rtu_get_u16_be(&p[i * 2]);
      }

      print_decoded(regs);

      // Drop the parsed frame from RX buffer
      const size_t crc_len = MODBUS_DISABLE_CRC ? 0 : 2;
      const size_t total_len = 3 + (size_t)(READ_QUANTITY * 2) + crc_len;
      // (Reuse internal helper behavior by just resetting for single outstanding request)
      g_mb.rx_len = 0;

      g_state = APP_IDLE;
      return;
    }

    if ((now - g_mb.request_started_ms) > RESPONSE_TIMEOUT_MS) {
#if MODBUS_DEBUG
      const uint32_t age = now - g_mb.request_started_ms;
      printf("Modbus timeout (no valid response) age=%lu ms, usb_total=%lu, rx_len=%u\n",
             (unsigned long)age, (unsigned long)g_usb_rx_bytes_total, (unsigned)g_mb.rx_len);
      if (g_mb.rx_len) {
        printf("Modbus RX buf: ");
        hexdump_line(g_mb.rx_buf, g_mb.rx_len, 64);
        printf("\n");
      }
      printf("----\n");
#else
      printf("Modbus timeout (no valid response)\n----\n");
#endif
      g_mb.rx_len = 0;
      g_state = APP_IDLE;
      return;
    }
  }
}

// TinyUSB Host CDC callbacks
void tuh_cdc_mount_cb(uint8_t idx) {
  g_cdc_idx = idx;

  // Configure UART settings on the adapter (supported by ACM + CH34x; some other chips may vary).
  cdc_line_coding_t lc = {
    .bit_rate = MODBUS_BAUDRATE,
    .stop_bits = CDC_LINE_CODING_STOP_BITS_1,
    .parity = CDC_LINE_CODING_PARITY_NONE,
    .data_bits = 8,
  };

  tuh_cdc_set_line_coding(idx, &lc, NULL, 0);
  tuh_cdc_connect(idx, NULL, 0); // assert DTR/RTS

#if MODBUS_DEBUG
  printf("USB-serial mounted (cdc idx=%u) line=%lu %u%c%u\n",
         idx,
         (unsigned long)lc.bit_rate,
         (unsigned)lc.data_bits,
         (lc.parity == CDC_LINE_CODING_PARITY_NONE) ? 'N' :
         (lc.parity == CDC_LINE_CODING_PARITY_ODD) ? 'O' :
         (lc.parity == CDC_LINE_CODING_PARITY_EVEN) ? 'E' : '?',
         (lc.stop_bits == CDC_LINE_CODING_STOP_BITS_1) ? 1 : 2);
#else
  printf("USB-serial mounted (cdc idx=%u)\n", idx);
#endif
  g_state = APP_IDLE;
  g_next_poll_ms = ms_now();
}

void tuh_cdc_umount_cb(uint8_t idx) {
  if (g_cdc_idx == idx) g_cdc_idx = 0xFF;
  printf("USB-serial unmounted (cdc idx=%u)\n", idx);
  g_state = APP_WAIT_USB;
}

void tuh_cdc_rx_cb(uint8_t idx) {
  // We also poll in app_task(); keep this empty to avoid duplicating logic.
  (void)idx;
}

int main(void) {
  log_uart_init();
  printf("pico_usb_modbus_host starting...\n");

  // Optional: Wi-Fi STA + DHCP (Pico W)
  // If this fails, we keep running the USB host application.
  (void)wifi_init_and_connect();

  modbus_rtu_master_config_t cfg = {
    .slave_id = MODBUS_SLAVE_ID,
    .disable_crc = (MODBUS_DISABLE_CRC != 0),
  };
  modbus_rtu_master_init(&g_mb, cfg);

  // Initialize TinyUSB host stack (RP2040 USB controller as host)
  tuh_init(0);

  while (true) {
    tuh_task();

    if (g_state != APP_WAIT_USB && g_cdc_idx != 0xFF && tuh_cdc_mounted(g_cdc_idx)) {
      app_task();
    }

    // Keep loop responsive.
    sleep_ms(1);
  }
}