#include <WiFi.h>
#include <PubSubClient.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <U8g2_for_Adafruit_GFX.h>
#include <ArduinoJson.h>
#include <Adafruit_NeoPixel.h>
#include <DHT.h>

int BUTTON_ONE = 15;
int BUTTON_TWO = 4;
int SPEAKER_OUT = 23;
int LED_DATA = 5;
int LED_COUNT = 16;

int SOUND_HZ = 2000;
int SOUND_CH = 8;

#define W 128
#define H 64
#define OLED_NO_PIN -1
Adafruit_SSD1306 screen(W, H, &Wire, OLED_NO_PIN);
U8G2_FOR_ADAFRUIT_GFX text_gfx;

#define DHT_P 18
#define DHT_MODEL DHT22
DHT env_sensor(DHT_P, DHT_MODEL);

Adafruit_NeoPixel lights(LED_COUNT, LED_DATA, NEO_GRB + NEO_KHZ800);

const char* NET_NAME = "Wokwi-GUEST";
const char* NET_PASS = "";
const char* MQTT_SERVER = "test.mosquitto.org";
const uint16_t MQTT_PORT_NUM = 1883;

WiFiClient net_client;
PubSubClient mqtt_client(net_client);

String my_id = "";
String topic_a, topic_b, topic_c, topic_d, topic_e, topic_f;

volatile uint32_t last_num = 0;
volatile uint32_t current_num = 0;

int count_waiting() {
  long diff = (long)last_num - (long)current_num;
  return (diff < 0) ? 0 : (int)diff;
}

bool check_button_one() {
  static int prev = HIGH;
  static uint32_t last_time = 0;
  int now = digitalRead(BUTTON_ONE);
  uint32_t time_now = millis();
  if (now != prev && (time_now - last_time) > 30) {
    prev = now;
    last_time = time_now;
    if (prev == LOW) return true;
  }
  return false;
}

bool check_button_two() {
  static int prev = HIGH;
  static uint32_t last_time = 0;
  int now = digitalRead(BUTTON_TWO);
  uint32_t time_now = millis();
  if (now != prev && (time_now - last_time) > 30) {
    prev = now;
    last_time = time_now;
    if (prev == LOW) return true;
  }
  return false;
}

void make_sound(uint16_t length = 90, uint16_t gap = 0) {
  ledcWriteTone(SPEAKER_OUT, SOUND_HZ);
  delay(length);
  ledcWriteTone(SPEAKER_OUT, 0);
  if (gap) delay(gap);
}

float smooth = 0.25f;
float arrive_rate = 0.1f;
float serve_rate = 2.0f;
bool auto_rate = true;
const int MAX_TIMES = 8;
float time_list[MAX_TIMES];
int time_count = 0;
uint32_t last_serve_time = 0;
uint32_t arrived_since = 0;
uint32_t last_check_time = 0;
const uint32_t CHECK_MS = 5000;
float wait_time = 0.0f;
float use_level = 0.0f;
String status_col = "green";
float air_temp = NAN;
float air_wet = NAN;
long net_strength = 0;

void update_screen() {
  screen.clearDisplay();
  text_gfx.setCursor(0, 18);
  text_gfx.setFont(u8g2_font_logisoso18_tr);
  text_gfx.print("Now: ");
  text_gfx.printf("%03u", current_num);

  text_gfx.setCursor(0, 40);
  text_gfx.setFont(u8g2_font_profont22_tr);
  text_gfx.print("Wait: ");
  text_gfx.printf("%02d", count_waiting());

  text_gfx.setCursor(0, 60);
  text_gfx.setFont(u8g2_font_profont15_tr);
  text_gfx.printf("Time: %.1fm  %s", wait_time, status_col.c_str());

  screen.display();
}

void set_lights(float val) {
  val = constrain(val, 0.0f, 1.0f);
  uint8_t red=0,green=0,blue=0;
  if (val < 0.6f) { green = 255 * (0.2f + val*0.8f); }
  else if (val < 0.85f) { red = 200; green = 200; }
  else { red = 255; }
  for (int i=0;i<LED_COUNT;i++) lights.setPixelColor(i, lights.Color(red,green,blue));
  lights.show();
}

String make_id(const uint8_t* mac) {
  char buf[20];
  sprintf(buf, "%02x%02x%02x%02x%02x%02x", mac[0],mac[1],mac[2],mac[3],mac[4],mac[5]);
  return String(buf);
}

