#include "mqtt_handler.h"
#include "config.h"
#include "modbus.h"
#include <ArduinoJson.h>
#include <esp_wifi.h>

WiFiClient espClient;
PubSubClient client(espClient);

// --- MQTT ACK 전송 함수 ---
void publish_ack(const char* command_id, const char* status) {
  if (!client.connected()) return;
  
  JsonDocument doc;
  doc["command_id"] = command_id;
  doc["status"] = status;
  doc["timestamp"] = millis();
  
  char buffer[128];
  serializeJson(doc, buffer, sizeof(buffer));
  client.publish("/line1/ack", buffer);
  Serial.printf("[MQTT ACK PUBLISH] Published status %s for command %s to /line1/ack\n", status, command_id);
}

// --- MQTT 수신 콜백: [[id, val], ...] 배치 파싱 및 제어 적용 ---
void on_mqtt_message(char* topic, byte* payload, unsigned int length) {
  // 수신된 페이로드 로그 출력
  char log_buf[256];
  unsigned int log_len = (length < 255) ? length : 255;
  memcpy(log_buf, payload, log_len);
  log_buf[log_len] = '\0';
  Serial.printf("\n[MQTT RECEIVED] Topic: %s | Len: %d | Payload: %s\n", topic, length, log_buf);

  // 명령(cmd) 토픽 처리
  if (strcmp(topic, "/line1/cmd") == 0) {
    char cmd_str[32];
    unsigned int len = (length < 31) ? length : 31;
    memcpy(cmd_str, payload, len);
    cmd_str[len] = '\0';
    
    char* clean_cmd = cmd_str;
    if (clean_cmd[0] == '"') clean_cmd++;
    int clean_len = strlen(clean_cmd);
    if (clean_len > 0 && clean_cmd[clean_len - 1] == '"') {
      clean_cmd[clean_len - 1] = '\0';
    }
    
    if (strcmp(clean_cmd, "sync") == 0) {
      Serial.println("[CMD PROCESS] Manual Full-Sync Request Received!");
      
      if (pending_command_active) {
        publish_ack(pending_command_id, "FAILED");
      }
      
      setup_buffers(); 
      currentState = STATE_INIT;
      lastStateChange = millis();
      initSequenceStarted = false; 
      
      if (client.connected()) {
        client.publish("/line1/event", "{\"id\":999,\"val\":1}");
        Serial.println("[CMD PROCESS] Published event 999 to broker.");
      }
      return;
    }
  }

  JsonDocument doc;
  DeserializationError error = deserializeJson(doc, payload, length);
  if (error) {
    Serial.printf("[MQTT ERROR] JSON Deserialization failed: %s\n", error.c_str());
    return;
  }
  
  const char* cmd_id = nullptr;
  JsonArray arr;
  bool has_envelope = false;

  // JSON 객체(새로운 Envelope 규격)와 JSON 배열(기존 규격) 둘 다 호환되도록 처리
  if (doc.is<JsonObject>()) {
    JsonObject obj = doc.as<JsonObject>();
    if (obj["command_id"].is<const char*>() && obj["items"].is<JsonArray>()) {
      cmd_id = obj["command_id"].as<const char*>();
      arr = obj["items"].as<JsonArray>();
      has_envelope = true;
      Serial.printf("[MQTT ENVELOPE] command_id: %s\n", cmd_id);
    } else {
      Serial.println("[MQTT ERROR] JSON Object requires string command_id and array items");
      return;
    }
  } else if (doc.is<JsonArray>()) {
    arr = doc.as<JsonArray>();
  } else {
    Serial.println("[MQTT ERROR] Payload is neither JSON Array nor Object");
    return;
  }

  const bool is_m_topic = strcmp(topic, "/line1/m/out") == 0;
  const bool is_d_topic = strcmp(topic, "/line1/d/out") == 0;
  bool valid_items = (is_m_topic || is_d_topic) && !arr.isNull() && arr.size() > 0;

  // ArduinoJson 7에서 중첩 배열을 JsonVariant -> JsonArray로 명시 변환해 검증한다.
  for (JsonVariant item_variant : arr) {
    if (!item_variant.is<JsonArray>()) {
      valid_items = false;
      break;
    }
    JsonArray item = item_variant.as<JsonArray>();
    if (item.size() < 2 || !item[0].is<int>()) {
      valid_items = false;
      break;
    }

    int id = item[0].as<int>();
    if (is_m_topic) {
      if (!item[1].is<int>()) {
        valid_items = false;
        break;
      }
      int val = item[1].as<int>();
      if (id < 0 || id >= M_OUT_COILS || (val != 0 && val != 1)) {
        valid_items = false;
        break;
      }
    } else if (is_d_topic) {
      if (id < 0 || id >= 4 || !item[1].is<float>()) {
        valid_items = false;
        break;
      }
    }
  }

  if (!valid_items) {
    Serial.printf("[MQTT ERROR] Invalid or empty items for topic %s; command will not be written to PLC\n", topic);
    if (has_envelope && cmd_id != nullptr) {
      publish_ack(cmd_id, "FAILED");
    }
    return;
  }

  // Envelope 규격 수신 시 상태를 RECEIVED로 즉시 업데이트하고 전송 대기 정보 설정
  if (has_envelope && cmd_id != nullptr) {
    if (!plcOnline || currentState == STATE_INIT) {
      Serial.printf("[MQTT REJECT] Command %s rejected because PLC is offline or in INIT state\n", cmd_id);
      publish_ack(cmd_id, "FAILED");
      return;
    }
    if (pending_command_active) {
      Serial.printf("[MQTT WARN] Command %s rejected because command %s is still active (BUSY)\n", cmd_id, pending_command_id);
      publish_ack(cmd_id, "BUSY");
      return;
    }
    strncpy(pending_command_id, cmd_id, sizeof(pending_command_id) - 1);
    pending_command_id[sizeof(pending_command_id) - 1] = '\0';
    pending_command_active = true;
    pending_command_time = millis();
    pending_command_type_is_m = is_m_topic;
    snapshot_pending_command_buffers(pending_command_type_is_m);
    
    publish_ack(pending_command_id, "RECEIVED");
  }

  Serial.printf("[MQTT PROCESS] Parsed array size: %d\n", arr.size());
  
  if (is_m_topic) {
    for (JsonVariant item_variant : arr) {
      JsonArray item = item_variant.as<JsonArray>();
      if (item.size() >= 2) {
        int id = item[0];
        int val = item[1];
        int word_idx = id / 16;
        int bit_idx = id % 16;
        
        if (word_idx >= 0 && word_idx < M_OUT_WORDS) {
          uint16_t before_word = M_Out_New[word_idx];
          Serial.printf("[M_OUT REG] Requested bit: %d | val: %d | Calculated word_idx: %d, bit_idx: %d\n", id, val, word_idx, bit_idx);
          if (val) {
            M_Out_New[word_idx] |= (1 << bit_idx);
            M_Out_Pending_Write[word_idx] |= (1 << bit_idx);
            M_Out_Pending_Clear[word_idx] &= ~(1 << bit_idx);
          } else {
            M_Out_New[word_idx] &= ~(1 << bit_idx);
            M_Out_Pending_Write[word_idx] &= ~(1 << bit_idx);
            M_Out_Pending_Clear[word_idx] &= ~(1 << bit_idx);
          }
          Serial.printf("[M_OUT PROCESS completed] Updated M_Out_New[%d] to 0x%04X (Pending Write: 0x%04X)\n", word_idx, M_Out_New[word_idx], M_Out_Pending_Write[word_idx]);
          Serial.printf(
            "[M CMD TRACE] command_id=%s ui_bit=M%d plc_target=M%04d%X coil_addr=%d requested=%d word_before=0x%04X word_after=0x%04X pending=0x%04X\n",
            pending_command_active ? pending_command_id : "legacy",
            id,
            200 + word_idx,
            bit_idx,
            M_OUT_START_ADDR + id,
            val,
            before_word,
            M_Out_New[word_idx],
            M_Out_Pending_Write[word_idx]
          );
        } else {
          Serial.printf("[M_OUT LIMIT ERROR] word_idx %d is out of bounds (max: %d)\n", word_idx, M_OUT_WORDS);
        }
      }
    }
  } 
  else if (is_d_topic) {
    for (JsonVariant item_variant : arr) {
      JsonArray item = item_variant.as<JsonArray>();
      if (item.size() >= 2) {
        int id = item[0];
        float val = item[1].as<float>();
        
        if (id >= 0 && id < 4) {
          uint32_t raw;
          memcpy(&raw, &val, sizeof(raw));
          uint16_t w_low = raw & 0xFFFF;
          uint16_t w_high = (raw >> 16) & 0xFFFF;
          
          D_Out_New[id * 2] = w_low;
          D_Out_New[id * 2 + 1] = w_high;
          D_Out_Pending_Groups |= (1U << id);
          
          Serial.printf("[D_OUT REG] Requested Group ID: %d | REAL: %.6g | Low Word: 0x%04X, High Word: 0x%04X\n", id, val, w_low, w_high);
        } else {
          Serial.printf("[D_OUT LIMIT ERROR] Group ID %d is out of bounds (max: 4)\n", id);
        }
      }
    }
  }
}

