#include <ArduinoWebsockets.h>
#include "esp_http_server.h"
#include "esp_timer.h"
#include "esp_camera.h"
#include "camera_index.h"
#include "Arduino.h"
#include "fd_forward.h"
#include "fr_forward.h"
#include "fr_flash.h"

using namespace websockets;
WebsocketsServer socket_server;

const char* ssid = "KOPI";
const char* password = "digoreng123";

#define ENROLL_CONFIRM_TIMES 5
#define FACE_ID_SAVE_NUMBER 5
#define flash_pin 4

#define CAMERA_MODEL_AI_THINKER
#include "camera_pins.h"

camera_fb_t * fb = NULL;

unsigned long recognition_start_time = 0;
bool recognition_active = false;
unsigned long detect_start_time = 0;
bool detect_active = false;

face_id_name_list st_face_list;
static dl_matrix3du_t *aligned_face = NULL;

typedef struct {
  uint8_t *image;
  box_array_t *net_boxes;
  dl_matrix3d_t *face_id;
} http_img_process_result;

static inline mtmn_config_t app_mtmn_config() {
  mtmn_config_t mtmn_config = {0};
  mtmn_config.type = FAST;
  mtmn_config.min_face = 100;
  mtmn_config.pyramid = 0.707;
  mtmn_config.pyramid_times = 4;
  mtmn_config.p_threshold.score = 0.6;
  mtmn_config.p_threshold.nms = 0.7;
  mtmn_config.p_threshold.candidate_number = 20;
  mtmn_config.r_threshold.score = 0.7;
  mtmn_config.r_threshold.nms = 0.7;
  mtmn_config.r_threshold.candidate_number = 10;
  mtmn_config.o_threshold.score = 0.7;
  mtmn_config.o_threshold.nms = 0.7;
  mtmn_config.o_threshold.candidate_number = 1;
  return mtmn_config;
}

mtmn_config_t mtmn_config = app_mtmn_config();

typedef enum {
  START_STREAM,
  START_DETECT,
  SHOW_FACES,
  START_RECOGNITION,
  START_ENROLL,
  ENROLL_COMPLETE,
} en_fsm_state;
en_fsm_state g_state;

typedef struct {
  char enroll_name[ENROLL_NAME_LEN];
} httpd_resp_value;

httpd_resp_value st_name;

WebsocketsClient activeClient;

void app_facenet_main() {
  face_id_name_init(&st_face_list, FACE_ID_SAVE_NUMBER, ENROLL_CONFIRM_TIMES);
  aligned_face = dl_matrix3du_alloc(1, FACE_WIDTH, FACE_HEIGHT, 3);
  read_face_id_from_flash_with_name(&st_face_list);
  print_all_faces();
}

static inline int do_enrollment(face_id_name_list *face_list, dl_matrix3d_t *new_id)
{
  ESP_LOGD(TAG, "START ENROLLING");
  int left_sample_face = enroll_face_id_to_flash_with_name(face_list, new_id, st_name.enroll_name);
  ESP_LOGD(TAG, "Face ID %s Enrollment: Sample %d",
           st_name.enroll_name,
           ENROLL_CONFIRM_TIMES - left_sample_face);
  return left_sample_face;
}

void app_httpserver_init() {
  httpd_config_t config = HTTPD_DEFAULT_CONFIG();
  httpd_handle_t camera_httpd = NULL;
  if (httpd_start(&camera_httpd, &config) == ESP_OK) {
    httpd_uri_t index_uri = {
      .uri       = "/",
      .method    = HTTP_GET,
      .handler   = [](httpd_req_t *req) -> esp_err_t {
        httpd_resp_set_type(req, "text/html");
        httpd_resp_set_hdr(req, "Content-Encoding", "gzip");
        return httpd_resp_send(req, (const char *)index_html_gz, index_html_gz_len);
      },
      .user_ctx  = NULL
    };
    httpd_register_uri_handler(camera_httpd, &index_uri);
  }
}

void send_face_list(WebsocketsClient &client) {
  String names = "";
  face_id_node *head = st_face_list.head;
  for (int i = 0; i < st_face_list.count; i++) {
    names += head->id_name;
    if (i < st_face_list.count - 1) names += ",";
    head = head->next;
  }

  if (client.available()) {
    client.send("facelist:" + names);  // hanya kirim satu kali
  }
}


