#include "esphome.h"

using namespace esphome;

static const char *TAG = "lwz";

/**
 * Global data storage for LWZ heat pump sensors.
 * This struct holds the latest parsed values from the UART stream.
 */
struct LWZStore {
  // System Info
  float firmware_version = NAN;
  
  // Temperature Sensors (°C)
  float collector_temp = NAN;
  float flow_temp = NAN;
  float return_temp = NAN;
  float hot_gas_temp = NAN;
  float dhw_temp = NAN;
  float flow_temp_hc2 = NAN;
  float inside_temp = NAN;
  float evaporator_temp = NAN;
  float condenser_temp = NAN;
  float outside_temp = NAN;
  float dew_point_temp = NAN;
  float compressor_temp = NAN;
  
  // Ventilation Fans (Hz / Speed)
  float extr_speed_actual = NAN;
  float vent_speed_actual = NAN;
  float expel_speed_actual = NAN;
  
  // System Pressure (bar) and Humidity (%)
  float rel_humidity = NAN;
  float p_nd = NAN;
  float p_hd = NAN;
  
  // Operational Counters (h)
  float compressor_heating = NAN;
  float compressor_cooling = NAN;
  float compressor_dhw = NAN;
  float booster_dhw = NAN;
  float booster_heating = NAN;
};

static LWZStore lwz_data;

/**
 * Core handler for LWZ serial communication.
 * Handles the DLE (Data Link Escape) protocol used by Stiebel Eltron / Tecalor.
 */
class LWZHandler {
 public:
  /**
   * Main entry point called by ESPHome's interval timer.
   */
  static void update(uart::UARTComponent *parent) {
    // 1. One-time query for firmware version if unknown
    if (std::isnan(lwz_data.firmware_version)) {
      ESP_LOGI(TAG, "Detecting Heat Pump Firmware Version...");
      poll_values(parent, 0xFD, 2);
    }

    ESP_LOGD(TAG, "Starting poll cycle...");
    
    // 2. Poll Actual Values (Temperatures, Fans, Pressures)
    poll_values(parent, 0xFB, 53);
    delay(500); 
    
    // 3. Poll Heating Values (Indoor temperature)
    poll_values(parent, 0xF4, 39);
    delay(500); 

    // 4. Poll Operational State (Running hours)
    poll_values(parent, 0x09, 10);
    
    ESP_LOGD(TAG, "Poll cycle finished.");
  }

 private:
  /**
   * Waits for a target byte from UART with a 1-second timeout.
   */
  static bool wait_for_byte(uart::UARTComponent *parent, uint8_t target, const char *label) {
    uint32_t start = millis();
    while (millis() - start < 1000) {
      if (parent->available()) {
        uint8_t b;
        parent->read_byte(&b);
        if (b == target) return true;
      }
      yield();
    }
    ESP_LOGW(TAG, "Timeout: Expected %02X for %s", target, label);
    return false;
  }

  /**
   * Reads a full DLE frame including escape handling and ETX check.
   */
  static bool read_frame(uart::UARTComponent *parent, std::vector<uint8_t> &frame) {
    bool escaping = false;
    uint32_t start = millis();
    while (millis() - start < 1500) {
      if (parent->available()) {
        uint8_t b;
        parent->read_byte(&b);
        if (escaping) {
          if (b == 0x03) return true;            // End of Text (ETX)
          if (b == 0x10) frame.push_back(0x10);  // Escaped 0x10 byte
          escaping = false;
        } else if (b == 0x10) {
          escaping = true;
        } else {
          frame.push_back(b);
        }
      }
      yield();
    }
    return false;
  }

  /**
   * Helper to parse 1 or 2 bytes into a fixed-point float.
   */
  static float get_fixed_point(uint8_t *p, size_t size, float divider) {
    int32_t val = 0;
    if (size == 1) val = (int8_t)p[0];
    else if (size == 2) val = (int16_t)((p[0] << 8) | p[1]);
    return val / divider;
  }

