#include <WiFi.h>
#include <Preferences.h>
#include "logs.h"
#include <Adafruit_NeoPixel.h>

#include <WiFiUdp.h>

WiFiUDP udp;
const int broadcastPort = 12345;
const char* broadcastMsg = "ESP32_SERVER_HERE";
bool serverFound = false;

// -------------------- НАСТРОЙКИ --------------------
char ssid[32] = "";
char password[32] = "";
char serverIP[32] = "";

#define NEOPIXEL_PIN 48
#define NUMPIXELS 1
#define BRIGHTNESS 40

//#define SERVER  // раскомментировать для сервера

// -------------------- ФУНКЦИИ --------------------

#ifndef SERVER
#undef NEOPIXEL_PIN
#define LED_PIN 8   // любой доступный цифровой пин для светодиода
#endif


#ifdef NEOPIXEL_PIN
Adafruit_NeoPixel pixels(NUMPIXELS, NEOPIXEL_PIN, NEO_GRB + NEO_KHZ800);
#endif
uint8_t brightness = 0; // 0 = выкл, 255 = макс
bool greenOn = false;

uint32_t sendInterval = 2000;    // 500 мс между отправками

LogModule logger;
Preferences prefs;   // для сохранения настроек

void saveSettings() {
    prefs.begin("network", false);
    prefs.putString("ssid", ssid);
    prefs.putString("pass", password);
    prefs.putString("ip", serverIP);
    prefs.putUInt("delay", sendInterval);  // сохраняем период отправки
    prefs.end();
    Serial.println("Settings saved");
}


void loadSettings() {
    prefs.begin("network", true);
    String s = prefs.getString("ssid", "Xiaomi 12");
    String p = prefs.getString("pass", "890089JJ");
    String ip = prefs.getString("ip", "10.251.43.192");
    sendInterval = prefs.getUInt("delay", 2000);  // восстанавливаем период
    prefs.end();

    s.toCharArray(ssid, sizeof(ssid));
    p.toCharArray(password, sizeof(password));
    ip.toCharArray(serverIP, sizeof(serverIP));

    Serial.println("Settings loaded:");
    Serial.print("SSID: "); Serial.println(ssid);
    Serial.print("PASS: "); Serial.println(password);
    Serial.print("IP: "); Serial.println(serverIP);
    Serial.print("Send interval: "); Serial.println(sendInterval);
}

void toggleGreen() {
#ifdef  NEOPIXEL_PIN
    greenOn = !greenOn; // меняем состояние
    if (greenOn) {
        pixels.setPixelColor(0, pixels.Color(0, BRIGHTNESS, 0)); // включаем зелёный
    } else {
        pixels.setPixelColor(0, pixels.Color(0, 0, 0));   // выключаем
    }
    pixels.show();
#else
    digitalWrite(LED_PIN, !digitalRead(LED_PIN));
#endif
}

void setRed() {
#ifdef  NEOPIXEL_PIN
    pixels.setPixelColor(0, pixels.Color(BRIGHTNESS, 0, 0));
    pixels.show();
#else
    digitalWrite(LED_PIN, HIGH); // включить
#endif
}

void setYellow() {
#ifdef  NEOPIXEL_PIN
    pixels.setPixelColor(0, pixels.Color(BRIGHTNESS, BRIGHTNESS, 0));
    pixels.show();
#else
    digitalWrite(LED_PIN, LOW);  // выключить (можно оставить так же)
#endif
}

void clearLED() {
#ifdef  NEOPIXEL_PIN
    pixels.setPixelColor(0, pixels.Color(0, 0, 0));
    pixels.show();
#else
    digitalWrite(LED_PIN, LOW);  // выключить (можно оставить так же)
#endif
}
// -------------------- РЕЖИМЫ --------------------
#ifdef SERVER
WiFiServer server(1234);
#else
WiFiClient client;
#endif
int run_client = 1;
int client_init = 0;

// -------------------- Wi-Fi --------------------
bool wifiConnecting = false;
bool wifiNeedReconnect = false;
uint32_t wifiStartMillis = 0;

void startWiFi() {
    Serial.println("Connecting to WiFi...");
    WiFi.disconnect(true, true);
    delay(1000);
    WiFi.begin(ssid, password);
    wifiConnecting = true;
    wifiStartMillis = millis();
}

