#include "epd.h"
#include "config.h"
#include <LittleFS.h>
#include <qrcode.h>
#include <esp_task_wdt.h>

static uint8_t s_row[EPD_WIDTH / 2];

static void wait_busy() {
    uint32_t start = millis();
    while (digitalRead(PIN_BUSY) == LOW) {
        if (millis() - start > 60000) return;  // 6-color refresh takes ~20s
        delay(5);
        esp_task_wdt_reset();  // feed WDT — display refresh can take ~20 s
    }
}
static void cmd(uint8_t c) {
    digitalWrite(PIN_DC, LOW);
    digitalWrite(PIN_CS, LOW);
    SPI.transfer(c);
    digitalWrite(PIN_CS, HIGH);
}
static void dat(uint8_t d) {
    digitalWrite(PIN_DC, HIGH);
    digitalWrite(PIN_CS, LOW);
    SPI.transfer(d);
    digitalWrite(PIN_CS, HIGH);
}

void epd_init() {
    digitalWrite(PIN_RST, HIGH); delay(20);
    digitalWrite(PIN_RST, LOW);  delay(10);
    digitalWrite(PIN_RST, HIGH); delay(20);
    wait_busy(); delay(30);

    cmd(0xAA);
    dat(0x49); dat(0x55); dat(0x20);
    dat(0x08); dat(0x09); dat(0x18);
    cmd(0x01); dat(0x3F);
    cmd(0x00); dat(0x5F); dat(0x69);
    cmd(0x03); dat(0x00); dat(0x54); dat(0x00); dat(0x44);
    cmd(0x05); dat(0x40); dat(0x1F); dat(0x1F); dat(0x2C);
    cmd(0x06); dat(0x6F); dat(0x1F); dat(0x17); dat(0x49);
    cmd(0x08); dat(0x6F); dat(0x1F); dat(0x1F); dat(0x22);
    cmd(0x30); dat(0x03);
    cmd(0x50); dat(0x3F);
    cmd(0x60); dat(0x02); dat(0x00);
    cmd(0x61); dat(0x03); dat(0x20); dat(0x01); dat(0xE0);
    cmd(0x84); dat(0x01);
    cmd(0xE3); dat(0x2F);
    cmd(0x04); wait_busy();
}

void epd_sleep() {
    cmd(0x02); dat(0x00); wait_busy();
    cmd(0x07); dat(0xA5);
}

static void epd_refresh() {
    cmd(0x04); wait_busy();
    cmd(0x12); dat(0x00); wait_busy();
}

void epd_fill(uint8_t color) {
    uint8_t byte = (color << 4) | color;
    cmd(0x10);
    digitalWrite(PIN_DC, HIGH);
    digitalWrite(PIN_CS, LOW);
    for (int i = 0; i < EPD_WIDTH * EPD_HEIGHT / 2; i++) SPI.transfer(byte);
    digitalWrite(PIN_CS, HIGH);
    epd_refresh();
}

void epd_draw_image_from_file(fs::File& f) {
    cmd(0x10);
    uint8_t buf[512];
    digitalWrite(PIN_DC, HIGH);
    digitalWrite(PIN_CS, LOW);
    while (f.available()) {
        size_t n = f.read(buf, sizeof(buf));
        SPI.writeBytes(buf, n);
    }
    digitalWrite(PIN_CS, HIGH);
    epd_refresh();
}

