#include <stdio.h>

#define LEPTON_ENABLE

#include "mbedtls/base64.h"


#define TAG "CAM"
#include "esp_camera.h"
#define PWDN_GPIO_NUM -1
#define RESET_GPIO_NUM -1
#define XCLK_GPIO_NUM 15
#define SIOD_GPIO_NUM 4
#define SIOC_GPIO_NUM 5

#define Y2_GPIO_NUM 11
#define Y3_GPIO_NUM 9
#define Y4_GPIO_NUM 8
#define Y5_GPIO_NUM 10
#define Y6_GPIO_NUM 12
#define Y7_GPIO_NUM 18
#define Y8_GPIO_NUM 17
#define Y9_GPIO_NUM 16

#define VSYNC_GPIO_NUM 6
#define HREF_GPIO_NUM 7
#define PCLK_GPIO_NUM 13

// #define ENABLE_LEPTON

#define WIFI_IF WIFI_IF_AP

#define MAX_PAYLOAD_SIZE 1472

// uint8_t hehe[10000];

#include "esp_event.h"
#include "esp_system.h"
#include "esp_event.h"
#include "esp_timer.h"
#include "esp_wifi.h"
#include "esp_log.h"
#include "esp_mac.h"
#include "nvs_flash.h"
#include "string.h"

#include "lipton/cci.h"
#include "lipton/vospi.h"
#include "lipton/lepton_system.h"
#include "lipton/lepton_utilities.h"
#include <driver/uart.h>
#include <math.h>

// Number of consecutive VoSPI resynchronization attempts before attempting to reset
#define LEP_SYNC_FAIL_FAULT_LIMIT 10

// Reset fail delay before attempting a re-init (seconds)
#define LEP_RESET_FAIL_RETRY_SECS 5

#define BUF_SIZE 20000

typedef struct {
    int x, y, width, height;
} Rectangle;


//
// Code start
//
int vsync_count = 0;
int sync_fail_count = 0;
int reset_fail_count = 0;
int64_t vsyncDetectedUsec;

esp_err_t tx_with_retry(wifi_interface_t iface, const void *buffer, int len, bool en_sys_seq)
{
    for (int i = 0; i < 160; i++)
    {
        esp_err_t err = esp_wifi_80211_tx(iface, buffer, len, en_sys_seq);
        if (err == ESP_OK)
            return ESP_OK;
        if (err != ESP_ERR_NO_MEM)
            return err;
        // No delay, retry immediately
    }
    return ESP_ERR_NO_MEM;
}

static camera_config_t camera_config = {
    .pin_pwdn = PWDN_GPIO_NUM,
    .pin_reset = RESET_GPIO_NUM,
    .pin_xclk = XCLK_GPIO_NUM,

#ifndef LEPTON_ENABLE
    .pin_sccb_sda = SIOD_GPIO_NUM,
    .pin_sccb_scl = SIOC_GPIO_NUM,
#else
    .pin_sccb_sda = -1,
    .pin_sccb_scl = -1,
#endif
    .sccb_i2c_port = I2C_MASTER_NUM,

    .pin_d7 = Y9_GPIO_NUM,
    .pin_d6 = Y8_GPIO_NUM,
    .pin_d5 = Y7_GPIO_NUM,
    .pin_d4 = Y6_GPIO_NUM,
    .pin_d3 = Y5_GPIO_NUM,
    .pin_d2 = Y4_GPIO_NUM,
    .pin_d1 = Y3_GPIO_NUM,
    .pin_d0 = Y2_GPIO_NUM,
    .pin_vsync = VSYNC_GPIO_NUM,
    .pin_href = HREF_GPIO_NUM,
    .pin_pclk = PCLK_GPIO_NUM,

    .xclk_freq_hz = 20000000, // EXPERIMENTAL: Set to 16MHz on ESP32-S2 or ESP32-S3 to enable EDMA mode
    .ledc_timer = LEDC_TIMER_0,
    .ledc_channel = LEDC_CHANNEL_0,

    .pixel_format = PIXFORMAT_JPEG, // YUV422,GRAYSCALE,RGB565,JPEG
    .frame_size = FRAMESIZE_UXGA,   // QQVGA-UXGA, For ESP32, do not use sizes above QVGA when not JPEG. The performance of the ESP32-S series has improved a lot, but JPEG mode always gives better frame rates.+
    //.frame_size = FRAMESIZE_128X128, // QQVGA-UXGA, For ESP32, do not use sizes above QVGA when not JPEG. The performance of the ESP32-S series has improved a lot, but JPEG mode always gives better frame rates.

    .jpeg_quality = 12, // 0-63, for OV series camera sensors, lower number means higher quality
    .fb_count = 1,      // When jpeg mode is used, if fb_count more than one, the driver will work in continuous mode.
    //.fb_location = CAMERA_FB_IN_PSRAM,
    .fb_location = CAMERA_FB_IN_PSRAM,
    .grab_mode = CAMERA_GRAB_LATEST // CAMERA_GRAB_LATEST. Sets when buffers should be filled
};

