#include "low_level_display_st7796.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <cstdlib>  // For abs()

// RGB565 color definitions
#define COLOR_BLACK   0x0000
#define COLOR_WHITE   0xFFFF

LowLevelDisplayST7796::LowLevelDisplayST7796(const st7796_config* cfg, int w, int h, bool invert)
    : config(cfg), width(w), height(h), initialized(false), rgb_buffer(nullptr), invert_color(invert),
      prev_bit_buffer(nullptr), dirty_rect_enabled(true) {
    for (int i = 0; i < MAX_DIRTY_RECTS; i++) {
        dirty_rects[i].reset();
    }
}

LowLevelDisplayST7796::~LowLevelDisplayST7796() {
    if (rgb_buffer) {
        free(rgb_buffer);
        rgb_buffer = nullptr;
    }
    if (prev_bit_buffer) {
        free(prev_bit_buffer);
        prev_bit_buffer = nullptr;
    }
}

bool LowLevelDisplayST7796::init() {
    if (initialized) {
        return true;
    }
    
    st7796_init(config, width, height);
    
    // Allocate RGB565 buffer once (reused for all draw operations)
    size_t buffer_size = width * height * sizeof(uint16_t);
    rgb_buffer = (uint16_t *)malloc(buffer_size);
    if (!rgb_buffer) {
        printf("Error: Failed to allocate %zu bytes for RGB buffer\n", buffer_size);
        return false;
    }
    printf("ST7796 display initialized: %dx%d (RGB buffer: %zu bytes)\n", width, height, buffer_size);
    
    initialized = true;
    return true;
}

void LowLevelDisplayST7796::clear(bool white) {
    bool out_white = invert_color ? !white : white;
    st7796_fill(out_white ? COLOR_WHITE : COLOR_BLACK);
}

void LowLevelDisplayST7796::draw_pixel(int x, int y, bool white) {
    bool out_white = invert_color ? !white : white;
    st7796_draw_pixel(x, y, out_white ? COLOR_WHITE : COLOR_BLACK);
}

