Xteink-X4-crosspoint-reader/lib/Epub/Epub/converters/JpegToFramebufferConverter.cpp

#include "JpegToFramebufferConverter.h"
#include <GfxRenderer.h>
#include <HardwareSerial.h>
#include <SdFat.h>
#include <SDCardManager.h>
#include <picojpeg.h>
#include <cstdio>
#include <cstring>

struct JpegContext {
  FsFile& file;
  uint8_t buffer[512];
  size_t bufferPos;
  size_t bufferFilled;
  JpegContext(FsFile& f) : file(f), bufferPos(0), bufferFilled(0) {}
};

// Cache buffer for storing 2-bit pixels during decode
struct PixelCache {
  uint8_t* buffer;
  int width;
  int height;
  int bytesPerRow;
  int originX;  // config.x - to convert screen coords to cache coords
  int originY;  // config.y

  PixelCache() : buffer(nullptr), width(0), height(0), bytesPerRow(0), originX(0), originY(0) {}

  bool allocate(int w, int h, int ox, int oy) {
    width = w;
    height = h;
    originX = ox;
    originY = oy;
    bytesPerRow = (w + 3) / 4;  // 2 bits per pixel, 4 pixels per byte
    size_t bufferSize = bytesPerRow * h;
    buffer = (uint8_t*)malloc(bufferSize);
    if (buffer) {
      memset(buffer, 0, bufferSize);
      Serial.printf("[%lu] [JPG] Allocated cache buffer: %d bytes for %dx%d\n", millis(), bufferSize, w, h);
    }
    return buffer != nullptr;
  }

  void setPixel(int screenX, int screenY, uint8_t value) {
    if (!buffer) return;
    int localX = screenX - originX;
    int localY = screenY - originY;
    if (localX < 0 || localX >= width || localY < 0 || localY >= height) return;

    int byteIdx = localY * bytesPerRow + localX / 4;
    int bitShift = 6 - (localX % 4) * 2;  // MSB first: pixel 0 at bits 6-7
    buffer[byteIdx] = (buffer[byteIdx] & ~(0x03 << bitShift)) | ((value & 0x03) << bitShift);
  }

  bool writeToFile(const std::string& cachePath) {
    if (!buffer) return false;

    FsFile cacheFile;
    if (!SdMan.openFileForWrite("IMG", cachePath, cacheFile)) {
      Serial.printf("[%lu] [JPG] Failed to open cache file for writing: %s\n", millis(), cachePath.c_str());
      return false;
    }

    uint16_t w = width;
    uint16_t h = height;
    cacheFile.write(&w, 2);
    cacheFile.write(&h, 2);
    cacheFile.write(buffer, bytesPerRow * height);
    cacheFile.close();

    Serial.printf("[%lu] [JPG] Cache written: %s (%dx%d, %d bytes)\n", millis(),
                  cachePath.c_str(), width, height, 4 + bytesPerRow * height);
    return true;
  }

  ~PixelCache() {
    if (buffer) {
      free(buffer);
      buffer = nullptr;
    }
  }
};

static int16_t ditherErrors[512][3];

bool JpegToFramebufferConverter::getDimensionsStatic(const std::string& imagePath, ImageDimensions& out) {
  FsFile file;
  if (!SdMan.openFileForRead("JPG", imagePath, file)) {
    Serial.printf("[%lu] [JPG] Failed to open file for dimensions: %s\n", millis(), imagePath.c_str());
    return false;
  }

  JpegContext context(file);
  pjpeg_image_info_t imageInfo;

  int status = pjpeg_decode_init(&imageInfo, jpegReadCallback, &context, 0);
  file.close();

  if (status != 0) {
    Serial.printf("[%lu] [JPG] Failed to init JPEG for dimensions: %d\n", millis(), status);
    return false;
  }

  out.width = imageInfo.m_width;
  out.height = imageInfo.m_height;
  Serial.printf("[%lu] [JPG] Image dimensions: %dx%d\n", millis(), out.width, out.height);
  return true;
}