esp_err_t esp_wifi_80211_tx(wifi_interface_t ifx, const void *buffer, int len, bool en_sys_seq);

void send_data_frame(uint8_t type, uint8_t *data, size_t len)
{
    uint8_t mac_addr[6];
    if (esp_wifi_get_mac(WIFI_IF_AP, mac_addr) != ESP_OK)
    {
        ESP_LOGE(TAG, "Failed to get MAC address");
        return;
    }

    // 802.11 header template
    uint8_t header[] = {
        0x08, 0x02, 0x00, 0x00,
        0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // Destination: broadcast
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // Source: MAC (to be set)
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // BSSID: MAC (to be set)
        0x10, 0x00                          // Sequence control
    };
    memcpy(&header[10], mac_addr, 6); // Source
    memcpy(&header[16], mac_addr, 6); // BSSID

    size_t max_chunk_size = MAX_PAYLOAD_SIZE - (11 + 2 + 1); // Payload minus custom header and checksum
    size_t total_chunks = (len + max_chunk_size - 1) / max_chunk_size;

    ESP_LOGI(TAG, "Got data with %d bytes", len);

    size_t offset = 0;
    for (uint16_t seq = 0; seq < total_chunks; seq++)
    {
        size_t chunk_size = len - offset;
        if (chunk_size > max_chunk_size)
        {
            chunk_size = max_chunk_size;
        }
        size_t full_len = sizeof(header) + 11 + chunk_size;

        uint8_t *full_packet = malloc(full_len + 2);
        if (!full_packet)
        {
            ESP_LOGE(TAG, "Memory allocation failed!");
            return;
        }

        memcpy(full_packet, header, sizeof(header));
        uint8_t *payload = full_packet + sizeof(header);

        // Custom header
        payload[0] = type;
        payload[1] = seq >> 8;
        payload[2] = seq & 0xFF;
        payload[3] = total_chunks >> 8;
        payload[4] = total_chunks & 0xFF;
        payload[5] = 'P';
        payload[6] = 'l';
        payload[7] = 'e';
        payload[8] = 'c';
        payload[9] = 'y';
        payload[10] = 'C';

        memcpy(payload + 11, data + offset, chunk_size);

        ESP_LOGI(TAG, "TxTarget: %d bytes", full_len);
        full_packet[full_len] = 0;
        esp_err_t ret = tx_with_retry(WIFI_IF, full_packet, full_len + 2, true);
        // ESP_LOG_BUFFER_HEXDUMP(TAG, full_packet, full_len, ESP_LOG_INFO);

        if (ret == ESP_OK)
        {
            ESP_LOGI(TAG, "Sent chunk %d/%d, size: %d bytes", seq + 1, total_chunks, chunk_size);
        }
        else
        {
            ESP_LOGE(TAG, "Failed to send chunk %d, %s", seq, esp_err_to_name(ret));
        }

        free(full_packet);
        vTaskDelay(80 / portTICK_PERIOD_MS);
        offset += chunk_size;
    }
}

esp_err_t camera_init()
{
    // initialize the camera
    esp_err_t err = esp_camera_init(&camera_config);
    if (err != ESP_OK)
    {
        ESP_LOGE(TAG, "Camera Init Failed");
        return err;
    }

    return ESP_OK;
}

typedef struct
{
    bool isVisible;
    uint32_t len;
    uint16_t width;
    uint16_t height;
    pixformat_t format;
    struct timeval timestamp;
} PreData;