void send_feedback(String message) {
  if (activeClient.available()) {
    activeClient.send(message);
  } else {
    Serial.println(message);
  }
}

bool search_face_id_in_flash_by_name(face_id_name_list *face_list, const char *name) {
  if (!name || strlen(name) == 0) return false;

  face_id_node *current = face_list->head;
  while (current != NULL) {
    if (strncmp(current->id_name, name, ENROLL_NAME_LEN) == 0) {
      return true;  // ditemukan
    }
    current = current->next;
  }
  return false;  // tidak ditemukan
}

void handle_command(String command) {
  command.trim();

  if (command == "stream") {
    g_state = START_STREAM;
    send_feedback("STREAMING");
  }
  else if (command == "detect") {
    g_state = START_DETECT;
    detect_start_time = millis();
    detect_active = true;
    digitalWrite(flash_pin, HIGH);
    send_feedback("DETECTING");
  }
  else if (command == "flash:on") {
    digitalWrite(flash_pin, HIGH);
    send_feedback("FLASH ON");
  }
  else if (command == "flash:off") {
    digitalWrite(flash_pin, LOW);
    send_feedback("FLASH OFF");
  }
  else if (command.startsWith("capture:")) {
    String person = command.substring(8);
    person.trim();  
    if (person.length() == 0) {
      send_feedback("FACE DATA CORRUPTED: NAME EMPTY");
      Serial.println("ERROR: Tidak boleh kosong. Nama wajah tidak valid.");
      return;
    }
  
    // Cek karakter valid
    bool valid = true;
    for (unsigned int i = 0; i < person.length(); i++) {
      char c = person.charAt(i);
      if (c < 32 || c > 126) {
        valid = false;
        break;
      }
    }
    if (!valid) {
      send_feedback("FACE DATA CORRUPTED: INVALID CHAR");
      Serial.println("ERROR: Nama wajah mengandung karakter tidak valid.");
      return;
    }
  
    if (search_face_id_in_flash_by_name(&st_face_list, person.c_str())) {
      send_feedback("NAME ALREADY EXISTS");
      Serial.println("ERROR: Nama sudah ada.");
      return;
    }

    person.toCharArray(st_name.enroll_name, ENROLL_NAME_LEN - 1);
    st_name.enroll_name[ENROLL_NAME_LEN - 1] = '\0';
  
    g_state = START_ENROLL;
    digitalWrite(flash_pin, HIGH);
    send_feedback("CAPTURING");
  }

  else if (command == "recognise" || command == "DETEKSI") {
    Serial.println("READY");
    g_state = START_RECOGNITION;
    recognition_start_time = millis();
    recognition_active = true;
    digitalWrite(flash_pin, HIGH);
    send_feedback("RECOGNISING");
  }

  else if (command == "getpin") {
    Serial.println("getpin");  
  }

  else if (command.startsWith("newpin:")) {
    String newPin = command.substring(7);
    newPin.trim();
    if (st_face_list.count > 0) {
      send_feedback("PIN TIDAK BISA DIUBAH");
      return;
    }
  
    if (newPin.length() == 6 && newPin.toInt() >= 0) {
      Serial.println("SET_PIN:" + newPin);
      send_feedback("PIN BARU TERKIRIM");
    } else {
      send_feedback("PIN INVALID");
    }
  }

  else if (command.startsWith("remove:")) {
    String person = command.substring(7);
    char person_buf[ENROLL_NAME_LEN * FACE_ID_SAVE_NUMBER];
    strncpy(person_buf, person.c_str(), sizeof(person_buf));
    person_buf[sizeof(person_buf) - 1] = '\0';

    bool result = delete_face_id_in_flash_with_name(&st_face_list, person_buf);
    if (result) {
        send_feedback("remove:ok");
        if (st_face_list.count == 0) {
          Serial.println("resetpin"); 
        }
    } else {
        send_feedback("remove:error");
    }
  }

  else if (command == "delete:allfaces") {
    wipe_all_faces_manual();
    send_face_list(activeClient);
  }

  else if (command == "getfacelist") {
    String names = "";
    face_id_node *node = st_face_list.head;
    while (node != nullptr) {
      names += node->id_name;
      names += ",";
      node = node->next;
    }
  
    Serial.println("facelist:" + names);
    if (activeClient.available()) {
      activeClient.send("facelist:" + names);
    }
  }
}