  /**
   * Standard polling sequence for a specific block ID.
   */
  static void poll_values(uart::UARTComponent *parent, uint8_t request_id, size_t expected_len) {
    // Clear buffer
    while(parent->available()) { uint8_t dummy; parent->read_byte(&dummy); }

    // Phase 1: Wakeup
    bool ack = false;
    for (int r = 0; r < 3 && !ack; r++) {
      parent->write_byte(0x02); 
      if (wait_for_byte(parent, 0x10, "Wakeup ACK")) ack = true;
      else delay(100);
    }
    if (!ack) return;

    // Phase 2: Command
    uint8_t checksum = (1 + request_id) & 0xFF;
    uint8_t cmd[] = {0x01, 0x00, checksum, request_id, 0x10, 0x03};
    parent->write_array(cmd, 6);

    if (!wait_for_byte(parent, 0x10, "Cmd ACK 1")) return;
    if (!wait_for_byte(parent, 0x02, "Cmd ACK 2")) return;

    // Phase 3: Trigger & Read
    parent->write_byte(0x10); 
    std::vector<uint8_t> frame;
    if (!read_frame(parent, frame)) return;

    /**
     * LWZ Protocol Quirk Fix:
     * Firmware versions (e.g. 4.19, 4.39) incorrectly insert a 0x18 byte after every 0x2B byte.
     * We filter this out to maintain data alignment.
     */
    for (size_t i = 0; i + 1 < frame.size(); i++) {
      if (frame[i] == 0x2B && frame[i+1] == 0x18) {
        frame.erase(frame.begin() + i + 1);
      }
    }

    if (frame.size() < 4 + expected_len) return;
    
    uint8_t *p = &frame[4]; // Data starts after DLE header

    // Mapping based on Firmware Version 4.19 / 4.39
    if (request_id == 0xFB) {
      lwz_data.collector_temp    = get_fixed_point(p + 0, 2, 10.0f);
      lwz_data.flow_temp         = get_fixed_point(p + 4, 2, 10.0f);
      lwz_data.return_temp       = get_fixed_point(p + 6, 2, 10.0f);
      lwz_data.hot_gas_temp      = get_fixed_point(p + 8, 2, 10.0f);
      lwz_data.dhw_temp          = get_fixed_point(p + 10, 2, 10.0f);
      lwz_data.flow_temp_hc2     = get_fixed_point(p + 12, 2, 10.0f);
      lwz_data.compressor_temp   = get_fixed_point(p + 14, 2, 10.0f);
      lwz_data.evaporator_temp   = get_fixed_point(p + 16, 2, 10.0f);
      lwz_data.condenser_temp    = get_fixed_point(p + 18, 2, 10.0f);
      lwz_data.extr_speed_actual = get_fixed_point(p + 29, 2, 1.0f);
      lwz_data.vent_speed_actual = get_fixed_point(p + 31, 2, 1.0f);
      lwz_data.expel_speed_actual = get_fixed_point(p + 33, 2, 1.0f);
      lwz_data.outside_temp      = get_fixed_point(p + 35, 2, 10.0f);
      lwz_data.rel_humidity      = get_fixed_point(p + 37, 2, 10.0f);
      lwz_data.dew_point_temp    = get_fixed_point(p + 39, 2, 10.0f);
      lwz_data.p_nd              = get_fixed_point(p + 41, 2, 100.0f);
      lwz_data.p_hd              = get_fixed_point(p + 43, 2, 100.0f);
    } else if (request_id == 0xF4) {
      lwz_data.inside_temp       = get_fixed_point(p + 32, 2, 10.0f);
    } else if (request_id == 0x09) {
      lwz_data.compressor_heating = (uint16_t)(p[0] << 8 | p[1]);
      lwz_data.compressor_cooling = (uint16_t)(p[2] << 8 | p[3]);
      lwz_data.compressor_dhw     = (uint16_t)(p[4] << 8 | p[5]);
      lwz_data.booster_dhw        = (uint16_t)(p[6] << 8 | p[7]);
      lwz_data.booster_heating    = (uint16_t)(p[8] << 8 | p[9]);
    } else if (request_id == 0xFD) {
      lwz_data.firmware_version   = get_fixed_point(p, 2, 100.0f);
      ESP_LOGI(TAG, "Heat Pump Version Detected: %.2f", lwz_data.firmware_version);
    }
    
    ESP_LOGD(TAG, "Parsed block ID %02X", request_id);

    // Reset sequence
    parent->write_byte(0x10); 
    parent->write_byte(0x02); 
  }
};