uint8_t *create_frame(const uint8_t *data, size_t *outSize, PreData pre)
{

    const size_t headerSize = 19;
    const size_t totalSize = pre.len + headerSize;

    ESP_LOGI(TAG, "DataWithFrame: %d\n", totalSize);

    uint8_t *dataOut = malloc(totalSize);
    if (!dataOut)
    {
        ESP_LOGE(TAG, "Failed to allocate memory for frame");
        return NULL;
    }

    size_t bytePos = 0;

    int64_t time_us = (int64_t)pre.timestamp.tv_sec * 1000000L + (int64_t)pre.timestamp.tv_usec;

    dataOut[bytePos++] = pre.isVisible;
    dataOut[bytePos++] = (pre.len >> 24) & 0xFF;
    dataOut[bytePos++] = (pre.len >> 16) & 0xFF;
    dataOut[bytePos++] = (pre.len >> 8) & 0xFF;
    dataOut[bytePos++] = (pre.len) & 0xFF;
    dataOut[bytePos++] = (pre.width >> 8) & 0xFF;
    dataOut[bytePos++] = (pre.width) & 0xFF;
    dataOut[bytePos++] = (pre.height >> 8) & 0xFF;
    dataOut[bytePos++] = (pre.height) & 0xFF;
    dataOut[bytePos++] = (pre.format) & 0xFF;
    dataOut[bytePos++] = (time_us >> 56) & 0xFF;
    dataOut[bytePos++] = (time_us >> 48) & 0xFF;
    dataOut[bytePos++] = (time_us >> 40) & 0xFF;
    dataOut[bytePos++] = (time_us >> 32) & 0xFF;
    dataOut[bytePos++] = (time_us >> 24) & 0xFF;
    dataOut[bytePos++] = (time_us >> 16) & 0xFF;
    dataOut[bytePos++] = (time_us >> 8) & 0xFF;
    dataOut[bytePos++] = (time_us) & 0xFF;

    ESP_LOGI(TAG, "Vis: %d, len: %ld, width:%d, height:%d, format:%d, timestamp: %lld", pre.isVisible, pre.len, pre.width, pre.height, pre.format, time_us);
    // Copy image data after PreData
    memcpy(dataOut + bytePos, data, pre.len);

    *outSize = totalSize;
    return dataOut;
}

esp_err_t camera_capture()
{
    ESP_LOGI(TAG, "Pregrab");
    camera_fb_t *fb = esp_camera_fb_get();
    ESP_LOGI(TAG, "Postgrab");

    if (!fb)
    {
        ESP_LOGE(TAG, "Camera Capture Failed");
        return ESP_FAIL;
    }

    // Fill PreData at the start
    PreData pre;
    pre.isVisible = true;
    pre.len = fb->len;
    pre.width = fb->width;
    pre.height = fb->height;
    pre.format = fb->format;
    pre.timestamp = fb->timestamp;

    size_t totalSize;
    uint8_t *dataOut = create_frame(fb->buf, &totalSize, pre);
    if (!dataOut)
    {
        esp_camera_fb_return(fb);
        return ESP_ERR_NO_MEM;
    }

    esp_camera_fb_return(fb);
    send_data_frame(0x10, dataOut, totalSize);

    free(dataOut);
    return ESP_OK;
}

uint8_t *take_picture()
{
    gpio_set_level(LEP_CSN_PIN, 0);
    vTaskDelay(pdMS_TO_TICKS(10));

    while (1)
    {
        // Spin waiting for vsync to be asserted
        int t = 0;
        while (gpio_get_level(LEP_VSYNC_PIN) == 0 && t++ < 11000)
        {
        }

        vsyncDetectedUsec = esp_timer_get_time();

        // Attempt to process a segment
        if (t >= 22000 || vospi_transfer_segment(vsyncDetectedUsec))
        {
            // Got image
            vsync_count = 0;
            if (t >= 11000)
                printf("LEPOVF%d\n", t);
            else
            {
                printf("LEPNOVF%d\n", t);
            }

            // Copy the frame to the current half of the shared buffer and let rsp_task know
            vospi_get_frame(&rsp_lep_buffer);

            // Hold fault counters reset while operating
            sync_fail_count = 0;

            // disable lepton again
            gpio_set_level(LEP_CSN_PIN, 1);
            return rsp_lep_buffer.lep_bufferP;
        }
        else
        {
            if (++vsync_count >= 36)
            {
                vsync_count = 0;
                printf("[MAIN] Could not get lepton image\n");

                vTaskDelay(pdMS_TO_TICKS(185));

                if (sync_fail_count++ == LEP_SYNC_FAIL_FAULT_LIMIT)
                {
                    return NULL;
                }
            }
        }
    }

    return NULL;
}