void handle_message(WebsocketsClient &client, WebsocketsMessage msg) {
  handle_command(msg.data());
}

void handle_serial() {
  if (Serial.available()) {
    String line = Serial.readStringUntil('\n');
    line.trim();

    if (line.startsWith("PIN:")) {
      // Kirim ke WebSocket
      if (activeClient.available()) {
        activeClient.send("pin:" + line.substring(4));  // kirim hanya angkanya
      }
    } else {
      handle_command(line);  // tetap proses perintah lainnya
    }
  }
}

void handle_websocket() {
  if (!activeClient.available()) {
    auto client = socket_server.accept();
    if (client.available()) {
      Serial.println("Web Client Connected");
      activeClient = client;
      activeClient.onMessage(handle_message);
      send_feedback("READY WEBSOCKET");
    }
  }
  if (activeClient.available()) {
    activeClient.poll();
  }
}

void handle_camera() {
  dl_matrix3du_t *image_matrix = dl_matrix3du_alloc(1, 320, 240, 3);
  http_img_process_result out_res = {0};
  out_res.image = image_matrix->item;

  fb = esp_camera_fb_get();
  if (!fb) {
    Serial.println("Camera capture failed");
    dl_matrix3du_free(image_matrix);
    return;
  }

  if (g_state == START_ENROLL)
  {
    fmt2rgb888(fb->buf, fb->len, fb->format, out_res.image);
    out_res.net_boxes = face_detect(image_matrix, &mtmn_config);
  
    if (out_res.net_boxes)
    {
    if (align_face(out_res.net_boxes, image_matrix, aligned_face) == ESP_OK)
    {
      out_res.face_id = get_face_id(aligned_face);
    
      if (out_res.face_id && out_res.face_id->item) {
        bool valid_face = false;
        for (int i = 0; i < FACE_ID_SIZE; i++) {
          float val = out_res.face_id->item[i];
          if (!isnan(val) && fabs(val) > 1e-5) {
            valid_face = true;
            break;
          }
        }
    
        if (valid_face) {
        static int last_enroll_count = -1;
        int left_sample_face = do_enrollment(&st_face_list, out_res.face_id);
        
        if (left_sample_face != last_enroll_count) {
          last_enroll_count = left_sample_face;
          char enrolling_message[64];
          sprintf(enrolling_message, "SAMPLE NUMBER %d FOR %s", ENROLL_CONFIRM_TIMES - left_sample_face, st_name.enroll_name);
          send_feedback(enrolling_message);
        }
            
        if (left_sample_face == 0) {
          g_state = START_STREAM;
          last_enroll_count = -1;
          digitalWrite(flash_pin, LOW);
        
          // Reload daftar wajah untuk cek jumlah
          read_face_id_from_flash_with_name(&st_face_list);
        
          send_feedback("capture:done"); // beri tahu website bahwa capture selesai
          send_face_list(activeClient);
        
          if (st_face_list.count == 1) {
            Serial.println("Waiting for newpin...");
          }
          
        }

        } else {
          send_feedback("FACE DATA INVALID, PLEASE RETRY");
          Serial.println("Face ID tidak valid: seluruh vektor nol/NaN");
        }
    
        dl_matrix3d_free(out_res.face_id);
      } else {
        send_feedback("FAILED TO GET FACE ID");
        Serial.println("Face ID pointer null atau kosong");
      }
    }
    else
    {
      send_feedback("FACE ALIGNMENT FAILED");
      Serial.println("align_face gagal");
    }
   }
  
    if (out_res.net_boxes)
      free(out_res.net_boxes);
  }
  
  if (g_state == START_RECOGNITION && recognition_active) {
    fmt2rgb888(fb->buf, fb->len, fb->format, out_res.image);
    out_res.net_boxes = face_detect(image_matrix, &mtmn_config);
  
    if (out_res.net_boxes) {
      if (align_face(out_res.net_boxes, image_matrix, aligned_face) == ESP_OK) {
        out_res.face_id = get_face_id(aligned_face);
  
        if (out_res.face_id && st_face_list.count > 0) {
          face_id_node *f = recognize_face_with_name(&st_face_list, out_res.face_id);
  
          if (f) {
            Serial.printf("DIKENALI: %s\n", f->id_name);
            char recognised_message[64];
            sprintf(recognised_message, "DOOR OPEN FOR %s", f->id_name);
            send_feedback(recognised_message);
            g_state = START_STREAM;
            recognition_active = false;
            digitalWrite(flash_pin, LOW);
          } else if (millis() - recognition_start_time >= 5000) {
            Serial.println("TAKKENAL");
            send_feedback("FACE NOT RECOGNISED");
            g_state = START_STREAM;
            recognition_active = false;
            digitalWrite(flash_pin, LOW);
          }
        } else if (millis() - recognition_start_time >= 5000) {
          Serial.println("TAKKENAL");
          send_feedback("FACE NOT RECOGNISED");
          g_state = START_STREAM;
          recognition_active = false;
          digitalWrite(flash_pin, LOW);
        }
  
        if (out_res.face_id) {
          dl_matrix3d_free(out_res.face_id);
        }
      } else if (millis() - recognition_start_time >= 5000) {
        Serial.println("TAKADA");
        send_feedback("NO FACE DETECTED");
        g_state = START_STREAM;
        recognition_active = false;
        digitalWrite(flash_pin, LOW);
      }
  
      free(out_res.net_boxes);
    } else if (millis() - recognition_start_time >= 5000) {
      Serial.println("TAKADA");
      send_feedback("NO FACE DETECTED");
      g_state = START_STREAM;
      recognition_active = false;
      digitalWrite(flash_pin, LOW);
    }
  }

  if (g_state == START_DETECT && detect_active) {
    fmt2rgb888(fb->buf, fb->len, fb->format, out_res.image);
    out_res.net_boxes = face_detect(image_matrix, &mtmn_config);
  
    if (out_res.net_boxes) {
      Serial.println("Wajah Terdeteksi");
      send_feedback("FACE DETECTED");
      g_state = START_STREAM;
      detect_active = false;
      digitalWrite(flash_pin, LOW);
    } 
    else {
      if (millis() - detect_start_time >= 5000) {
        Serial.println("Tidak ada wajah");
        send_feedback("FACE NOT RECOGNISED");
        g_state = START_STREAM;
        detect_active = false;
        digitalWrite(flash_pin, LOW);
      }
    }
  }
  
  if (activeClient.available()) {
      activeClient.sendBinary((const char *)fb->buf, fb->len);
  }

  esp_camera_fb_return(fb);
  fb = NULL;
  dl_matrix3du_free(image_matrix);
}