// --- MQTT 초기화 및 최초 구독 연결 설정 ---
void setup_mqtt() {
  client.setServer(mqtt_server, 1883);
  client.setCallback(on_mqtt_message);
  client.setBufferSize(2048); 
  Serial.println("[MQTT SETUP] Server configured to 1883. Buffer size set to 2048.");
}

// --- MQTT 연결 유지 및 비차단 재연결 함수 ---
void maintain_mqtt() {
  if (WiFi.status() != WL_CONNECTED) {
    return; 
  }
  
  if (!client.connected()) {
    unsigned long now = millis();
    if (now - lastMqttRetry >= 5000) {
      lastMqttRetry = now;
      Serial.print("[MQTT LOOP] Attempting connection...");
      
      uint8_t mac[6];
      WiFi.macAddress(mac);
      char clientId[30];
      snprintf(clientId, sizeof(clientId), "ESP32_Gateway_%02X%02X%02X", mac[3], mac[4], mac[5]);
      
      // Last Will & Testament 설정 (어플리케이션 비정상 종료 시 브로커가 OFFLINE 메시지 자동 배포)
      if (client.connect(clientId, mqtt_user, mqtt_password, "/line1/status", 1, true, "{\"gateway\": \"OFFLINE\"}")) {
        Serial.println(" OK (connected)");
        
        // 연결 성공 직후 retained SYNCING 상태 전송 (초기화 중이므로)
        char status_msg[160];
        snprintf(status_msg, sizeof(status_msg),
                 "{\"gateway\": \"SYNCING\", \"plc_state\": \"%s\", \"firmware\": \"%s\"}",
                 plcOnline ? "ONLINE" : "OFFLINE", FIRMWARE_VERSION);
        client.publish("/line1/status", status_msg, true);
        
        // MQTT 연결 복구 시 Full-Sync 강제 수행 상태로 전이 (동기 정합성 확보)
        Serial.println("[MQTT LOOP] Reconnected successfully. Triggering Full-Sync...");
        setup_buffers();
        currentState = STATE_INIT;
        lastStateChange = millis();
        initSequenceStarted = false;
        
        client.publish("/line1/event", "{\"id\":999,\"val\":1}");
        
        client.subscribe("/line1/m/out");
        client.subscribe("/line1/d/out");
        client.subscribe("/line1/cmd");
        Serial.println("[MQTT LOOP] Subscribed to /line1/m/out, /line1/d/out, /line1/cmd successfully.");
      } else {
        Serial.print(" FAILED, rc=");
        Serial.println(client.state());
      }
    }
  } else {
    client.loop();
    
    // 주기적인 (5초 간격) Retained Heartbeat 발행
    static unsigned long lastHeartbeat = 0;
    if (millis() - lastHeartbeat >= 5000) {
      char status_msg[160];
      const char* gw_state = (currentState == STATE_INIT) ? "SYNCING" : "ONLINE";
      snprintf(status_msg, sizeof(status_msg),
               "{\"gateway\": \"%s\", \"plc_state\": \"%s\", \"firmware\": \"%s\"}",
               gw_state, plcOnline ? "ONLINE" : "OFFLINE", FIRMWARE_VERSION);
      client.publish("/line1/status", status_msg, true);
      lastHeartbeat = millis();
    }
  }

  // --- 명령 처리 타임아웃 감시 ---
  // M 명령: 3초. D 명령: D 쓰기 + Apply ON + Apply OFF의 여러 통신 슬롯을 포함하므로 6초.
  unsigned long cmd_timeout = pending_command_type_is_m ? 3000 : 6000;
  if (pending_command_active && (millis() - pending_command_time > cmd_timeout)) {
    Serial.printf("[TIMEOUT] Command %s expired (%lums). Force unlocking and rolling back buffers.\n",
                  pending_command_id, cmd_timeout);
    publish_ack(pending_command_id, "FAILED");

    // Old는 모멘터리 ON 상태일 수 있으므로, 명령 수신 직전 스냅샷으로 복원한다.
    rollback_pending_command_buffers();

    // D Apply 펄스가 걸려 있었다면 Apply 비트와 대기 상태를 정리한다.
    if (d_apply_pending && d_apply_word_idx >= 0 && d_apply_word_idx < M_OUT_WORDS) {
      M_Out_New[d_apply_word_idx] &= ~d_apply_word_mask;
      M_Out_Pending_Write[d_apply_word_idx] &= ~d_apply_word_mask;
      M_Out_Pending_Clear[d_apply_word_idx] &= ~d_apply_word_mask;
    }
    d_apply_pending = false;
    d_apply_word_mask = 0;
    d_apply_word_idx = -1;

    pending_command_active = false;
    pending_command_id[0] = '\0';
    pending_command_time = 0;
  }
}