void send_status(const char* stat="online") {
  StaticJsonDocument<256> data;
  data["stat"] = stat;
  data["ver"] = "lr20-p1-1.0.2";
  data["signal"] = net_strength;
  String output;
  serializeJson(data, output);
  mqtt_client.publish(topic_a.c_str(), output.c_str(), true);
}

void send_data() {
  StaticJsonDocument<512> data;
  data["time"] = (uint32_t)(millis()/1000);
  data["queue"] = count_waiting();
  data["last"] = last_num;
  data["now"] = current_num;
  data["wait_m"] = wait_time;
  data["arrive_m"] = arrive_rate;
  data["serve_m"] = serve_rate;
  data["busy"] = use_level;
  data["color"] = status_col;
  if (!isnan(air_temp) && !isnan(air_wet)) {
    JsonObject env = data.createNestedObject("air");
    env["temp"] = air_temp;
    env["wet"] = air_wet;
  }
  data["wifi"] = net_strength;
  String output;
  serializeJson(data, output);
  mqtt_client.publish(topic_b.c_str(), output.c_str(), false);
}

void send_ai_request() {
  StaticJsonDocument<768> data;
  data["kind"] = "digest_prompt";
  data["language"] = "en";
  data["place"] = "small_business_queue";
  data["desc"] = "Summarize the last period and suggest 3 actionable tips.";
  JsonObject numbers = data.createNestedObject("numbers");
  numbers["arrive_m"] = arrive_rate;
  numbers["serve_m"] = serve_rate;
  numbers["queue"] = count_waiting();
  numbers["now"] = current_num;
  numbers["last"] = last_num;
  numbers["wait_m"] = wait_time;
  numbers["busy"] = use_level;
  numbers["color"] = status_col;
  data["request"] = "Return JSON with {brief, staffing, schedule, risk}. Keep it concise.";
  String output;
  serializeJson(data, output);
  mqtt_client.publish(topic_d.c_str(), output.c_str(), false);
}

void got_mqtt(char* topic, byte* msg, unsigned int len) {
  StaticJsonDocument<512> data;
  DeserializationError err = deserializeJson(data, msg, len);
  if (err) return;

  if (data.containsKey("next")) {
    uint32_t n = data["next"];
    while (n--) {
      current_num++;
      make_sound();
      uint32_t ms = millis();
      if (last_serve_time != 0) {
        float diff = (ms - last_serve_time) / 1000.0f;
        if (time_count < MAX_TIMES) time_list[time_count++] = diff;
        else { for (int i=1;i<MAX_TIMES;i++) time_list[i-1]=time_list[i]; time_list[MAX_TIMES-1] = diff; }
      }
      last_serve_time = ms;
    }
  }
  if (data.containsKey("ticket")) {
    uint32_t n = data["ticket"];
    arrived_since += n;
    last_num += n;
  }
  if (data.containsKey("beep")) {
    make_sound();
  }
  if (data.containsKey("set")) {
    JsonObject s = data["set"];
    if (s.containsKey("alpha")) {
      float a = s["alpha"];
      if (a>0 && a<1) smooth = a;
    }
    if (s.containsKey("mu")) {
      serve_rate = max(0.1f, (float)s["mu"]);
    }
    if (s.containsKey("mu_auto")) {
      auto_rate = s["mu_auto"];
    }
  }
  if (data.containsKey("pixels")) {
    JsonObject p = data["pixels"];
    String mode = p["mode"] | "bar";
    if (mode == "off") {
      for (int i=0;i<LED_COUNT;i++) lights.setPixelColor(i, 0);
      lights.show();
    } else if (p.containsKey("level")) {
      set_lights((float)p["level"]);
    }
  }
  if (data.containsKey("digest")) {
    send_ai_request();
  }
}

void connect_wifi() {
  if (WiFi.status() == WL_CONNECTED) return;
  WiFi.mode(WIFI_STA);
  WiFi.begin(NET_NAME, NET_PASS);
  uint32_t start = millis();
  while (WiFi.status() != WL_CONNECTED && (millis()-start)<15000) {
    delay(250);
  }
}

void connect_mqtt() {
  if (mqtt_client.connected()) return;
  while (!mqtt_client.connected()) {
    String client_name = my_id + "-" + String((uint32_t)random(0xFFFF), HEX);
    if (mqtt_client.connect(client_name.c_str())) {
      mqtt_client.subscribe(topic_c.c_str());
      send_status("online");
    } else {
      delay(1000);
    }
  }
}

