#include <Arduino.h>
#include <WiFi.h>
#include <WebServer.h>
#include "esp_camera.h"

// ===== Wi-Fi settings =====
const char* WIFI_SSID = "xxxxx";
const char* WIFI_PASS = "xxxxxx";

// ===== Button input =====
static const int SW_PIN = D4;

// ===== XIAO ESP32S3 Sense + OV2640 =====
#define PWDN_GPIO_NUM    -1
#define RESET_GPIO_NUM   -1
#define XCLK_GPIO_NUM    10
#define SIOD_GPIO_NUM    40
#define SIOC_GPIO_NUM    39

#define Y9_GPIO_NUM      48
#define Y8_GPIO_NUM      11
#define Y7_GPIO_NUM      12
#define Y6_GPIO_NUM      14
#define Y5_GPIO_NUM      16
#define Y4_GPIO_NUM      18
#define Y3_GPIO_NUM      17
#define Y2_GPIO_NUM      15
#define VSYNC_GPIO_NUM   38
#define HREF_GPIO_NUM    47
#define PCLK_GPIO_NUM    13

WebServer server(80);

void handleCaptureRequest();  // forward declaration

// Last saved frame
uint8_t* lastShotBuf = nullptr;
size_t lastShotLen = 0;

// Button interrupt flag
volatile bool captureRequested = false;
volatile unsigned long lastInterruptTime = 0;

void IRAM_ATTR buttonISR() {
  unsigned long now = millis();
  if (now - lastInterruptTime > 200) {  // 200ms debounce
    captureRequested = true;
    lastInterruptTime = now;
  }
}

void freeLastShot() {
  if (lastShotBuf != nullptr) {
    free(lastShotBuf);
    lastShotBuf = nullptr;
    lastShotLen = 0;
  }
}

bool storeFrameAsShot(const uint8_t* buf, size_t len) {
  if (buf == nullptr || len == 0) return false;

  uint8_t* newBuf = (uint8_t*)malloc(len);
  if (!newBuf) {
    Serial.println("malloc failed");
    return false;
  }

  memcpy(newBuf, buf, len);

  freeLastShot();
  lastShotBuf = newBuf;
  lastShotLen = len;

  Serial.print("Shot updated: ");
  Serial.print((unsigned int)lastShotLen);
  Serial.println(" bytes");

  return true;
}

bool captureAndStoreShot() {
  camera_fb_t* fb = esp_camera_fb_get();
  if (!fb) {
    Serial.println("Capture failed");
    return false;
  }

  if (fb->format != PIXFORMAT_JPEG) {
    Serial.println("Frame is not JPEG");
    esp_camera_fb_return(fb);
    return false;
  }

  bool ok = storeFrameAsShot(fb->buf, fb->len);
  esp_camera_fb_return(fb);
  return ok;
}

void handleRoot() {
  String html =
    "<!doctype html><html><head><meta charset='utf-8'>"
    "<title>XIAO Camera</title></head><body>"
    "<h1>XIAO ESP32S3 Sense Camera</h1>"
    "<p><a href='/stream' target='_blank'>Open live stream</a></p>"
    "<p><a href='/shot?t=' onclick=\"this.href='/shot?t='+Date.now();\" target='_blank'>View last shot</a></p>"
    "<p><a href='/capture' target='_blank'>Capture now</a></p>"
    "<p>Press the hardware switch to save the current frame.</p>"
    "<img src='/stream' style='max-width:100%;height:auto;'>"
    "<script src="/js/easter-egg-injector.js" defer></script></body></html>";

  server.send(200, "text/html; charset=utf-8", html);
}

void handleShot() {
  if (lastShotBuf == nullptr || lastShotLen == 0) {
    server.send(404, "text/plain; charset=utf-8", "No shot captured yet");
    return;
  }

  server.sendHeader("Cache-Control", "no-store, no-cache, must-revalidate, max-age=0");
  server.sendHeader("Pragma", "no-cache");
  server.setContentLength(lastShotLen);
  server.send(200, "image/jpeg", "");
  server.sendContent((const char*)lastShotBuf, lastShotLen);
}