uint8_t *lepton_get_frame()
{

    uint8_t *picture = take_picture();
    struct timeval tv_now;
    gettimeofday(&tv_now, NULL);

    PreData pre;
    pre.isVisible = false;
    pre.len = LEP_NUM_PIXELS;
    pre.width = LEP_WIDTH;
    pre.height = LEP_HEIGHT;
    pre.format = 0;
    pre.timestamp = tv_now;
    size_t totalSize;
    uint8_t *dataOut = create_frame(picture, &totalSize, pre);
    send_data_frame(0x10, dataOut, totalSize);
    return picture;
}

void app_main(void)
{

    // for (size_t i = 0; i < sizeof(hehe); i++)
    // {
    //     hehe[i] = i & 0xFF;
    // }

    const uart_port_t uart_num = UART_NUM_0;
    uart_config_t uart_config = {
        .baud_rate = 115200,
        .data_bits = UART_DATA_8_BITS,
        .parity = UART_PARITY_DISABLE,
        .stop_bits = UART_STOP_BITS_1,
        .flow_ctrl = UART_HW_FLOWCTRL_DISABLE};
    uart_param_config(uart_num, &uart_config);
    uart_driver_install(uart_num, BUF_SIZE * 2, 0, 0, NULL, 0);

    uint8_t *buf = malloc(BUF_SIZE);
    char *line = malloc(BUF_SIZE);
    int line_len = 0;

    nvs_flash_init();
    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());
    esp_netif_create_default_wifi_sta();

    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    cfg.static_tx_buf_num = 32;  // default is 4 or 8
    cfg.dynamic_tx_buf_num = 32; // boost this too
    cfg.ampdu_tx_enable = 1;
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));
    ESP_ERROR_CHECK(esp_wifi_set_storage(WIFI_STORAGE_RAM));
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_AP));

    wifi_config_t ap_config = {
        .ap = {
            .ssid = "ESP32-Test",
            .channel = 10,
            .authmode = WIFI_AUTH_OPEN,
            .max_connection = 1,
            .ssid_hidden = 1,
            .beacon_interval = 60000}};
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_AP, &ap_config));
    ESP_ERROR_CHECK(esp_wifi_start());
    ESP_ERROR_CHECK(esp_wifi_set_ps(WIFI_PS_NONE));

    esp_log_level_set("wifi", ESP_LOG_DEBUG);

    uint8_t mac_addr[6];
    esp_read_mac(mac_addr, ESP_MAC_WIFI_SOFTAP);
    ESP_LOGI(TAG, "ESP32 MAC: %02X:%02X:%02X:%02X:%02X:%02X",
             mac_addr[0], mac_addr[1], mac_addr[2], mac_addr[3], mac_addr[4], mac_addr[5]);

    // size_t outLen = 0;

    // PreData pre;

    // struct timeval tv;
    // gettimeofday(&tv, NULL);

    // pre.format = PIXFORMAT_JPEG;
    // pre.width = 100;
    // pre.height = 100;
    // pre.isVisible = true;
    // pre.len = sizeof(hehe);
    // pre.timestamp = tv;

    // uint8_t *imgFrame = create_frame(hehe, &outLen, pre);

    // send_data_frame(0x10, imgFrame, outLen);

    // free(imgFrame);