void handleWiFi() {
    static uint32_t lastCheck = 0;
    const uint32_t interval = 500;

    if (millis() - lastCheck < interval) return;
    lastCheck = millis();

    if (wifiNeedReconnect) {
        wifiNeedReconnect = false;
        startWiFi();
        return;
    }

    wl_status_t status = WiFi.status();
    if (status == WL_CONNECTED) {
        if (wifiConnecting) {
            wifiConnecting = false;
            clearLED();
            Serial.println("WiFi connected");
            Serial.print("IP: "); Serial.println(WiFi.localIP());
#ifdef SERVER
            server.begin();
            setupUDPBroadcast();
#else
            startUDPListener();
            if(!client.connected())
            {
              if (client.connect(serverIP, 1234)) {
                  clearLED();
                  client_init = 1;
                  Serial.println("Connected to server");
              }
              else
              {
                setRed();
                Serial.println("Server connect error. Reconnecting...");
                wifiConnecting = false;
              }
            }
#endif
        }

    } else {
        setRed();
        if (wifiConnecting && millis() - wifiStartMillis > 10000) {  // 10 секунд таймаут
            Serial.println("WiFi connect timeout. Reconnecting...");
            startWiFi();
        } else if (!wifiConnecting) {
            Serial.println("WiFi disconnected. Scheduling reconnect...");
            wifiNeedReconnect = true;
        } else {
            Serial.print(".");
        }
    }
}

#ifndef SERVER // client
// -------------------- CLIENT --------------------
bool checkClientAlive() {
    if (!client.connected()) return false;

    client.write(""); // отправляем пустой пакет
    if (!client.connected()) {
        client.stop();
        

        client_init = 0;
        wifiConnecting = true;
        return false;
    }
    return true;
}
void startUDPListener() {
    udp.begin(broadcastPort);
    Serial.println("UDP listener started");
}

bool checkUDPBroadcast() {
    int packetSize = udp.parsePacket();
    if (packetSize) {
        char incoming[128];
        int len = udp.read(incoming, 127);
        if (len > 0) incoming[len] = 0;
        String msg = String(incoming);
        msg.trim();

        if (msg == "ESP32_SERVER_HERE") {
            IPAddress senderIP = udp.remoteIP();
            Serial.print("Server found at IP: "); Serial.println(senderIP);
            senderIP.toString().toCharArray(serverIP, sizeof(serverIP));
            serverFound = true;
            return true;
        }
    }
    return false;
}

#else //SERVER
// -------------------- SERVER --------------------

void setupUDPBroadcast() {
    udp.begin(broadcastPort);
    Serial.println("UDP broadcast ready");
}

void sendUDPBroadcast() {
    udp.beginPacket("255.255.255.255", broadcastPort);
    udp.write((const uint8_t*)broadcastMsg, strlen(broadcastMsg));

    udp.endPacket();
}

static WiFiClient clientConn;
static uint32_t lastClientMillis = 0;
const uint32_t CLIENT_TIMEOUT = 10000; // 10 секунд таймаут