void print_all_faces() {
  face_id_node *head = st_face_list.head;
  int index = 0;

  while (head) {
    Serial.printf("[%d] Nama: ", index++);
    
    // Tampilkan karakter satu per satu secara heksadesimal dan normal
    for (int i = 0; i < ENROLL_NAME_LEN; i++) {
      char c = head->id_name[i];
      if (c == '\0') break;
      if (c < 32 || c > 126) {
        Serial.printf("[0x%02X]", c);  // karakter non-printable
      } else {
        Serial.print(c);  // karakter biasa
      }
    }
    Serial.println();
    head = head->next;
  }

  if (index == 0) {
    Serial.println("Tidak ada wajah yang tersimpan.");
  }
}

void wipe_all_faces_manual() {
  face_id_node *f = st_face_list.head;
  while (f != NULL) {
    delete_face_id_in_flash_with_name(&st_face_list, f->id_name);
    f = f->next;
  }
  face_id_name_init(&st_face_list, FACE_ID_SAVE_NUMBER, ENROLL_CONFIRM_TIMES);
  Serial.println("Semua data wajah dihapus manual.");
}

void remove_corrupted_faces() {
  int corrupt_count = 0;
  bool found_corrupt;

  do {
    found_corrupt = false;

    face_id_node *current = st_face_list.head;
    while (current != NULL) {
      bool corrupt = true;

      // Cek semua byte = 0xFF
      for (int i = 0; i < ENROLL_NAME_LEN; i++) {
        if ((uint8_t)current->id_name[i] != 0xFF) {
          corrupt = false;
          break;
        }
      }

      // Cek juga karakter ASCII printable
      if (!corrupt) {
        for (int i = 0; i < ENROLL_NAME_LEN; i++) {
          char c = current->id_name[i];
          if (c == '\0') break;
          if (c < 32 || c > 126) {
            corrupt = true;
            break;
          }
        }
      }

      if (corrupt) {
        Serial.print("Menghapus wajah korup: ");
        for (int i = 0; i < ENROLL_NAME_LEN; i++) {
          Serial.printf("[0x%02X]", (uint8_t)current->id_name[i]);
        }
        Serial.println();

        delete_face_id_in_flash_with_name(&st_face_list, current->id_name);
        corrupt_count++;

        // Reload list wajah dan mulai ulang dari head
        read_face_id_from_flash_with_name(&st_face_list);
        found_corrupt = true;
        break;
      }

      current = current->next;
    }

  } while (found_corrupt);

  Serial.printf("Total wajah korup yang dihapus: %d\n", corrupt_count);
}