void LowLevelDisplayST7796::draw_buffer(const uint8_t* bit_buffer) {
    if (!bit_buffer || !rgb_buffer) return;

    // Calculate buffer size
    size_t bit_buffer_size = (width * height + 7) / 8;

    // If dirty rectangle tracking is enabled and we have a previous buffer
    if (dirty_rect_enabled && prev_bit_buffer) {
        // Reset all dirty rectangles
        for (int i = 0; i < MAX_DIRTY_RECTS; i++) {
            dirty_rects[i].reset();
        }

        // Split screen into 4 quadrants
        int mid_x = width / 2;
        int mid_y = height / 2;

        // Use bitwise XOR to quickly detect changed bytes
        for (size_t byte_idx = 0; byte_idx < bit_buffer_size; byte_idx++) {
            uint8_t diff = bit_buffer[byte_idx] ^ prev_bit_buffer[byte_idx];

            // If this byte has changes
            if (diff != 0) {
                // Calculate pixel coordinates for this byte
                int pixel_idx = byte_idx * 8;
                int base_x = pixel_idx % width;
                int base_y = pixel_idx / width;

                // Check each changed bit/pixel in this byte
                for (int bit = 0; bit < 8 && (pixel_idx + bit) < (width * height); bit++) {
                    if (diff & (0x80 >> bit)) {
                        int x = base_x + bit;
                        int y = base_y;

                        // Adjust coordinates if we wrapped to next row
                        if (x >= width) {
                            x -= width;
                            y++;
                        }

                        // Route to appropriate quadrant based on X and Y position
                        // Quadrant 0: Top-left (x < mid_x, y < mid_y)
                        // Quadrant 1: Top-right (x >= mid_x, y < mid_y)
                        // Quadrant 2: Bottom-left (x < mid_x, y >= mid_y)
                        // Quadrant 3: Bottom-right (x >= mid_x, y >= mid_y)
                        int rect_idx = ((y >= mid_y) ? 2 : 0) + ((x >= mid_x) ? 1 : 0);
                        dirty_rects[rect_idx].expand(x, y);
                    }
                }
            }
        }

        // Check if we have any valid dirty rectangles
        int valid_rects = 0;
        for (int i = 0; i < MAX_DIRTY_RECTS; i++) {
            if (dirty_rects[i].is_valid) {
                valid_rects++;
            }
        }

        // If there are no changes, skip the update
        if (valid_rects == 0) {
            return;
        }

        // Optimization: Merge adjacent rectangles if beneficial
        // Check pairs of rectangles and merge if they overlap or are close
        if (valid_rects >= 2) {
            // Try merging adjacent quadrants
            // Check top row (0,1) merge
            if (dirty_rects[0].is_valid && dirty_rects[1].is_valid) {
                int gap_x = dirty_rects[1].x0 - dirty_rects[0].x1;
                int gap_y = abs(dirty_rects[0].y0 - dirty_rects[1].y0) + abs(dirty_rects[0].y1 - dirty_rects[1].y1);

                if (gap_x < 30 && gap_y < 20) {
                    dirty_rects[0].merge(dirty_rects[1]);
                    dirty_rects[1].reset();
                    valid_rects--;
                }
            }

            // Check bottom row (2,3) merge
            if (dirty_rects[2].is_valid && dirty_rects[3].is_valid) {
                int gap_x = dirty_rects[3].x0 - dirty_rects[2].x1;
                int gap_y = abs(dirty_rects[2].y0 - dirty_rects[3].y0) + abs(dirty_rects[2].y1 - dirty_rects[3].y1);

                if (gap_x < 30 && gap_y < 20) {
                    dirty_rects[2].merge(dirty_rects[3]);
                    dirty_rects[3].reset();
                    valid_rects--;
                }
            }

            // Check left column (0,2) merge
            if (dirty_rects[0].is_valid && dirty_rects[2].is_valid) {
                int gap_y = dirty_rects[2].y0 - dirty_rects[0].y1;
                int gap_x = abs(dirty_rects[0].x0 - dirty_rects[2].x0) + abs(dirty_rects[0].x1 - dirty_rects[2].x1);

                if (gap_y < 30 && gap_x < 20) {
                    dirty_rects[0].merge(dirty_rects[2]);
                    dirty_rects[2].reset();
                    valid_rects--;
                }
            }

            // Check right column (1,3) merge
            if (dirty_rects[1].is_valid && dirty_rects[3].is_valid) {
                int gap_y = dirty_rects[3].y0 - dirty_rects[1].y1;
                int gap_x = abs(dirty_rects[1].x0 - dirty_rects[3].x0) + abs(dirty_rects[1].x1 - dirty_rects[3].x1);

                if (gap_y < 30 && gap_x < 20) {
                    dirty_rects[1].merge(dirty_rects[3]);
                    dirty_rects[3].reset();
                    valid_rects--;
                }
            }

            // Final pass: merge any remaining valid rectangles if they're very close
            for (int i = 0; i < MAX_DIRTY_RECTS - 1; i++) {
                if (!dirty_rects[i].is_valid) continue;
                for (int j = i + 1; j < MAX_DIRTY_RECTS; j++) {
                    if (!dirty_rects[j].is_valid) continue;

                    DirtyRect merged = dirty_rects[i];
                    merged.merge(dirty_rects[j]);

                    int combined_area = dirty_rects[i].get_area() + dirty_rects[j].get_area();
                    int merged_area = merged.get_area();

                    // Merge if the combined overhead is less than 40%
                    if (merged_area < combined_area * 1.4f) {
                        dirty_rects[i] = merged;
                        dirty_rects[j].reset();
                        valid_rects--;
                        break;  // Move to next i
                    }
                }
            }
        }

        // Copy current buffer to previous buffer for next frame comparison
        memcpy(prev_bit_buffer, bit_buffer, bit_buffer_size);

        // Process each valid dirty rectangle
        for (int rect_idx = 0; rect_idx < MAX_DIRTY_RECTS; rect_idx++) {
            if (!dirty_rects[rect_idx].is_valid) continue;

            DirtyRect& rect = dirty_rects[rect_idx];

            // Convert only the dirty rectangle region to RGB565
            for (int y = rect.y0; y <= rect.y1; y++) {
                for (int x = rect.x0; x <= rect.x1; x++) {
                    int byte_index = (y * width + x) / 8;
                    int bit_index = 7 - (x % 8);
                    bool pixel_white = (bit_buffer[byte_index] >> bit_index) & 0x01;
                    bool out_white = invert_color ? !pixel_white : pixel_white;
                    rgb_buffer[y * width + x] = out_white ? COLOR_WHITE : COLOR_BLACK;
                }
            }

            // Draw only this dirty rectangle
            st7796_set_window(rect.x0, rect.y0, rect.x1, rect.y1);

            // Calculate size of dirty region
            int dirty_width = rect.get_width();
            int dirty_height = rect.get_height();

            // Write only the dirty rectangle pixels
            // We need to extract rows from the full rgb_buffer
            for (int row = 0; row < dirty_height; row++) {
                int buffer_offset = (rect.y0 + row) * width + rect.x0;
                st7796_write_raw((const uint8_t*)&rgb_buffer[buffer_offset], dirty_width * 2);
            }
        }

    } else {
        // Full screen update (original behavior)
        // Convert 1-bit buffer to RGB565 using persistent buffer
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                int byte_index = (y * width + x) / 8;
                int bit_index = 7 - (x % 8);
                bool pixel_white = (bit_buffer[byte_index] >> bit_index) & 0x01;
                bool out_white = invert_color ? !pixel_white : pixel_white;
                rgb_buffer[y * width + x] = out_white ? COLOR_WHITE : COLOR_BLACK;
            }
        }
        // Draw entire buffer at once
        st7796_set_cursor(0, 0);
        // Use raw write for speed.
        // Since we only use 0x0000 (Black) and 0xFFFF (White), endianness doesn't matter.
        // 0x0000 -> 0x00, 0x00 (LE) -> Display sees 0x00, 0x00 (0x0000 correct)
        // 0xFFFF -> 0xFF, 0xFF (LE) -> Display sees 0xFF, 0xFF (0xFFFF correct)
        st7796_write_raw((const uint8_t*)rgb_buffer, width * height * 2);

        // If dirty rect is enabled, store this buffer for next comparison
        if (dirty_rect_enabled && prev_bit_buffer) {
            memcpy(prev_bit_buffer, bit_buffer, bit_buffer_size);
        }
    }
}