// --- M영역 변경 감지 및 배치 MQTT 전송 (JsonDocument 힙 고갈 방지를 위해 3개 워드씩 Chunk 분할 전송, 워드 경계 일치로 변경 유실 차단) ---
bool detect_and_publish_M() {
  if (!client.connected()) return false;
  
  bool overall_success = true;
  JsonDocument doc;
  JsonArray arr = doc.to<JsonArray>();

  int pending_words[M_IN_WORDS];
  int pending_word_count = 0;

  for (int i = 0; i < M_IN_WORDS; i++) {
    // Full-Sync 대기 상태인 경우 0xFFFF를 대입하여 전체 비트 송신 강제
    uint16_t diff = full_sync_pending ? 0xFFFF : (M_In_New[i] ^ M_In_Old[i]);
    if (diff != 0) {
      pending_words[pending_word_count++] = i;
      for (int j = 0; j < 16; j++) {
        if ((diff >> j) & 1) {
          int bit_id = i * 16 + j;
          int val = (M_In_New[i] >> j) & 1;
          
          Serial.printf("[LOG M_IN CHANGE] PLC M-Read Bit Offset %d (M%d) changed to %d\n", bit_id, 1000 + bit_id, val);
          
          JsonArray pair = arr.add<JsonArray>();
          pair.add(bit_id);
          pair.add(val);
        }
      }
    }
    
    // 워드 경계 정렬을 위해 루프 본문 완료 시점에 3개 워드 단위로 청크 발행 판단
    if (pending_word_count >= 3) {
      char buffer[2048];
      serializeJson(doc, buffer, sizeof(buffer));
      
      // 발행 성공 시에만 해당 청크에 완전히 포함된 워드들의 Old 버퍼를 New와 동기화
      bool ok = client.publish("/line1/m/in", buffer);
      Serial.printf("[M_IN PUBLISH Chunk] size=%d success=%d payload=%s\n", arr.size(), ok, buffer);
      if (ok) {
        for (int w = 0; w < pending_word_count; w++) {
          M_In_Old[pending_words[w]] = M_In_New[pending_words[w]];
        }
      } else {
        overall_success = false;
      }
      pending_word_count = 0;
      
      doc.clear();
      arr = doc.to<JsonArray>();
    }
  }

  if (arr.size() > 0) {
    char buffer[2048];
    serializeJson(doc, buffer, sizeof(buffer));
    
    // 발행 성공 시에만 남은 워드들의 Old 버퍼를 New와 동기화
    bool ok = client.publish("/line1/m/in", buffer);
    Serial.printf("[M_IN PUBLISH Chunk Last] size=%d success=%d payload=%s\n", arr.size(), ok, buffer);
    if (ok) {
      for (int w = 0; w < pending_word_count; w++) {
        M_In_Old[pending_words[w]] = M_In_New[pending_words[w]];
      }
    } else {
      overall_success = false;
    }
  }
  return overall_success;
}