uint8_t JpegToFramebufferConverter::applyAtkinsonDithering(uint8_t gray, int x, int y, int width) {
  int16_t error = gray - (gray < 128 ? 0 : 255);
  uint8_t newGray = gray - error;
  
  int8_t fraction = error >> 3;
  
  if (x + 1 < width && y + 1 < 512) ditherErrors[y + 1][(x + 1) % 3] += fraction;
  if (x + 2 < width && y + 1 < 512) ditherErrors[y + 1][(x + 2) % 3] += fraction;
  if (x + 1 < width) ditherErrors[y][(x + 1) % 3] += fraction;
  if (x + 2 < width) ditherErrors[y][(x + 2) % 3] += fraction;
  if (x - 1 >= 0 && x + 1 < width && y + 1 < 512) ditherErrors[y + 1][(x - 1 + 1) % 3] += fraction;
  if (x - 1 >= 0 && y + 1 < 512) ditherErrors[y + 1][(x - 1) % 3] += fraction;
  if (x + 1 < width && y + 2 < 512) ditherErrors[y + 2][(x + 1) % 3] += fraction;
  
  int16_t adjustedGray = newGray + ditherErrors[y][x % 3];
  if (adjustedGray < 0) adjustedGray = 0;
  if (adjustedGray > 255) adjustedGray = 255;
  
  uint8_t outputGray;
  if (adjustedGray < 64) {
    outputGray = 0;
  } else if (adjustedGray < 128) {
    outputGray = 1;
  } else if (adjustedGray < 192) {
    outputGray = 2;
  } else {
    outputGray = 3;
  }
  
  ditherErrors[y][x % 3] = adjustedGray - (outputGray * 85);
  
  return outputGray;
}