void LowLevelDisplayST7796::refresh() {
    // ST7796 updates immediately, no refresh needed
}

void LowLevelDisplayST7796::set_backlight(bool on) {
    // Use brightness control: on = 100%, off = 0%
    st7796_set_brightness(on ? 100 : 0);
}

void LowLevelDisplayST7796::set_brightness(uint8_t brightness) {
    st7796_set_brightness(brightness);
}

uint8_t LowLevelDisplayST7796::get_brightness() const {
    return st7796_get_brightness();
}

void LowLevelDisplayST7796::sleep() {
    st7796_sleep();
}

void LowLevelDisplayST7796::wake() {
    st7796_wake();
}

void LowLevelDisplayST7796::set_rotation(uint8_t rotation) {
    // ST7796 driver doesn't have rotation control yet
    // TODO: Add MADCTL register manipulation for rotation
    (void)rotation;
}

void LowLevelDisplayST7796::enable_dirty_rect(bool enabled) {
    dirty_rect_enabled = enabled;

    if (enabled && !prev_bit_buffer) {
        // Allocate buffer to store previous frame for change detection
        size_t bit_buffer_size = (width * height + 7) / 8;  // 1 bit per pixel
        prev_bit_buffer = (uint8_t *)malloc(bit_buffer_size);
        if (prev_bit_buffer) {
            // Initialize to all zeros (black screen)
            memset(prev_bit_buffer, 0, bit_buffer_size);
            printf("ST7796: Dirty rectangle tracking enabled (buffer: %zu bytes, max rects: %d)\n",
                   bit_buffer_size, MAX_DIRTY_RECTS);
        } else {
            printf("Error: Failed to allocate %zu bytes for dirty rect buffer\n", bit_buffer_size);
            dirty_rect_enabled = false;
        }
    } else if (!enabled && prev_bit_buffer) {
        // Disable and free tracking buffer
        free(prev_bit_buffer);
        prev_bit_buffer = nullptr;
        for (int i = 0; i < MAX_DIRTY_RECTS; i++) {
            dirty_rects[i].reset();
        }
        printf("ST7796: Dirty rectangle tracking disabled\n");
    }
}

void LowLevelDisplayST7796::on_idle_2min() {
    if (!is_dimmed && !is_sleeping) {
        saved_brightness = get_brightness();
        set_brightness(5); // Dim to 5%
        is_dimmed = true;
        printf("TFT: Dimmed to 5%%\n");
    }
}

void LowLevelDisplayST7796::on_idle_10min() {
    if (!is_sleeping) {
        sleep();
        is_sleeping = true;
        is_dimmed = true; // Sleep implies dimmed
        printf("TFT: Entered sleep mode\n");
    }
}

void LowLevelDisplayST7796::on_user_interaction() {
    if (is_sleeping) {
        wake();
        // Restore brightness if we have a saved value, or default to 100
        set_brightness(saved_brightness > 0 ? saved_brightness : 100);
        is_sleeping = false;
        is_dimmed = false;
        printf("TFT: Woke from sleep\n");
    } else if (is_dimmed) {
        set_brightness(saved_brightness > 0 ? saved_brightness : 100);
        is_dimmed = false;
        printf("TFT: Restored brightness\n");
    }
}