void setup() {
  Serial.begin(9600);
  Serial.setDebugOutput(true);

  camera_config_t config;
  config.ledc_channel = LEDC_CHANNEL_0;
  config.ledc_timer = LEDC_TIMER_0;
  config.pin_d0 = Y2_GPIO_NUM;
  config.pin_d1 = Y3_GPIO_NUM;
  config.pin_d2 = Y4_GPIO_NUM;
  config.pin_d3 = Y5_GPIO_NUM;
  config.pin_d4 = Y6_GPIO_NUM;
  config.pin_d5 = Y7_GPIO_NUM;
  config.pin_d6 = Y8_GPIO_NUM;
  config.pin_d7 = Y9_GPIO_NUM;
  config.pin_xclk = XCLK_GPIO_NUM;
  config.pin_pclk = PCLK_GPIO_NUM;
  config.pin_vsync = VSYNC_GPIO_NUM;
  config.pin_href = HREF_GPIO_NUM;
  config.pin_sscb_sda = SIOD_GPIO_NUM;
  config.pin_sscb_scl = SIOC_GPIO_NUM;
  config.pin_pwdn = PWDN_GPIO_NUM;
  config.pin_reset = RESET_GPIO_NUM;
  config.xclk_freq_hz = 20000000;
  config.pixel_format = PIXFORMAT_JPEG;

  if (psramFound()) {
    config.frame_size = FRAMESIZE_UXGA;
    config.jpeg_quality = 10;
    config.fb_count = 2;
  } else {
    config.frame_size = FRAMESIZE_SVGA;
    config.jpeg_quality = 12;
    config.fb_count = 1;
  }

  esp_err_t err = esp_camera_init(&config);
  if (err != ESP_OK) {
    Serial.printf("Camera init failed with error 0x%x", err);
    return;
  }

  sensor_t * s = esp_camera_sensor_get();
  s->set_framesize(s, FRAMESIZE_QVGA);
  s->set_vflip(s, 1);
  s->set_hmirror(s, 1);
  pinMode(flash_pin, OUTPUT);
  digitalWrite(flash_pin, LOW);

  WiFi.begin(ssid, password);
  unsigned long startAttemptTime = millis();
  const unsigned long wifiTimeout = 10000;

  while (WiFi.status() != WL_CONNECTED && millis() - startAttemptTime < wifiTimeout) {
    delay(500);
    Serial.println("MENGHUBUNGKAN WIFI...");
  }

  if (WiFi.status() == WL_CONNECTED) {
    Serial.println("READY");
    Serial.println(WiFi.localIP());
  } else {
    Serial.println("WIFI NOT CONNECTED");
  }

  //  cek apakah ada data yang tidak valid
  read_face_id_from_flash_with_name(&st_face_list);
  remove_corrupted_faces();

  if (st_face_list.count == 0) {
    delay(1200);
    Serial.println("resetpin"); 
  }

  app_httpserver_init();
  app_facenet_main();
  socket_server.listen(82);
  g_state = START_STREAM;

}

void loop() {
  handle_serial();
  handle_camera();
  handle_websocket();
}