/*
 * Copyright (C) EdgeTX
 *
 * Based on code named
 *   opentx - https://github.com/opentx/opentx
 *   th9x - http://code.google.com/p/th9x
 *   er9x - http://code.google.com/p/er9x
 *   gruvin9x - http://code.google.com/p/gruvin9x
 *
 * License GPLv2: http://www.gnu.org/licenses/gpl-2.0.html
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as
 * published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 */

#include "edgetx.h"

#if defined(COLORLCD)
// Scaling for RBG565 to RGB888
uint8_t rbScale[32] = {
  0,8,16,24,32,41,49,57,65,74,82,90,98,106,115,123,131,139,148,156,164,172,180,189,197,205,213,222,230,238,246,255
};
uint8_t gScale[64] = {
  0,4,8,12,16,20,24,28,32,36,40,44,48,52,56,60,64,68,72,76,80,85,89,93,97,101,105,109,113,117,121,125,
  129,133,137,141,145,149,153,157,161,165,170,174,178,182,186,190,194,198,202,206,210,214,218,222,226,230,234,238,242,246,250,255
};
#endif

#if defined(COLORLCD) && (defined(STM32H7) || defined(STM32H7RS))

#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb/stb_image_write.h"

static void convert_RGB565_to_RGB888(uint8_t* dst, lv_img_dsc_t* src, coord_t w, coord_t h)
{
  for (int y = 0; y < h; y++) {
    for (int x = 0; x < w; x++) {
      lv_color_t pixel = lv_img_buf_get_px_color(src, x, y, {});
      *(dst++) = rbScale[pixel.ch.red];
      *(dst++) = gScale[pixel.ch.green];
      *(dst++) = rbScale[pixel.ch.blue];
    }
  }
}

static const char* writeResult = nullptr;

static void writeScreenToFile(void* context, void* data, int size)
{
  const char* file = (const char*)context;
  FIL fp;
  UINT written = 0;

  FRESULT result = f_open(&fp, file, FA_CREATE_ALWAYS | FA_WRITE);
  if (result == FR_OK) {
    result = f_write(&fp, data, size, &written);
  }
  f_close(&fp);

  if (result != FR_OK || (int)written != size) {
    writeResult = SDCARD_ERROR(result);
  }
}

static const char* _writeScreenshot()
{
  char filename[42]; // /SCREENSHOTS/screen-2013-01-01-123540.bmp

  if (sdIsFull()) {
    return STR_SDCARD_FULL_EXT;
  }

  // check and create folder here
  char* s = strAppend(filename, SCREENSHOTS_PATH);
  const char * error = sdCheckAndCreateDirectory(filename);
  if (error) {
    return error;
  }

  lv_img_dsc_t* snapshot = lv_snapshot_take(lv_scr_act(), LV_IMG_CF_TRUE_COLOR);
  if (!snapshot) {
    return "LVGL error creating snapshot";
  }

  s = strAppend(s, "/screen");
  s = strAppendDate(s, true);
  strAppend(s, ".png");

  writeResult = nullptr;

  // allocate enough for 3 channels
  auto pixels = LCD_W * LCD_H;
  auto stride = LCD_W * 3;
  uint8_t* img = (uint8_t*)malloc(pixels * 3);
  convert_RGB565_to_RGB888(img, snapshot, LCD_W, LCD_H);
  stbi_write_png_to_func(writeScreenToFile, filename, LCD_W, LCD_H, 3, img, stride);
  free(img);

  lv_snapshot_free(snapshot);

  return writeResult;
}

#else

#if defined(COLORLCD)
constexpr uint16_t BMP_HEADERSIZE = 0x36;
constexpr uint8_t BMP_BITS_PER_PIXEL = 32;
#else
constexpr uint16_t BMP_HEADERSIZE = 0x76;
constexpr uint8_t BMP_BITS_PER_PIXEL = 4;
#endif

constexpr uint32_t BMP_FILESIZE = uint32_t(BMP_HEADERSIZE + (LCD_W * LCD_H) * BMP_BITS_PER_PIXEL / 8);