float average(float *arr, int n) {
  if (n<=0) return NAN;
  double sum=0;
  for (int i=0;i<n;i++) sum += arr[i];
  return (float)(sum/n);
}

void update_serve_rate() {
  if (!auto_rate) return;
  float avg = average(time_list, time_count);
  if (!isnan(avg) && avg>0.5f) {
    serve_rate = 60.0f / avg;
  }
}

void update_arrive_rate() {
  uint32_t now = millis();
  if (last_check_time == 0) { last_check_time = now; return; }
  uint32_t diff = now - last_check_time;
  if (diff >= CHECK_MS) {
    float instant = (arrived_since * 60000.0f) / diff;
    arrive_rate = smooth * instant + (1.0f - smooth) * arrive_rate;
    arrived_since = 0;
    last_check_time = now;
  }
}

void update_calc() {
  float serve = max(0.1f, serve_rate);
  float arrive = max(0.0f, arrive_rate);
  use_level = min(0.99f, arrive / serve);

  int waiting = count_waiting();
  wait_time = (waiting > 0) ? (waiting / serve) : 0.0f;

  if (use_level < 0.6f) status_col = "green";
  else if (use_level < 0.85f) status_col = "yellow";
  else status_col = "red";

  set_lights(use_level);
}

uint32_t last_ui = 0;
uint32_t last_data = 0;
uint32_t last_ai = 0;

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

  pinMode(BUTTON_ONE, INPUT_PULLUP);
  pinMode(BUTTON_TWO, INPUT_PULLUP);
  pinMode(SPEAKER_OUT, OUTPUT);

  ledcAttach(SPEAKER_OUT, SOUND_HZ, SOUND_CH);

  lights.begin();
  lights.clear();
  lights.show();

  Wire.begin();
  if (!screen.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("OLED fail"));
    for(;;) delay(1000);
  }
  text_gfx.begin(screen);
  screen.clearDisplay();
  screen.display();

  env_sensor.begin();

  uint8_t mac[6];
  WiFi.macAddress(mac);
  char mac_string[20];
  sprintf(mac_string, "%02x%02x%02x%02x%02x%02x", mac[0],mac[1],mac[2],mac[3],mac[4],mac[5]);
  my_id = String("esp32-queue-") + mac_string;

  topic_a = "sb/" + my_id + "/state";
  topic_b = "sb/" + my_id + "/telemetry";
  topic_c = "sb/" + my_id + "/cmd";
  topic_d = "sb/" + my_id + "/llm_prompt";
  topic_e = "sb/" + my_id + "/schedule";
  topic_f = "sb/" + my_id + "/control";

  mqtt_client.setServer(MQTT_SERVER, MQTT_PORT_NUM);
  mqtt_client.setCallback(got_mqtt);

  connect_wifi();
  connect_mqtt();

  text_gfx.setFont(u8g2_font_profont22_tr);
  text_gfx.setCursor(0, 24);
  text_gfx.print("Queue System");
  text_gfx.setCursor(0, 46);
  text_gfx.print(my_id);
  screen.display();
  delay(900);

  last_ui = last_data = last_ai = millis();
}

void loop() {
  connect_wifi();
  connect_mqtt();
  mqtt_client.loop();

  if (check_button_one()) {
    last_num++;
    arrived_since++;
  }
  if (check_button_two()) {
    current_num++;
    make_sound();
    uint32_t ms = millis();
    if (last_serve_time != 0) {
      float diff = (ms - last_serve_time) / 1000.0f;
      if (time_count < MAX_TIMES) time_list[time_count++] = diff;
      else { for (int i=1;i<MAX_TIMES;i++) time_list[i-1]=time_list[i]; time_list[MAX_TIMES-1] = diff; }
    }
    last_serve_time = ms;
  }

  static uint32_t last_env = 0;
  if (millis() - last_env > 3000) {
    air_temp = env_sensor.readTemperature();
    air_wet  = env_sensor.readHumidity();
    last_env = millis();
  }

  update_arrive_rate();
  update_serve_rate();
  update_calc();

  if (millis() - last_ui > 150) {
    update_screen();
    last_ui = millis();
  }

  if (millis() - last_data > 5000) {
    net_strength = WiFi.RSSI();
    send_data();
    last_data = millis();
  }

  if (millis() - last_ai > 60000) {
    send_ai_request();
    last_ai = millis();
  }
}