#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include "esp_camera.h"

// ===== Camera model & pins (XIAO ESP32S3 Sense) =====
#define CAMERA_MODEL_XIAO_ESP32S3
#include "camera_pins.h"

// ===== OLED =====
#define SCREEN_WIDTH   128
#define SCREEN_HEIGHT  64
#define SCREEN_ADDRESS 0x3C
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// ===== Status LED (board LED) =====
#define PIN_LED 21

// ===== Capture timing =====
const unsigned long CAPTURE_INTERVAL_MS = 60000UL;  // 1 minute
unsigned long lastCapture = 0;

bool camera_ready = false;

// A small working buffer for scaled grayscale (float for FS dithering)
static float gray_buffer[SCREEN_WIDTH * SCREEN_HEIGHT];

// --------- Draw helpers ---------
static void drawTextCentered(const char* line1, const char* line2 = nullptr) {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  int16_t x1, y1; uint16_t w, h;
  if (line1) {
    display.getTextBounds(line1, 0, 0, &x1, &y1, &w, &h);
    display.setCursor((SCREEN_WIDTH - w)/2, (SCREEN_HEIGHT/2) - (line2 ? 12 : h/2));
    display.println(line1);
  }
  if (line2) {
    display.getTextBounds(line2, 0, 0, &x1, &y1, &w, &h);
    display.setCursor((SCREEN_WIDTH - w)/2, (SCREEN_HEIGHT/2) + 2);
    display.println(line2);
  }
  display.display();
}

// --------- Image processing: scale -> contrast stretch -> Floyd–Steinberg dither -> draw ---------
static void processAndDisplayImage(const camera_fb_t* fb) {
  // Expect PIXFORMAT_GRAYSCALE
  const uint8_t* src = fb->buf;
  const int sw = fb->width;
  const int sh = fb->height;

  // 1) Downsample to 128×64 via simple box averaging
  const float sx = (float)sw / SCREEN_WIDTH;
  const float sy = (float)sh / SCREEN_HEIGHT;

  for (int y = 0; y < SCREEN_HEIGHT; ++y) {
    int start_y = (int)(y * sy);
    int end_y   = (int)((y + 1) * sy); if (end_y <= start_y) end_y = start_y + 1;

    for (int x = 0; x < SCREEN_WIDTH; ++x) {
      int start_x = (int)(x * sx);
      int end_x   = (int)((x + 1) * sx); if (end_x <= start_x) end_x = start_x + 1;

      uint32_t sum = 0;
      int count = 0;
      for (int yy = start_y; yy < end_y; ++yy) {
        const int row = yy * sw;
        for (int xx = start_x; xx < end_x; ++xx) {
          sum += src[row + xx];
          ++count;
        }
      }
      gray_buffer[y * SCREEN_WIDTH + x] = (count ? (float)sum / count : 0.0f);
    }
  }

  // 2) Contrast stretch (linear)
  float gmin = 255.0f, gmax = 0.0f;
  const int N = SCREEN_WIDTH * SCREEN_HEIGHT;
  for (int i = 0; i < N; ++i) {
    float v = gray_buffer[i];
    if (v < gmin) gmin = v;
    if (v > gmax) gmax = v;
  }
  float range = gmax - gmin;
  if (range < 1.0f) range = 1.0f;
  for (int i = 0; i < N; ++i) {
    gray_buffer[i] = (gray_buffer[i] - gmin) * (255.0f / range);
  }

  // 3) Floyd–Steinberg dithering on the working buffer
  for (int y = 0; y < SCREEN_HEIGHT; ++y) {
    for (int x = 0; x < SCREEN_WIDTH; ++x) {
      int idx = y * SCREEN_WIDTH + x;
      float oldp = gray_buffer[idx];
      float newp = (oldp >= 128.0f) ? 255.0f : 0.0f;
      gray_buffer[idx] = newp;
      float err = oldp - newp;

      if (x + 1 < SCREEN_WIDTH)                         gray_buffer[idx + 1]                 += err * 7.0f / 16.0f;
      if (y + 1 < SCREEN_HEIGHT) {
        if (x > 0)                                      gray_buffer[idx + SCREEN_WIDTH - 1]  += err * 3.0f / 16.0f;
                                                        gray_buffer[idx + SCREEN_WIDTH]      += err * 5.0f / 16.0f;
        if (x + 1 < SCREEN_WIDTH)                       gray_buffer[idx + SCREEN_WIDTH + 1]  += err * 1.0f / 16.0f;
      }
    }
  }

  // 4) Draw to OLED
  display.clearDisplay();
  for (int y = 0; y < SCREEN_HEIGHT; ++y) {
    for (int x = 0; x < SCREEN_WIDTH; ++x) {
      if (gray_buffer[y * SCREEN_WIDTH + x] > 0.5f) {
        display.drawPixel(x, y, SSD1306_WHITE);
      }
    }
  }
  display.display();
}

// ===== Setup =====
void setup() {
  Serial.begin(115200);
  delay(250);

  pinMode(PIN_LED, OUTPUT);
  digitalWrite(PIN_LED, HIGH);

  // OLED init
  Wire.begin();
  delay(50);
  if (!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
    Serial.println("SSD1306 init failed");
  } else {
    drawTextCentered("Camera test", "1-min auto capture");
  }

  // Camera init (PSRAM frame buffers)
  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;

  // Small-ish frame is fine for 128×64 preview; you can bump to QVGA later
  config.frame_size   = FRAMESIZE_QQVGA;      // 160×120
  config.pixel_format = PIXFORMAT_GRAYSCALE;  // direct grayscale → easy preview
  config.fb_location  = CAMERA_FB_IN_PSRAM;   // use PSRAM
  config.jpeg_quality = 12;
  config.fb_count     = 1;
  config.grab_mode    = CAMERA_GRAB_LATEST;

  esp_err_t err = esp_camera_init(&config);
  if (err == ESP_OK) {
    camera_ready = true;
    sensor_t *s = esp_camera_sensor_get();
    // Optional orientation tweaks
    s->set_vflip(s, 1);
    s->set_hmirror(s, 0);

    Serial.println("Camera ready (PSRAM)");
    drawTextCentered("Camera ready", "Capturing every 60s");
    lastCapture = millis(); // start the interval
  } else {
    char buf[40];
    snprintf(buf, sizeof(buf), "Err=0x%X", (unsigned)err);
    Serial.printf("Camera init failed: 0x%X\n", err);
    drawTextCentered("Camera init failed", buf);
  }
}

// ===== Loop =====
void loop() {
  if (!camera_ready) {
    delay(250);
    return;
  }

  const unsigned long now = millis();
  if (now - lastCapture >= CAPTURE_INTERVAL_MS) {
    lastCapture = now;

    // indicate capture
    digitalWrite(PIN_LED, LOW);
    drawTextCentered("Capturing...", nullptr);

    camera_fb_t* fb = esp_camera_fb_get();
    digitalWrite(PIN_LED, HIGH);

    if (!fb) {
      Serial.println("Capture failed");
      drawTextCentered("Capture failed", nullptr);
    } else {
      Serial.printf("Captured %dx%d, len=%d\n", fb->width, fb->height, fb->len);

      // "Delete last picture" = replace previous OLED image: clear + redraw new frame
      processAndDisplayImage(fb);

      // release frame buffer back to driver (don’t store -> effectively delete)
      esp_camera_fb_return(fb);
    }
  }

  delay(30);
}