#include <lgfx.h>
#include <lvgl.h>

#include "esp_err.h"
#include "esp_heap_caps.h"
#include "esp_log.h"
#include "esp_task_wdt.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"

#include "ff.h"
#include "storage.h"
#include <stdlib.h>

#define DISPLAY_WIDTH (480)
#define DISPLAY_HEIGHT (320)

const unsigned int lvBufferSize = DISPLAY_WIDTH * DISPLAY_HEIGHT / 10 * (LV_COLOR_DEPTH / 8);
uint8_t lvBuffer1[lvBufferSize];
uint8_t lvBuffer2[lvBufferSize];

static const char *TAG = "main";

LGFX tft;

// Image cache variables
static uint8_t *image_cache = NULL;
static size_t image_cache_size = 0;
static lv_image_dsc_t cached_image_dsc;

void setup_tft(void)
{
    tft.begin();
    tft.setRotation(1);
    tft.setBrightness(255);
}

void flush(lv_display_t *display, const lv_area_t *area, unsigned char *data)
{
    uint32_t w = lv_area_get_width(area);
    uint32_t h = lv_area_get_height(area);
    lv_draw_sw_rgb565_swap(data, w * h);
    tft.pushImageDMA(area->x1, area->y1, w, h, (uint16_t *)data);
    lv_display_flush_ready(display);
}

void my_touch_read(lv_indev_t *indev_driver, lv_indev_data_t *data)
{
    uint16_t touchX, touchY;
    bool touched = tft.getTouch(&touchX, &touchY);
    if (!touched)
    {
        data->state = LV_INDEV_STATE_REL;
    }
    else
    {
        data->state = LV_INDEV_STATE_PR;
        data->point.x = touchX;
        data->point.y = touchY;
    }
}

bool load_image_to_cache(const char *filename)
{
    FIL file;
    FRESULT result = f_open(&file, filename, FA_READ);
    if (result != FR_OK)
    {
        ESP_LOGE(TAG, "Failed to open %s: %d", filename, result);
        return false;
    }

    // Determine file size
    FSIZE_t file_size = f_size(&file);
    ESP_LOGI(TAG, "Image file size: %d bytes", (int)file_size);

    // Allocate memory for entire image (prefer PSRAM)
    image_cache = (uint8_t *)heap_caps_malloc(file_size, MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT);
    if (image_cache == NULL)
    {
        ESP_LOGW(TAG, "Failed to allocate PSRAM, trying internal RAM");
        // Fallback to internal RAM
        image_cache = (uint8_t *)malloc(file_size);
    }

    if (image_cache == NULL)
    {
        ESP_LOGE(TAG, "Failed to allocate memory for image cache");
        f_close(&file);
        return false;
    }

    // Check which memory type is used
    bool is_psram = esp_ptr_external_ram(image_cache);
    ESP_LOGI(TAG, "Image cache allocated in %s: %d bytes", is_psram ? "PSRAM" : "internal RAM", (int)file_size);

    // Read entire image in one go
    UINT bytes_read;
    result = f_read(&file, image_cache, file_size, &bytes_read);
    f_close(&file);

    if (result != FR_OK || bytes_read != file_size)
    {
        ESP_LOGE(TAG, "Failed to read image file completely");
        free(image_cache);
        image_cache = NULL;
        return false;
    }

    image_cache_size = file_size;

    // Read image header from first bytes (LVGL BIN format)
    if (image_cache_size >= 4)
    {
        // Parse LVGL BIN header
        lv_image_header_t *header = (lv_image_header_t *)image_cache;

        // Set up image descriptor for LVGL with correct header data
        cached_image_dsc.header.w = header->w;
        cached_image_dsc.header.h = header->h;
        cached_image_dsc.header.cf = header->cf;
        cached_image_dsc.header.flags = header->flags;
        cached_image_dsc.header.stride = header->stride;
        cached_image_dsc.data_size = image_cache_size;
        cached_image_dsc.data = image_cache;

        ESP_LOGI(TAG, "Image header: %dx%d, format: %d", header->w, header->h, header->cf);
    }
    else
    {
        // Fallback for unknown format
        cached_image_dsc.header.w = 0;
        cached_image_dsc.header.h = 0;
        cached_image_dsc.header.cf = LV_COLOR_FORMAT_UNKNOWN;
        cached_image_dsc.header.flags = 0;
        cached_image_dsc.header.stride = 0;
        cached_image_dsc.data_size = image_cache_size;
        cached_image_dsc.data = image_cache;
    }

    ESP_LOGI(TAG, "Image successfully loaded to cache (%d bytes)", (int)image_cache_size);
    return true;
}

void esp_lv_log_print(lv_log_level_t level, const char *buf)
{
    switch (level)
    {
    case LV_LOG_LEVEL_TRACE:
        ESP_LOGV("LVGL", "%s", buf);
        break;
    case LV_LOG_LEVEL_INFO:
        ESP_LOGI("LVGL", "%s", buf);
        break;
    case LV_LOG_LEVEL_WARN:
        ESP_LOGW("LVGL", "%s", buf);
        break;
    case LV_LOG_LEVEL_ERROR:
        ESP_LOGE("LVGL", "%s", buf);
        break;
    case LV_LOG_LEVEL_USER:
        ESP_LOGI("LVGL", "%s", buf);
        break;
    case LV_LOG_LEVEL_NONE:
        break;
    }
}

void setup()
{
    setup_tft();

    lv_init();

    lv_log_register_print_cb(esp_lv_log_print);

    fs_mount();

    static auto *display = lv_display_create(DISPLAY_WIDTH, DISPLAY_HEIGHT);
    lv_display_set_color_format(display, LV_COLOR_FORMAT_RGB565);
    lv_display_set_flush_cb(display, flush);
    lv_display_set_buffers(display, lvBuffer1, lvBuffer2, lvBufferSize, LV_DISPLAY_RENDER_MODE_PARTIAL);

    static auto *lvInput = lv_indev_create();
    lv_indev_set_type(lvInput, LV_INDEV_TYPE_POINTER);
    lv_indev_set_read_cb(lvInput, my_touch_read);

    ESP_LOGI(TAG, "create image");

    // Load image once into cache
    if (load_image_to_cache("/poster.bin"))
    {
        lv_obj_t *img = lv_image_create(lv_screen_active());
        lv_image_set_src(img, &cached_image_dsc); // Use cached image data
        lv_obj_center(img);
    }
    else
    {
        ESP_LOGE(TAG, "Failed to load image from cache");
    }

    ESP_LOGI(TAG, "image created");
}

void loop()
{
    lv_tick_inc(10);
    lv_timer_handler();
    vTaskDelay(pdMS_TO_TICKS(10));
}

void cleanup_image_cache(void)
{
    if (image_cache != NULL)
    {
        // Use heap_caps_free since we used heap_caps_malloc
        if (esp_ptr_external_ram(image_cache))
        {
            ESP_LOGI(TAG, "Freeing PSRAM image cache");
        }
        else
        {
            ESP_LOGI(TAG, "Freeing internal RAM image cache");
        }
        heap_caps_free(image_cache);
        image_cache = NULL;
        image_cache_size = 0;
        ESP_LOGI(TAG, "Image cache cleaned up");
    }
}

void lvgl_task(void *pvParameter)
{
    setup();
    while (1)
    {
        loop();
    }

    cleanup_image_cache();
    fs_unmount();
}

extern "C" void app_main(void)
{
    xTaskCreatePinnedToCore(lvgl_task, "lvgl_task", 8192, NULL, 5, NULL, 1);
}