bool JpegToFramebufferConverter::decodeToFramebuffer(
  const std::string& imagePath,
  GfxRenderer& renderer,
  const RenderConfig& config
) {
  Serial.printf("[%lu] [JPG] Decoding JPEG: %s\n", millis(), imagePath.c_str());

  FsFile file;
  if (!SdMan.openFileForRead("JPG", imagePath, file)) {
    Serial.printf("[%lu] [JPG] Failed to open file: %s\n", millis(), imagePath.c_str());
    return false;
  }

  memset(ditherErrors, 0, sizeof(ditherErrors));

  JpegContext context(file);
  pjpeg_image_info_t imageInfo;

  int status = pjpeg_decode_init(&imageInfo, jpegReadCallback, &context, 0);
  if (status != 0) {
    Serial.printf("[%lu] [JPG] picojpeg init failed: %d\n", millis(), status);
    file.close();
    return false;
  }

  if (!validateImageDimensions(imageInfo.m_width, imageInfo.m_height, "JPEG")) {
    file.close();
    return false;
  }

  // Calculate scale factor to fit within maxWidth/maxHeight
  float scaleX = (config.maxWidth > 0 && imageInfo.m_width > config.maxWidth)
                   ? (float)config.maxWidth / imageInfo.m_width
                   : 1.0f;
  float scaleY = (config.maxHeight > 0 && imageInfo.m_height > config.maxHeight)
                   ? (float)config.maxHeight / imageInfo.m_height
                   : 1.0f;
  float scale = (scaleX < scaleY) ? scaleX : scaleY;
  if (scale > 1.0f) scale = 1.0f;

  int destWidth = (int)(imageInfo.m_width * scale);
  int destHeight = (int)(imageInfo.m_height * scale);

  Serial.printf("[%lu] [JPG] JPEG %dx%d -> %dx%d (scale %.2f), scan type: %d, MCU: %dx%d\n", millis(),
                imageInfo.m_width, imageInfo.m_height, destWidth, destHeight, scale,
                imageInfo.m_scanType, imageInfo.m_MCUWidth, imageInfo.m_MCUHeight);

  if (!imageInfo.m_pMCUBufR || !imageInfo.m_pMCUBufG || !imageInfo.m_pMCUBufB) {
    Serial.printf("[%lu] [JPG] Null buffer pointers in imageInfo\n", millis());
    file.close();
    return false;
  }

  const int screenWidth = renderer.getScreenWidth();
  const int screenHeight = renderer.getScreenHeight();

  // Allocate pixel cache if cachePath is provided
  PixelCache cache;
  bool caching = !config.cachePath.empty();
  if (caching) {
    if (!cache.allocate(destWidth, destHeight, config.x, config.y)) {
      Serial.printf("[%lu] [JPG] Failed to allocate cache buffer, continuing without caching\n", millis());
      caching = false;
    }
  }

  int mcuX = 0;
  int mcuY = 0;

  while (mcuY < imageInfo.m_MCUSPerCol) {
    status = pjpeg_decode_mcu();
    if (status == PJPG_NO_MORE_BLOCKS) {
      break;
    }
    if (status != 0) {
      Serial.printf("[%lu] [JPG] MCU decode failed: %d\n", millis(), status);
      file.close();
      return false;
    }

    // Source position in image coordinates
    int srcStartX = mcuX * imageInfo.m_MCUWidth;
    int srcStartY = mcuY * imageInfo.m_MCUHeight;

    switch (imageInfo.m_scanType) {
      case PJPG_GRAYSCALE:
        for (int row = 0; row < 8; row++) {
          int srcY = srcStartY + row;
          int destY = config.y + (int)(srcY * scale);
          if (destY >= screenHeight || destY >= config.y + destHeight) continue;
          for (int col = 0; col < 8; col++) {
            int srcX = srcStartX + col;
            int destX = config.x + (int)(srcX * scale);
            if (destX >= screenWidth || destX >= config.x + destWidth) continue;
            uint8_t gray = imageInfo.m_pMCUBufR[row * 8 + col];
            uint8_t dithered = config.useDithering ? applyAtkinsonDithering(gray, destX, destY, screenWidth) : gray / 85;
            if (dithered > 3) dithered = 3;
            renderer.drawPixel(destX, destY, dithered < 2);
            if (caching) cache.setPixel(destX, destY, dithered);
          }
        }
        break;

      case PJPG_YH1V1:
        for (int row = 0; row < 8; row++) {
          int srcY = srcStartY + row;
          int destY = config.y + (int)(srcY * scale);
          if (destY >= screenHeight || destY >= config.y + destHeight) continue;
          for (int col = 0; col < 8; col++) {
            int srcX = srcStartX + col;
            int destX = config.x + (int)(srcX * scale);
            if (destX >= screenWidth || destX >= config.x + destWidth) continue;
            uint8_t r = imageInfo.m_pMCUBufR[row * 8 + col];
            uint8_t g = imageInfo.m_pMCUBufG[row * 8 + col];
            uint8_t b = imageInfo.m_pMCUBufB[row * 8 + col];
            uint8_t gray = (uint8_t)((r * 77 + g * 150 + b * 29) >> 8);
            uint8_t dithered = config.useDithering ? applyAtkinsonDithering(gray, destX, destY, screenWidth) : gray / 85;
            if (dithered > 3) dithered = 3;
            renderer.drawPixel(destX, destY, dithered < 2);
            if (caching) cache.setPixel(destX, destY, dithered);
          }
        }
        break;

      case PJPG_YH2V1:
        for (int row = 0; row < 8; row++) {
          int srcY = srcStartY + row;
          int destY = config.y + (int)(srcY * scale);
          if (destY >= screenHeight || destY >= config.y + destHeight) continue;
          for (int col = 0; col < 16; col++) {
            int srcX = srcStartX + col;
            int destX = config.x + (int)(srcX * scale);
            if (destX >= screenWidth || destX >= config.x + destWidth) continue;
            int blockIndex = (col < 8) ? 0 : 1;
            int pixelIndex = row * 8 + (col % 8);
            uint8_t r = imageInfo.m_pMCUBufR[blockIndex * 64 + pixelIndex];
            uint8_t g = imageInfo.m_pMCUBufG[blockIndex * 64 + pixelIndex];
            uint8_t b = imageInfo.m_pMCUBufB[blockIndex * 64 + pixelIndex];
            uint8_t gray = (uint8_t)((r * 77 + g * 150 + b * 29) >> 8);
            uint8_t dithered = config.useDithering ? applyAtkinsonDithering(gray, destX, destY, screenWidth) : gray / 85;
            if (dithered > 3) dithered = 3;
            renderer.drawPixel(destX, destY, dithered < 2);
            if (caching) cache.setPixel(destX, destY, dithered);
          }
        }
        break;

      case PJPG_YH1V2:
        for (int row = 0; row < 16; row++) {
          int srcY = srcStartY + row;
          int destY = config.y + (int)(srcY * scale);
          if (destY >= screenHeight || destY >= config.y + destHeight) continue;
          for (int col = 0; col < 8; col++) {
            int srcX = srcStartX + col;
            int destX = config.x + (int)(srcX * scale);
            if (destX >= screenWidth || destX >= config.x + destWidth) continue;
            int blockIndex = (row < 8) ? 0 : 1;
            int pixelIndex = (row % 8) * 8 + col;
            uint8_t r = imageInfo.m_pMCUBufR[blockIndex * 128 + pixelIndex];
            uint8_t g = imageInfo.m_pMCUBufG[blockIndex * 128 + pixelIndex];
            uint8_t b = imageInfo.m_pMCUBufB[blockIndex * 128 + pixelIndex];
            uint8_t gray = (uint8_t)((r * 77 + g * 150 + b * 29) >> 8);
            uint8_t dithered = config.useDithering ? applyAtkinsonDithering(gray, destX, destY, screenWidth) : gray / 85;
            if (dithered > 3) dithered = 3;
            renderer.drawPixel(destX, destY, dithered < 2);
            if (caching) cache.setPixel(destX, destY, dithered);
          }
        }
        break;

      case PJPG_YH2V2:
        for (int row = 0; row < 16; row++) {
          int srcY = srcStartY + row;
          int destY = config.y + (int)(srcY * scale);
          if (destY >= screenHeight || destY >= config.y + destHeight) continue;
          for (int col = 0; col < 16; col++) {
            int srcX = srcStartX + col;
            int destX = config.x + (int)(srcX * scale);
            if (destX >= screenWidth || destX >= config.x + destWidth) continue;
            int blockX = (col < 8) ? 0 : 1;
            int blockY = (row < 8) ? 0 : 1;
            int blockIndex = blockY * 2 + blockX;
            int pixelIndex = (row % 8) * 8 + (col % 8);
            int blockOffset = blockIndex * 64;
            uint8_t r = imageInfo.m_pMCUBufR[blockOffset + pixelIndex];
            uint8_t g = imageInfo.m_pMCUBufG[blockOffset + pixelIndex];
            uint8_t b = imageInfo.m_pMCUBufB[blockOffset + pixelIndex];
            uint8_t gray = (uint8_t)((r * 77 + g * 150 + b * 29) >> 8);
            uint8_t dithered = config.useDithering ? applyAtkinsonDithering(gray, destX, destY, screenWidth) : gray / 85;
            if (dithered > 3) dithered = 3;
            renderer.drawPixel(destX, destY, dithered < 2);
            if (caching) cache.setPixel(destX, destY, dithered);
          }
        }
        break;
    }

    mcuX++;
    if (mcuX >= imageInfo.m_MCUSPerRow) {
      mcuX = 0;
      mcuY++;
    }
  }

  Serial.printf("[%lu] [JPG] Decoding complete\n", millis());
  file.close();

  // Write cache file if caching was enabled
  if (caching) {
    cache.writeToFile(config.cachePath);
  }

  return true;
}

unsigned char JpegToFramebufferConverter::jpegReadCallback(
  unsigned char* pBuf,
  unsigned char buf_size,
  unsigned char* pBytes_actually_read,
  void* pCallback_data
) {
  JpegContext* context = reinterpret_cast<JpegContext*>(pCallback_data);
  
  if (context->bufferPos >= context->bufferFilled) {
    int readCount = context->file.read(context->buffer, sizeof(context->buffer));
    if (readCount <= 0) {
      *pBytes_actually_read = 0;
      return 0;
    }
    context->bufferFilled = readCount;
    context->bufferPos = 0;
  }
  
  unsigned int bytesAvailable = context->bufferFilled - context->bufferPos;
  unsigned int bytesToCopy = (bytesAvailable < buf_size) ? bytesAvailable : buf_size;
  
  memcpy(pBuf, &context->buffer[context->bufferPos], bytesToCopy);
  context->bufferPos += bytesToCopy;
  *pBytes_actually_read = bytesToCopy;
  
  return 0;
}

bool JpegToFramebufferConverter::supportsFormat(const std::string& extension) const {
  std::string ext = extension;
  for (auto& c : ext) {
    c = tolower(c);
  }
  return (ext == ".jpg" || ext == ".jpeg");
}