void epd_draw_image_with_border(fs::File& f, uint8_t color, int thickness) {
    const size_t expected = (size_t)EPD_WIDTH * EPD_HEIGHT / 2;
    if (f.size() != expected) {
        epd_fill(color);
        return;
    }

    const uint8_t pair = (color << 4) | color;

    cmd(0x10);
    for (int y = 0; y < EPD_HEIGHT; y++) {
        f.read(s_row, sizeof(s_row));

        if (y < thickness || y >= EPD_HEIGHT - thickness) {
            // Top/bottom band — solid color across the row.
            for (size_t i = 0; i < sizeof(s_row); i++) s_row[i] = pair;
        } else {
            // Middle band — overlay border on left/right edges.
            for (int x = 0; x < thickness; x++) {
                if (x & 1) s_row[x/2] = (s_row[x/2] & 0xF0) | color;
                else       s_row[x/2] = (s_row[x/2] & 0x0F) | (color << 4);
            }
            for (int x = EPD_WIDTH - thickness; x < EPD_WIDTH; x++) {
                if (x & 1) s_row[x/2] = (s_row[x/2] & 0xF0) | color;
                else       s_row[x/2] = (s_row[x/2] & 0x0F) | (color << 4);
            }
        }

        digitalWrite(PIN_DC, HIGH); digitalWrite(PIN_CS, LOW);
        SPI.writeBytes(s_row, sizeof(s_row));
        digitalWrite(PIN_CS, HIGH);
    }
    epd_refresh();
}

void epd_draw_qr(QRCode* qr, uint8_t cellPx, uint8_t bg, uint8_t fg) {
    int qrPx = qr->size * cellPx;
    int offX  = (EPD_WIDTH  - qrPx) / 2;
    int offY  = (EPD_HEIGHT - qrPx) / 2;

    cmd(0x10);
    for (int y = 0; y < EPD_HEIGHT; y++) {
        for (int x = 0; x < EPD_WIDTH; x += 2) {
            auto nibble = [&](int px) -> uint8_t {
                int qx = (px - offX) / cellPx, qy = (y - offY) / cellPx;
                if (qx >= 0 && qx < qr->size && qy >= 0 && qy < qr->size)
                    return qrcode_getModule(qr, qx, qy) ? fg : bg;
                return bg;
            };
            s_row[x/2] = (nibble(x) << 4) | nibble(x+1);
        }
        digitalWrite(PIN_DC, HIGH); digitalWrite(PIN_CS, LOW);
        SPI.writeBytes(s_row, sizeof(s_row));
        digitalWrite(PIN_CS, HIGH);
    }
    epd_refresh();
}

// Stream background from LittleFS, overlaying QR at (qr_x, qr_y) with given cell size.
// Falls back to a solid fill if the file is missing.
static void draw_from_lfs(const char* path, uint8_t fallback_color,
                           QRCode* qr, int qr_x, int qr_y, int qr_cell) {
    File f = LittleFS.open(path, "r");
    if (!f) { epd_fill(fallback_color); return; }

    int qr_px = qr->size * qr_cell;

    cmd(0x10);
    for (int y = 0; y < EPD_HEIGHT; y++) {
        f.read(s_row, sizeof(s_row));

        if (y >= qr_y && y < qr_y + qr_px) {
            int qy = (y - qr_y) / qr_cell;
            int x0 = max(qr_x, 0), x1 = min(qr_x + qr_px, EPD_WIDTH);
            for (int x = x0; x < x1; x++) {
                uint8_t c = qrcode_getModule(qr, (x - qr_x) / qr_cell, qy)
                            ? COLOR_BLACK : COLOR_WHITE;
                if (x & 1) s_row[x/2] = (s_row[x/2] & 0xF0) | c;
                else       s_row[x/2] = (s_row[x/2] & 0x0F) | (c << 4);
            }
        }

        digitalWrite(PIN_DC, HIGH); digitalWrite(PIN_CS, LOW);
        SPI.writeBytes(s_row, sizeof(s_row));
        digitalWrite(PIN_CS, HIGH);
    }
    f.close();
    epd_refresh();
}

void epd_draw_ap_screen(QRCode* qr) {
    // AP_QR_X=563, AP_QR_Y=185, AP_QR_CELL=5  (must match gen_screens.py)
    draw_from_lfs("/ap_bg.bin", COLOR_YELLOW, qr, 563, 185, 5);
}

void epd_draw_setup_screen(QRCode* qr) {
    // SETUP_QR_X=553, SETUP_QR_Y=175, SETUP_QR_CELL=5  (must match gen_screens.py)
    draw_from_lfs("/setup_bg.bin", COLOR_GREEN, qr, 553, 175, 5);
}