const uint8_t BMP_HEADER[] = {
  'B',  'M',
  /* file size */ BMP_FILESIZE & 0xFF, (BMP_FILESIZE >> 8) & 0xFF, (BMP_FILESIZE >> 16) & 0xFF, (BMP_FILESIZE >> 24) & 0xFF,
  0x00, 0x00, 0x00, 0x00,
  /* header size */ BMP_HEADERSIZE, 0x00, 0x00, 0x00,
  /* extra header size */ 0x28, 0x00, 0x00, 0x00,
  /* width */ LCD_W & 0xFF, LCD_W >> 8, 0x00, 0x00,
  /* height */ LCD_H & 0xFF, LCD_H >> 8, 0x00, 0x00,
  /* planes */ 0x01, 0x00,
  /* depth */ BMP_BITS_PER_PIXEL, 0x00,
  /* compression */ 0x00, 0x00, 0x00, 0x00,
  /* image size */ 0x00, 0x00, 0x00, 0x00,
  /* hor res 72dpi */ 0x13, 0x0B, 0x00, 0x00,
  /* ver res 72 dpi */ 0x13, 0x0B, 0x00, 0x00,
  /* num colors */ 0x00, 0x00, 0x00, 0x00,
  /* important colors */ 0x00, 0x00, 0x00, 0x00,
#if !defined(COLORLCD)
  /* grayscale palette */
#if defined(OLED_SCREEN)
  0x00, 0x00, 0x00, 0x00, 0x11, 0x11, 0x11, 0x00, 0x22, 0x22, 0x22, 0x00, 0x33, 0x33, 0x33, 0x00,
  0x44, 0x44, 0x44, 0x00, 0x55, 0x55, 0x55, 0x00, 0x66, 0x66, 0x66, 0x00, 0x77, 0x77, 0x77, 0x00,
  0x88, 0x88, 0x88, 0x00, 0x99, 0x99, 0x99, 0x00, 0xaa, 0xaa, 0xaa, 0x00, 0xbb, 0xbb, 0xbb, 0x00,
  0xcc, 0xcc, 0xcc, 0x00, 0xdd, 0xdd, 0xdd, 0x00, 0xee, 0xee, 0xee, 0x00, 0xff, 0xff, 0xff, 0x00,
#else
  0xff, 0xff, 0xff, 0x00, 0xee, 0xee, 0xee, 0x00, 0xdd, 0xdd, 0xdd, 0x00, 0xcc, 0xcc, 0xcc, 0x00,
  0xbb, 0xbb, 0xbb, 0x00, 0xaa, 0xaa, 0xaa, 0x00, 0x99, 0x99, 0x99, 0x00, 0x88, 0x88, 0x88, 0x00,
  0x77, 0x77, 0x77, 0x00, 0x66, 0x66, 0x66, 0x00, 0x55, 0x55, 0x55, 0x00, 0x44, 0x44, 0x44, 0x00,
  0x33, 0x33, 0x33, 0x00, 0x22, 0x22, 0x22, 0x00, 0x11, 0x11, 0x11, 0x00, 0x00, 0x00, 0x00, 0x00,
#endif
#endif
};

static const char* _writeScreenshot()
{
  FIL bmpFile;
  UINT written;
  char filename[42]; // /SCREENSHOTS/screen-2013-01-01-123540.bmp

  if (sdIsFull()) {
    return STR_SDCARD_FULL_EXT;
  }

  // check and create folder here
  strcpy(filename, SCREENSHOTS_PATH);
  const char * error = sdCheckAndCreateDirectory(filename);
  if (error) {
    return error;
  }

#if defined(RTCLOCK)
  char * tmp = strAppend(&filename[sizeof(SCREENSHOTS_PATH)-1], "/screen");
  tmp = strAppendDate(tmp, true);
  strcpy(tmp, BMP_EXT);
#endif

  FRESULT result = f_open(&bmpFile, filename, FA_CREATE_ALWAYS | FA_WRITE);
  if (result != FR_OK) {
    return SDCARD_ERROR(result);
  }

  result = f_write(&bmpFile, BMP_HEADER, sizeof(BMP_HEADER), &written);
  if (result != FR_OK || written != sizeof(BMP_HEADER)) {
    f_close(&bmpFile);
    return SDCARD_ERROR(result);
  }

#if defined(COLORLCD)
  lv_img_dsc_t* snapshot = lv_snapshot_take(lv_scr_act(), LV_IMG_CF_TRUE_COLOR);
  if (!snapshot) {
    f_close(&bmpFile);
    return "LVGL error creating snapshot";
  }

  auto w = snapshot->header.w;
  auto h = snapshot->header.h;

  for (int y = h - 1; y >= 0; y--) {
    for (uint32_t x = 0; x < w; x++) {

      lv_color_t pixel = lv_img_buf_get_px_color(snapshot, x, y, {});

      uint32_t dst = (0xFF << 24)
          | (rbScale[pixel.ch.red] << 16)
          | (gScale[pixel.ch.green] << 8)
          | rbScale[pixel.ch.blue];

      result = f_write(&bmpFile, &dst, sizeof(dst), &written);
      if (result != FR_OK || written != sizeof(dst)) {
        lv_snapshot_free(snapshot);
        f_close(&bmpFile);
        return SDCARD_ERROR(result);
      }
    }
  }

  lv_snapshot_free(snapshot);

#else // stdlcd

#if LCD_W == 128
  pixel_t inv = g_eeGeneral.invertLCD ? 0xFF : 0;
#endif

  for (int y=LCD_H-1; y>=0; y-=1) {
    for (int x=0; x<8*((LCD_W+7)/8); x+=2) {
#if LCD_W == 128
      pixel_t byte = (getPixel(x+1, y) | (getPixel(x, y) << 4)) ^ inv;
#else
      pixel_t byte = getPixel(x+1, y) | (getPixel(x, y) << 4);
#endif
      result = f_write(&bmpFile, &byte, 1, &written);
      if (result != FR_OK || written != 1) {
        f_close(&bmpFile);
        return SDCARD_ERROR(result);
      }
    }
  }
#endif

  f_close(&bmpFile);

  return nullptr;
}

#endif

void writeScreenshot()
{
  const char* res = _writeScreenshot();
  if (res)
    POPUP_WARNING(res);
}