void handleServer() {
    if (clientConn && clientConn.connected()) {
        if (millis() - lastClientMillis > CLIENT_TIMEOUT) {
            Serial.println("Client timeout, closing connection");
            clientConn.stop();
        }
    }

    if (!clientConn || !clientConn.connected()) {
        clientConn = server.available();
        if (clientConn) {
            lastClientMillis = millis();
            String clientIP = clientConn.remoteIP().toString();
            Serial.println("Client connected: " + clientIP);
        }
        return;
    }

    if (clientConn.available()) {
        lastClientMillis = millis(); // обновляем таймаут при активности
        String msg = clientConn.readStringUntil('\n');
        msg.trim();
        if (msg.length() == 0) return;

        LogEntry entry;
        entry.seq = 0;
        entry.ts = 0;
        entry.event_type = 0; // RECEIVE
        memset(entry.payload, 0, sizeof(entry.payload));

        int seqIndex = msg.indexOf("SEQ:");
        int tsIndex = msg.indexOf("TS:");
        int payloadIndex = msg.indexOf("PAYLOAD:");

        if (seqIndex >= 0 && tsIndex > seqIndex && payloadIndex > tsIndex) {
            String seqStr = msg.substring(seqIndex + 4, tsIndex); seqStr.trim();
            String tsStr = msg.substring(tsIndex + 3, payloadIndex); tsStr.trim();
            String payloadStr = msg.substring(payloadIndex + 8); payloadStr.trim();

            entry.seq = seqStr.toInt();
            entry.ts = strtoull(tsStr.c_str(), nullptr, 10);
            int len = min((int)payloadStr.length(), 16);
            payloadStr.toCharArray(entry.payload, len + 1);
        }

        // Сохраняем лог
        logger.writeLog(entry);
        Serial.print("Received: "); Serial.println(msg);

        toggleGreen();

        // Создаем SEND-запись
        LogEntry sendEntry = entry;
        sendEntry.ts = millis();
        sendEntry.event_type = 1; // SEND
        logger.writeLog(sendEntry);

        // Echo для клиента
        String echo = "SEQ:" + String(sendEntry.seq) +
                      " TS:" + String(sendEntry.ts) +
                      " EVT:SEND PAYLOAD:" + String(sendEntry.payload) + "\n";

        if (clientConn.print(echo)) {
            Serial.print("Sent: "); Serial.println(echo);
        } else {
            client_init = 0;
            Serial.println("Error sending to client");
            setRed();
        }
    }
}


#endif

// -------------------- UART ДЛЯ НАСТРОЕК --------------------
void handleUARTNetwork() {
    static String buffer;
    while (Serial.available()) {
        char c = Serial.read();
        if (c == '\n' || c == '\r') {
            buffer.trim();
            if (buffer.length() > 0) {
                int sep = buffer.indexOf(' ');
                if (sep > 0) {
                    String cmd = buffer.substring(0, sep);
                    String val = buffer.substring(sep + 1);
                    
                    if (cmd == "SSID") {
                        val.toCharArray((char*)ssid, 32);
                        Serial.print("\nSSID set to: "); Serial.println(ssid);                        
                        saveSettings();
                        wifiNeedReconnect = true; // для клиента тоже полезно
                    } 
                    else if (cmd == "PASS") {
                        val.toCharArray((char*)password, 32);
                        Serial.print("\nPassword set to: "); Serial.println(password);                 
                        saveSettings();
                        wifiNeedReconnect = true; // для клиента тоже полезно
                    } 
                    else if (cmd == "IP") {
                        val.toCharArray((char*)serverIP, 16);
                        Serial.print("\nServer IP set to: "); Serial.println(serverIP);         
                        saveSettings();
                        client_init = 0; // для клиента тоже полезно
                    }
                    else if (cmd == "DELAY") {
                        uint32_t newDelay = val.toInt();
                        if (newDelay > 0) {
                            sendInterval = newDelay;
                            Serial.print("\nSend interval set to: "); Serial.println(sendInterval);
                            saveSettings();
                        } else {
                            Serial.println("Invalid DELAY value");
                        }
                    }
                    else {
                        logger.handleUART(buffer[0]); // передаем неизвестные команды в логгер
                    }
                } else {
                    String cmd = buffer.substring(0, 1);
                    if (cmd == "s")
                    {
                      run_client = 0;
                    }
                    else if (cmd == "r")
                    {
                      run_client = 1;
                    }
                    else
                    {
                      logger.handleUART(buffer[0]); // нет пробела — считаем команду неизвестной
                    }
                }
            }
            buffer = "";
        } 
        else {
            buffer += c;
        }
    }
}



void setup() {
    Serial.begin(115200);
    delay(1000);

    logger.begin();
#ifdef  NEOPIXEL_PIN
    pixels.begin();
    pixels.show();
#else
  pinMode(LED_PIN, OUTPUT);
  digitalWrite(LED_PIN, HIGH); // по умолчанию выключен
#endif

    loadSettings();   // загружаем сохраненные настройки

#ifdef SERVER
    Serial.println("SERVER MODE: Enter SSID and PASS via UART, e.g.:\nSSID MyWiFi\nPASS 12345678");
#else
    Serial.println("CLIENT MODE: Enter server IP via UART, e.g.:\nIP 192.168.1.100");
#endif
}