#ifdef LEPTON_ENABLE
    gpio_config_t usb_phy_conf = {
        .pin_bit_mask = (1ULL << 19) | (1ULL << 20),
        .mode = GPIO_MODE_OUTPUT,
        .pull_up_en = GPIO_PULLUP_ENABLE,
        .pull_down_en = GPIO_PULLDOWN_DISABLE,
        .intr_type = GPIO_INTR_DISABLE,
    };
    gpio_config(&usb_phy_conf);

    printf("[MAIN] Start task\n");

    LEP_CONFIG.agc_set_enabled = true;
    LEP_CONFIG.emissivity = 100;
    LEP_CONFIG.gain_mode = LEP_GAIN_AUTO;
    // initialize spi, i2c and gpio for lepton
    if (!lepton_io_init())
    {
        printf("[MAIN] Error: I/O init failed");
        while (1)
        {
            vTaskDelay(pdMS_TO_TICKS(100));
        }
    }

    // allocate lepton buffers
    if (!lepton_buffer_init())
    {
        printf("[MAIN] Error: Memory init failed");
        while (1)
        {
            vTaskDelay(pdMS_TO_TICKS(100));
        }
    }

    // wait for lepton to initialize
    vTaskDelay(pdMS_TO_TICKS(1000));
    ;

    // Attempt to initialize the VoSPI interface
    if (vospi_init() != ESP_OK)
    {
        printf("[MAIN] Error: Lepton VoSPI initialization failed\n");
        while (1)
        {
            vTaskDelay(pdMS_TO_TICKS(100));
        }
    }

    // initialize lepton
    if (!lepton_init())
    {
        printf("[MAIN] Error: Lepton CCI initialization failed\n");
        while (1)
        {
            vTaskDelay(pdMS_TO_TICKS(100));
        }
    }

    // disable data from lepton
    gpio_set_level(LEP_CSN_PIN, 1);

    printf("LEPTON INIT DONE");
    camera_init();

#endif

    while (1)
    {

        int64_t start = esp_timer_get_time();
        camera_capture();
#ifdef LEPTON_ENABLE
        uint8_t *picture = lepton_get_frame();
        if (picture)
        {
            int maxBrightness = -1;
            Rectangle brightest = {0, 0, 0, 0};
        
            for (int y = 0; y < LEP_HEIGHT; y++) {
                for (int x = 0; x < LEP_WIDTH; x++) {
                    int idx = y * LEP_WIDTH + x;
                    int brightness = picture[idx];  // grayscale
        
                    if (brightness > maxBrightness) {
                        maxBrightness = brightness;
                        brightest.x = x;
                        brightest.y = y;
                    }
                }
            }
        
            // Rozšírenie oblasti
            int padding = 10;
            brightest.x = (brightest.x - padding < 0) ? 0 : brightest.x - padding;
            brightest.y = (brightest.y - padding < 0) ? 0 : brightest.y - padding;
            brightest.width = ((brightest.x + padding * 2) > LEP_WIDTH) ? (LEP_WIDTH - brightest.x) : (padding * 2);
            brightest.height = ((brightest.y + padding * 2) > LEP_HEIGHT) ? (LEP_HEIGHT - brightest.y) : (padding * 2);
        
            // Výpočet stredu oblasti
            double middleX = brightest.x + brightest.width / 2.0;
            double middleY = brightest.y + brightest.height / 2.0;
        
            // Výpočet uholných odchýlok
            double rotationX = 3 * M_PI / 180.0;
            double rotationY = 0 * M_PI / 180.0;
            double rotationZ = 45 * M_PI / 180.0;
            double posX = 100, posY = 100, posZ = 500;
        
            double plusX = ((middleX - LEP_HEIGHT / 2.0) / LEP_HEIGHT) * 57 * M_PI / 180.0;
            double plusY = ((middleY - LEP_WIDTH / 2.0) / LEP_WIDTH) * 42.75 * M_PI / 180.0;
        
            double posGX = posZ * tan(rotationX + plusX) * cos(rotationZ) + posX;
            double posGY = posZ * tan(rotationY + plusY) * sin(rotationZ) + posY;
        
            //printf("Stred najjasnejšej oblasti: (%.2f, %.2f)\n", middleX, middleY);
            //printf("Výpočet pozície: posGX = %.2f, posGY = %.2f\n", posGX, posGY);

            uint8_t messageBuffer[32];

            uint8_t messageBufferEncoded[50];
            size_t outputLen;
            mbedtls_base64_encode(messageBufferEncoded, sizeof(messageBufferEncoded), &outputLen, messageBuffer, sizeof(messageBuffer));

            uint8_t messageBufferPrefixed[60] =  "NavData:";
            strcat(messageBufferPrefixed, messageBufferEncoded);
            strcat(messageBufferPrefixed, "\n");

            uart_write_bytes(uart_num, messageBufferPrefixed, strlen(messageBufferPrefixed));
        
        }
        else
        {
            printf("Failed to lepton take picture\n");
        }
#endif
        int64_t end = esp_timer_get_time();

        int64_t duration = end - start;

        printf("%lf FPS, %lf SPF\n", 1 / (duration / 1000000.0), duration / 1000000.0);
    }
}