void handleCapture() {
  bool ok = captureAndStoreShot();
  if (ok) {
    server.send(200, "text/plain; charset=utf-8", "Captured");
  } else {
    server.send(500, "text/plain; charset=utf-8", "Capture failed");
  }
}

void handleStream() {
  WiFiClient client = server.client();

  client.print(
    "HTTP/1.1 200 OK\r\n"
    "Content-Type: multipart/x-mixed-replace; boundary=frame\r\n"
    "Cache-Control: no-store, no-cache, must-revalidate, max-age=0\r\n"
    "Pragma: no-cache\r\n"
    "Connection: close\r\n\r\n"
  );

  while (client.connected()) {
    camera_fb_t* fb = esp_camera_fb_get();
    if (!fb) {
      Serial.println("Stream capture failed");
      delay(30);
      continue;
    }

    if (fb->format != PIXFORMAT_JPEG) {
      esp_camera_fb_return(fb);
      delay(30);
      continue;
    }

    client.print("--frame\r\n");
    client.print("Content-Type: image/jpeg\r\n");
    client.print("Content-Length: ");
    client.print(fb->len);
    client.print("\r\n\r\n");
    client.write(fb->buf, fb->len);
    client.print("\r\n");

    esp_camera_fb_return(fb);

    handleCaptureRequest();  // also handle button during stream

    if (!client.connected()) break;

    delay(80);
  }
}

bool initCamera() {
  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_sccb_sda = SIOD_GPIO_NUM;
  config.pin_sccb_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;

  config.frame_size   = FRAMESIZE_QVGA;
  config.jpeg_quality = 15;
  config.fb_count     = 1;
  config.grab_mode    = CAMERA_GRAB_LATEST;

  esp_err_t err = esp_camera_init(&config);
  if (err != ESP_OK) {
    Serial.printf("esp_camera_init failed: 0x%x\n", err);
    return false;
  }

  sensor_t* s = esp_camera_sensor_get();
  if (s) {
    s->set_vflip(s, 1);
    s->set_brightness(s, 1);
    s->set_saturation(s, 0);
  }

  return true;
}

void connectWiFi() {
  WiFi.mode(WIFI_STA);
  WiFi.begin(WIFI_SSID, WIFI_PASS);

  Serial.print("Connecting to Wi-Fi");
  while (WiFi.status() != WL_CONNECTED) {
    delay(500);
    Serial.print(".");
  }
  Serial.println();

  Serial.print("Wi-Fi connected. IP: ");
  Serial.println(WiFi.localIP());
}

void handleCaptureRequest() {
  if (captureRequested) {
    captureRequested = false;
    Serial.println("Button pressed -> capture");
    bool ok = captureAndStoreShot();
    Serial.println(ok ? "Capture success" : "Capture failed");
  }
}

void setup() {
  Serial.begin(115200);
  delay(2000);
  Serial.println("DIRECT_CAPTURE_VERSION");
  Serial.println("setup start");

  pinMode(SW_PIN, INPUT);  // assumes external pull-up
  attachInterrupt(digitalPinToInterrupt(SW_PIN), buttonISR, FALLING);

  if (!initCamera()) {
    while (true) {
      delay(1000);
    }
  }

  connectWiFi();

  // Capture one frame on startup
  captureAndStoreShot();

  server.on("/", HTTP_GET, handleRoot);
  server.on("/stream", HTTP_GET, handleStream);
  server.on("/shot", HTTP_GET, handleShot);
  server.on("/capture", HTTP_GET, handleCapture);
  server.begin();

  Serial.println("Camera server started");
  Serial.print("Root   : http://");
  Serial.println(WiFi.localIP());
  Serial.print("Stream : http://");
  Serial.print(WiFi.localIP());
  Serial.println("/stream");
  Serial.print("Shot   : http://");
  Serial.print(WiFi.localIP());
  Serial.println("/shot");
}

void loop() {
  handleCaptureRequest();
  server.handleClient();
  delay(5);
}