void loop() {
    handleUARTNetwork();          // UART-функция для настройки сети и получения логов
    handleWiFi(); // проверка состояния Wi-Fi
    if (WiFi.status() != WL_CONNECTED) return;


    #ifdef SERVER
        static uint32_t lastBroadcast = 0;
        if (millis() - lastBroadcast > 2000) { // каждые 2 секунды
            lastBroadcast = millis();
            sendUDPBroadcast();
        }
        handleServer();
    #endif



#ifdef SERVER
    handleServer();
#else // CLIENT
  static uint32_t lastSendMillis = 0;   // время последней отправки
  static uint16_t seqNum = 1;
  

  if (run_client == 0) return;



  checkClientAlive();
  if (!serverFound) {
      checkUDPBroadcast(); // ищем сервер по UDP
  } else if (!client.connected() && serverFound) {
      client.stop(); // полностью закрываем предыдущий сокет
      delay(100);
      if (client.connect(serverIP, 1234)) {
          Serial.println("Connected to server via auto-discovered IP");
          client_init = 1;
          clearLED();
      } else {
          setRed();
          Serial.println("Failed to connect to server, will retry...");
          client_init = 0;
      }
  }

  if(client_init == 0) return;

  // проверяем, пора ли отправлять сообщение
  if (millis() - lastSendMillis >= sendInterval) {

      lastSendMillis = millis();  // фиксируем время отправки
      String msg = "SEQ:" + String(seqNum) + " TS:" + String(millis()) + " PAYLOAD:Hard!Text^?123";

      if (client.print(msg)) {
          Serial.println(); Serial.print("Sent: "); Serial.println(msg);
          toggleGreen();

          // Сохраняем SEND-запись
          LogEntry entry;
          entry.seq = seqNum;
          entry.ts = lastSendMillis;
          entry.event_type = 1; // SEND
          msg.substring(msg.indexOf("PAYLOAD:") + 8).toCharArray(entry.payload, 16);
          logger.writeLog(entry);

          // Ждем ответ от сервера с таймаутом 10 секунд
          uint32_t startWait = millis();
          bool ackReceived = false;

          while (millis() - startWait < 5000) { // 5 секунд
              if (client.available()) {
                  String resp = client.readStringUntil('\n');
                  resp.trim();
                  if (resp.length() == 0) continue;

                  Serial.print("Received: "); Serial.println(resp);

                  // Сохраняем RECEIVE-запись
                  LogEntry recvEntry;
                  int seqIndex = resp.indexOf("SEQ:");
                  int tsIndex = resp.indexOf("TS:");
                  int payloadIndex = resp.indexOf("PAYLOAD:");

                  if (seqIndex >= 0 && tsIndex > seqIndex && payloadIndex > tsIndex) {
                      String seqStr = resp.substring(seqIndex + 4, tsIndex);
                      String tsStr = resp.substring(tsIndex + 3, payloadIndex);
                      String payloadStr = resp.substring(payloadIndex + 8);

                      seqStr.trim();
                      tsStr.trim();
                      payloadStr.trim();

                      recvEntry.seq = seqStr.toInt();
                      recvEntry.ts = strtoull(tsStr.c_str(), nullptr, 10);
                      recvEntry.event_type = 0; // RECEIVE
                      memset(recvEntry.payload, 0, sizeof(recvEntry.payload));
                      int len = min((int)payloadStr.length(), 16);
                      payloadStr.toCharArray(recvEntry.payload, len + 1);

                      logger.writeLog(recvEntry);
                  }

                  toggleGreen();

                  // Проверяем, что это ответ на текущее сообщение
                  if (resp.startsWith("SEQ:" + String(seqNum))) {
                      ackReceived = true;
                      break;
                  }
              }
              delay(10); // небольшая пауза, чтобы не заблокировать loop
          }

          if (!ackReceived) {
              Serial.println("No response from server for 10s. Reconnecting...");
              client.stop();
              client_init = 0;
              setRed();
              return;
          }

          seqNum++; // увеличиваем только после получения ответа
      } else {
          Serial.println("Error sending");
          client.stop();
          client_init = 0;
          setRed();
      }
  }
#endif
}