// --- D영역 변경 감지 및 배치 MQTT 전송 ---
bool detect_and_publish_D() {
  if (!client.connected()) return false;
  
  bool overall_success = true;
  JsonDocument doc;
  JsonArray arr = doc.to<JsonArray>();

  int pending_groups[10];
  int pending_group_count = 0;

  for (int g = 0; g < 10; g++) {
    uint16_t w_low_new = D_In_New[g * 2];
    uint16_t w_high_new = D_In_New[g * 2 + 1];
    
    uint16_t w_low_old = D_In_Old[g * 2];
    uint16_t w_high_old = D_In_Old[g * 2 + 1];
    
    // Full-Sync 대기 상태인 경우 무조건 송신 강제
    if (full_sync_pending || w_low_new != w_low_old || w_high_new != w_high_old) {
      JsonArray pair = arr.add<JsonArray>();
      pair.add(g);

      if (g == 0 || g == 5) {
        // Pump restart time is an unsigned 16-bit value in the first word.
        // The second word is still read to preserve the fixed 32-bit group layout, but is not part of the value.
        pair.add(w_low_new);
        Serial.printf("[LOG D_IN CHANGE] PLC D-Read Group %d (D%d only, UINT16) changed from %u to %u\n",
                      g, 2000 + g * 2, w_low_old, w_low_new);
      } else {
        // All other groups contain an IEEE-754 REAL split into low/high PLC words.
        uint32_t raw_new = ((uint32_t)w_high_new << 16) | w_low_new;
        uint32_t raw_old = ((uint32_t)w_high_old << 16) | w_low_old;
        float val;
        float old_val;
        memcpy(&val, &raw_new, sizeof(val));
        memcpy(&old_val, &raw_old, sizeof(old_val));
        pair.add(val);
        Serial.printf("[LOG D_IN CHANGE] PLC D-Read Group %d (D%d-D%d, REAL) changed from %.6g to %.6g\n",
                      g, 2000 + g * 2, 2000 + g * 2 + 1, old_val, val);
      }
      
      pending_groups[pending_group_count++] = g;
    }
  }

  if (arr.size() > 0) {
    char buffer[1024];
    serializeJson(doc, buffer, sizeof(buffer));
    
    // 발행 성공 시에만 해당 그룹의 Old 버퍼를 New와 동기화
    bool ok = client.publish("/line1/d/in", buffer);
    Serial.printf("[D_IN PUBLISH Batch] size=%d success=%d payload=%s\n", arr.size(), ok, buffer);
    if (ok) {
      for (int w = 0; w < pending_group_count; w++) {
        int idx = pending_groups[w];
        D_In_Old[idx * 2] = D_In_New[idx * 2];
        D_In_Old[idx * 2 + 1] = D_In_New[idx * 2 + 1];
      }
    } else {
      overall_success = false;
    }
  }
  return overall_success;
}

// --- Read 변경사항 종합 프로세스 실행 ---
bool process_read_changes() {
  bool success_m = detect_and_publish_M();
  bool success_d = detect_and_publish_D();
  return success_m && success_d;
}