#include <SPI.h>
#include <Adafruit_GFX.h>
#include <Adafruit_ILI9341.h>
#include <math.h>
#include "arched_digits.h"

#define TFT_DC 2
#define TFT_CS 15

const int SCREEN_WIDTH = 320;
const int SCREEN_HEIGHT = 240;

Adafruit_ILI9341 display = Adafruit_ILI9341(TFT_CS, TFT_DC);

// Helper to get digit bitmap pointer and dimensions
typedef struct {
  const uint8_t* data;
  uint8_t width;
  uint8_t height;
  uint16_t pixels;
} DigitInfo;

// Get cyan color with specified brightness level (0-3)
// Pre-dithered levels: 0=black, 1=dim (85), 2=medium (170), 3=bright (255)
const uint8_t BRIGHTNESS_LEVELS[4] = {0, 85, 170, 255};

uint16_t getCyanBrightness(uint8_t level) {
  uint8_t brightness = BRIGHTNESS_LEVELS[level & 0x03];
  uint8_t r = 0;
  uint8_t g = brightness;
  uint8_t b = brightness;
  
  // Convert to RGB565
  return ((r & 0xF8) << 8) | ((g & 0xFC) << 3) | (b >> 3);
}

DigitInfo getDigit2d(int pos, char digit) {
  // 2-digit format positions
  DigitInfo result;
  if (pos == 0) { // tens
    if (digit == '8') {
      result.data = DIGIT_2d_pos0_8_D; result.width = DIGIT_2d_pos0_8_W;
      result.height = DIGIT_2d_pos0_8_H; result.pixels = DIGIT_2d_pos0_8_P;
      return result;
    }
    if (digit == '9') {
      result.data = DIGIT_2d_pos0_9_D; result.width = DIGIT_2d_pos0_9_W;
      result.height = DIGIT_2d_pos0_9_H; result.pixels = DIGIT_2d_pos0_9_P;
      return result;
    }
  } else if (pos == 1) { // ones
    switch(digit) {
      case '0': result.data = DIGIT_2d_pos1_0_D; result.width = DIGIT_2d_pos1_0_W; result.height = DIGIT_2d_pos1_0_H; result.pixels = DIGIT_2d_pos1_0_P; return result;
      case '1': result.data = DIGIT_2d_pos1_1_D; result.width = DIGIT_2d_pos1_1_W; result.height = DIGIT_2d_pos1_1_H; result.pixels = DIGIT_2d_pos1_1_P; return result;
      case '2': result.data = DIGIT_2d_pos1_2_D; result.width = DIGIT_2d_pos1_2_W; result.height = DIGIT_2d_pos1_2_H; result.pixels = DIGIT_2d_pos1_2_P; return result;
      case '3': result.data = DIGIT_2d_pos1_3_D; result.width = DIGIT_2d_pos1_3_W; result.height = DIGIT_2d_pos1_3_H; result.pixels = DIGIT_2d_pos1_3_P; return result;
      case '4': result.data = DIGIT_2d_pos1_4_D; result.width = DIGIT_2d_pos1_4_W; result.height = DIGIT_2d_pos1_4_H; result.pixels = DIGIT_2d_pos1_4_P; return result;
      case '5': result.data = DIGIT_2d_pos1_5_D; result.width = DIGIT_2d_pos1_5_W; result.height = DIGIT_2d_pos1_5_H; result.pixels = DIGIT_2d_pos1_5_P; return result;
      case '6': result.data = DIGIT_2d_pos1_6_D; result.width = DIGIT_2d_pos1_6_W; result.height = DIGIT_2d_pos1_6_H; result.pixels = DIGIT_2d_pos1_6_P; return result;
      case '7': result.data = DIGIT_2d_pos1_7_D; result.width = DIGIT_2d_pos1_7_W; result.height = DIGIT_2d_pos1_7_H; result.pixels = DIGIT_2d_pos1_7_P; return result;
      case '8': result.data = DIGIT_2d_pos1_8_D; result.width = DIGIT_2d_pos1_8_W; result.height = DIGIT_2d_pos1_8_H; result.pixels = DIGIT_2d_pos1_8_P; return result;
      case '9': result.data = DIGIT_2d_pos1_9_D; result.width = DIGIT_2d_pos1_9_W; result.height = DIGIT_2d_pos1_9_H; result.pixels = DIGIT_2d_pos1_9_P; return result;
    }
  } else if (pos == 3) { // decimal
    switch(digit) {
      case '0': result.data = DIGIT_2d_pos3_0_D; result.width = DIGIT_2d_pos3_0_W; result.height = DIGIT_2d_pos3_0_H; result.pixels = DIGIT_2d_pos3_0_P; return result;
      case '1': result.data = DIGIT_2d_pos3_1_D; result.width = DIGIT_2d_pos3_1_W; result.height = DIGIT_2d_pos3_1_H; result.pixels = DIGIT_2d_pos3_1_P; return result;
      case '2': result.data = DIGIT_2d_pos3_2_D; result.width = DIGIT_2d_pos3_2_W; result.height = DIGIT_2d_pos3_2_H; result.pixels = DIGIT_2d_pos3_2_P; return result;
      case '3': result.data = DIGIT_2d_pos3_3_D; result.width = DIGIT_2d_pos3_3_W; result.height = DIGIT_2d_pos3_3_H; result.pixels = DIGIT_2d_pos3_3_P; return result;
      case '4': result.data = DIGIT_2d_pos3_4_D; result.width = DIGIT_2d_pos3_4_W; result.height = DIGIT_2d_pos3_4_H; result.pixels = DIGIT_2d_pos3_4_P; return result;
      case '5': result.data = DIGIT_2d_pos3_5_D; result.width = DIGIT_2d_pos3_5_W; result.height = DIGIT_2d_pos3_5_H; result.pixels = DIGIT_2d_pos3_5_P; return result;
      case '6': result.data = DIGIT_2d_pos3_6_D; result.width = DIGIT_2d_pos3_6_W; result.height = DIGIT_2d_pos3_6_H; result.pixels = DIGIT_2d_pos3_6_P; return result;
      case '7': result.data = DIGIT_2d_pos3_7_D; result.width = DIGIT_2d_pos3_7_W; result.height = DIGIT_2d_pos3_7_H; result.pixels = DIGIT_2d_pos3_7_P; return result;
      case '8': result.data = DIGIT_2d_pos3_8_D; result.width = DIGIT_2d_pos3_8_W; result.height = DIGIT_2d_pos3_8_H; result.pixels = DIGIT_2d_pos3_8_P; return result;
      case '9': result.data = DIGIT_2d_pos3_9_D; result.width = DIGIT_2d_pos3_9_W; result.height = DIGIT_2d_pos3_9_H; result.pixels = DIGIT_2d_pos3_9_P; return result;
    }
  }
  result.data = NULL; result.width = 0; result.height = 0; result.pixels = 0;
  return result;
}

DigitInfo getDigit3d(int pos, char digit) {
  // 3-digit format positions
  DigitInfo result;
  if (pos == 0) {
    result.data = DIGIT_3d_pos0_1_D; result.width = DIGIT_3d_pos0_1_W;
    result.height = DIGIT_3d_pos0_1_H; result.pixels = DIGIT_3d_pos0_1_P;
    return result;
  }
  if (pos == 1) {
    result.data = DIGIT_3d_pos1_0_D; result.width = DIGIT_3d_pos1_0_W;
    result.height = DIGIT_3d_pos1_0_H; result.pixels = DIGIT_3d_pos1_0_P;
    return result;
  }
  if (pos == 2) { // ones (0-8)
    switch(digit) {
      case '0': result.data = DIGIT_3d_pos2_0_D; result.width = DIGIT_3d_pos2_0_W; result.height = DIGIT_3d_pos2_0_H; result.pixels = DIGIT_3d_pos2_0_P; return result;
      case '1': result.data = DIGIT_3d_pos2_1_D; result.width = DIGIT_3d_pos2_1_W; result.height = DIGIT_3d_pos2_1_H; result.pixels = DIGIT_3d_pos2_1_P; return result;
      case '2': result.data = DIGIT_3d_pos2_2_D; result.width = DIGIT_3d_pos2_2_W; result.height = DIGIT_3d_pos2_2_H; result.pixels = DIGIT_3d_pos2_2_P; return result;
      case '3': result.data = DIGIT_3d_pos2_3_D; result.width = DIGIT_3d_pos2_3_W; result.height = DIGIT_3d_pos2_3_H; result.pixels = DIGIT_3d_pos2_3_P; return result;
      case '4': result.data = DIGIT_3d_pos2_4_D; result.width = DIGIT_3d_pos2_4_W; result.height = DIGIT_3d_pos2_4_H; result.pixels = DIGIT_3d_pos2_4_P; return result;
      case '5': result.data = DIGIT_3d_pos2_5_D; result.width = DIGIT_3d_pos2_5_W; result.height = DIGIT_3d_pos2_5_H; result.pixels = DIGIT_3d_pos2_5_P; return result;
      case '6': result.data = DIGIT_3d_pos2_6_D; result.width = DIGIT_3d_pos2_6_W; result.height = DIGIT_3d_pos2_6_H; result.pixels = DIGIT_3d_pos2_6_P; return result;
      case '7': result.data = DIGIT_3d_pos2_7_D; result.width = DIGIT_3d_pos2_7_W; result.height = DIGIT_3d_pos2_7_H; result.pixels = DIGIT_3d_pos2_7_P; return result;
      case '8': result.data = DIGIT_3d_pos2_8_D; result.width = DIGIT_3d_pos2_8_W; result.height = DIGIT_3d_pos2_8_H; result.pixels = DIGIT_3d_pos2_8_P; return result;
    }
  } else if (pos == 4) { // decimal
    switch(digit) {
      case '0': result.data = DIGIT_3d_pos4_0_D; result.width = DIGIT_3d_pos4_0_W; result.height = DIGIT_3d_pos4_0_H; result.pixels = DIGIT_3d_pos4_0_P; return result;
      case '1': result.data = DIGIT_3d_pos4_1_D; result.width = DIGIT_3d_pos4_1_W; result.height = DIGIT_3d_pos4_1_H; result.pixels = DIGIT_3d_pos4_1_P; return result;
      case '2': result.data = DIGIT_3d_pos4_2_D; result.width = DIGIT_3d_pos4_2_W; result.height = DIGIT_3d_pos4_2_H; result.pixels = DIGIT_3d_pos4_2_P; return result;
      case '3': result.data = DIGIT_3d_pos4_3_D; result.width = DIGIT_3d_pos4_3_W; result.height = DIGIT_3d_pos4_3_H; result.pixels = DIGIT_3d_pos4_3_P; return result;
      case '4': result.data = DIGIT_3d_pos4_4_D; result.width = DIGIT_3d_pos4_4_W; result.height = DIGIT_3d_pos4_4_H; result.pixels = DIGIT_3d_pos4_4_P; return result;
      case '5': result.data = DIGIT_3d_pos4_5_D; result.width = DIGIT_3d_pos4_5_W; result.height = DIGIT_3d_pos4_5_H; result.pixels = DIGIT_3d_pos4_5_P; return result;
      case '6': result.data = DIGIT_3d_pos4_6_D; result.width = DIGIT_3d_pos4_6_W; result.height = DIGIT_3d_pos4_6_H; result.pixels = DIGIT_3d_pos4_6_P; return result;
      case '7': result.data = DIGIT_3d_pos4_7_D; result.width = DIGIT_3d_pos4_7_W; result.height = DIGIT_3d_pos4_7_H; result.pixels = DIGIT_3d_pos4_7_P; return result;
      case '8': result.data = DIGIT_3d_pos4_8_D; result.width = DIGIT_3d_pos4_8_W; result.height = DIGIT_3d_pos4_8_H; result.pixels = DIGIT_3d_pos4_8_P; return result;
      case '9': result.data = DIGIT_3d_pos4_9_D; result.width = DIGIT_3d_pos4_9_W; result.height = DIGIT_3d_pos4_9_H; result.pixels = DIGIT_3d_pos4_9_P; return result;
    }
  }
  result.data = NULL; result.width = 0; result.height = 0; result.pixels = 0;
  return result;
}

// Composite arched frequency from individual digit bitmaps
void drawArchedFrequency(float frequency) {
  int whole = (int)frequency;
  int decimal = (int)((frequency - whole) * 10 + 0.5);
  
  char freqStr[16];
  sprintf(freqStr, "%d.%d", whole, decimal);
  
  bool is3digit = (whole >= 100);
  
  // Get all digit parts
  DigitInfo parts[6];
  DigitInfo dot, fm;
  int partCount = 0;
  int totalWidth = 0;
  int maxHeight = 0;
  
  if (is3digit) {
    parts[0] = getDigit3d(0, '1');  // Always '1'
    parts[1] = getDigit3d(1, '0');  // Always '0'
    parts[2] = getDigit3d(2, freqStr[2]);  // ones
    dot = {DIGIT_3d_dot_D, DIGIT_3d_dot_W, DIGIT_3d_dot_H, DIGIT_3d_dot_P};
    parts[3] = getDigit3d(4, freqStr[4]);  // decimal
    fm = {DIGIT_3d_FM_D, DIGIT_3d_FM_W, DIGIT_3d_FM_H, DIGIT_3d_FM_P};
    partCount = 4;
  } else {
    parts[0] = getDigit2d(0, freqStr[0]);  // tens
    parts[1] = getDigit2d(1, freqStr[1]);  // ones
    dot = {DIGIT_2d_dot_D, DIGIT_2d_dot_W, DIGIT_2d_dot_H, DIGIT_2d_dot_P};
    parts[2] = getDigit2d(3, freqStr[3]);  // decimal
    fm = {DIGIT_2d_FM_D, DIGIT_2d_FM_W, DIGIT_2d_FM_H, DIGIT_2d_FM_P};
    partCount = 3;
  }
  
  // Calculate total size
  for (int i = 0; i < partCount; i++) {
    totalWidth += parts[i].width;
    if (parts[i].height > maxHeight) maxHeight = parts[i].height;
  }
  totalWidth += dot.width + fm.width;
  if (dot.height > maxHeight) maxHeight = dot.height;
  if (fm.height > maxHeight) maxHeight = fm.height;
  
  // Composite into buffer
  uint8_t* composite = (uint8_t*)malloc(totalWidth * maxHeight);
  if (!composite) {
    Serial.println(F("Composite buffer failed"));
    return;
  }
  memset(composite, 0, totalWidth * maxHeight);
  
  int xOffset = 0;
  
  // Copy digit bitmaps in order
  if (is3digit) {
    // '1' '0' 'X' '.' 'Y' ' FM'
    // Copy: parts[0], parts[1], parts[2], dot, parts[3], fm
    
    // First three digits: 1, 0, X
    for (int partIdx = 0; partIdx < 3; partIdx++) {
      const DigitInfo* part = &parts[partIdx];
      
      if (part->data) {
        for (uint16_t i = 0; i < part->pixels; i++) {
          int packed_idx = i / 4;
          uint8_t packed = pgm_read_byte(&part->data[packed_idx]);
          uint8_t shift = 6 - (i % 4) * 2;
          uint8_t level = (packed >> shift) & 0x03;
          
          int x = i % part->width;
          int y = i / part->width;
          composite[y * totalWidth + xOffset + x] = level;
        }
        xOffset += part->width;
      }
    }
    
    // Dot
    if (dot.data) {
      for (uint16_t i = 0; i < dot.pixels; i++) {
        int packed_idx = i / 4;
        uint8_t packed = pgm_read_byte(&dot.data[packed_idx]);
        uint8_t shift = 6 - (i % 4) * 2;
        uint8_t level = (packed >> shift) & 0x03;
        
        int x = i % dot.width;
        int y = i / dot.width;
        composite[y * totalWidth + xOffset + x] = level;
      }
      xOffset += dot.width;
    }
    
    // Last decimal digit
    if (parts[3].data) {
      for (uint16_t i = 0; i < parts[3].pixels; i++) {
        int packed_idx = i / 4;
        uint8_t packed = pgm_read_byte(&parts[3].data[packed_idx]);
        uint8_t shift = 6 - (i % 4) * 2;
        uint8_t level = (packed >> shift) & 0x03;
        
        int x = i % parts[3].width;
        int y = i / parts[3].width;
        composite[y * totalWidth + xOffset + x] = level;
      }
      xOffset += parts[3].width;
    }
    
    // FM
    if (fm.data) {
      for (uint16_t i = 0; i < fm.pixels; i++) {
        int packed_idx = i / 4;
        uint8_t packed = pgm_read_byte(&fm.data[packed_idx]);
        uint8_t shift = 6 - (i % 4) * 2;
        uint8_t level = (packed >> shift) & 0x03;
        
        int x = i % fm.width;
        int y = i / fm.width;
        composite[y * totalWidth + xOffset + x] = level;
      }
      xOffset += fm.width;
    }
  } else {
    // 2-digit: 'X' 'Y' '.' 'Z' ' FM'
    // Copy: parts[0], parts[1], dot, parts[2], fm
    
    // First two digits
    for (int partIdx = 0; partIdx < 2; partIdx++) {
      const DigitInfo* part = &parts[partIdx];
      
      if (part->data) {
        for (uint16_t i = 0; i < part->pixels; i++) {
          int packed_idx = i / 4;
          uint8_t packed = pgm_read_byte(&part->data[packed_idx]);
          uint8_t shift = 6 - (i % 4) * 2;
          uint8_t level = (packed >> shift) & 0x03;
          
          int x = i % part->width;
          int y = i / part->width;
          composite[y * totalWidth + xOffset + x] = level;
        }
        xOffset += part->width;
      }
    }
    
    // Dot
    if (dot.data) {
      for (uint16_t i = 0; i < dot.pixels; i++) {
        int packed_idx = i / 4;
        uint8_t packed = pgm_read_byte(&dot.data[packed_idx]);
        uint8_t shift = 6 - (i % 4) * 2;
        uint8_t level = (packed >> shift) & 0x03;
        
        int x = i % dot.width;
        int y = i / dot.width;
        composite[y * totalWidth + xOffset + x] = level;
      }
      xOffset += dot.width;
    }
    
    // Last decimal digit
    if (parts[2].data) {
      for (uint16_t i = 0; i < parts[2].pixels; i++) {
        int packed_idx = i / 4;
        uint8_t packed = pgm_read_byte(&parts[2].data[packed_idx]);
        uint8_t shift = 6 - (i % 4) * 2;
        uint8_t level = (packed >> shift) & 0x03;
        
        int x = i % parts[2].width;
        int y = i / parts[2].width;
        composite[y * totalWidth + xOffset + x] = level;
      }
      xOffset += parts[2].width;
    }
    
    // FM
    if (fm.data) {
      for (uint16_t i = 0; i < fm.pixels; i++) {
        int packed_idx = i / 4;
        uint8_t packed = pgm_read_byte(&fm.data[packed_idx]);
        uint8_t shift = 6 - (i % 4) * 2;
        uint8_t level = (packed >> shift) & 0x03;
        
        int x = i % fm.width;
        int y = i / fm.width;
        composite[y * totalWidth + xOffset + x] = level;
      }
      xOffset += fm.width;
    }
  }
  
  // Convert pre-dithered levels directly to RGB565 colors (no dithering needed)
  uint16_t* rgbBuffer = (uint16_t*)calloc(totalWidth * maxHeight, sizeof(uint16_t));
  if (!rgbBuffer) {
    free(composite);
    return;
  }
  
  for (int y = 0; y < maxHeight; y++) {
    for (int x = 0; x < totalWidth; x++) {
      int idx = y * totalWidth + x;
      uint8_t level = composite[idx];
      rgbBuffer[idx] = (level > 0) ? getCyanBrightness(level) : ILI9341_BLACK;
    }
  }
  
  free(composite);
  
  // Center on screen
  long sumX = 0, sumY = 0, count = 0;
  for (int y = 0; y < maxHeight; y++) {
    for (int x = 0; x < totalWidth; x++) {
      int idx = y * totalWidth + x;
      if (rgbBuffer[idx] != ILI9341_BLACK) {
        sumX += x;
        sumY += y;
        count++;
      }
    }
  }
  
  int contentCenterX = count > 0 ? sumX / count : totalWidth / 2;
  int contentCenterY = count > 0 ? sumY / count : maxHeight / 2;
  
  int startX = (display.width() / 2) - contentCenterX;
  int startY = (display.height() / 2) - contentCenterY;
  
  display.drawRGBBitmap(startX, startY, rgbBuffer, totalWidth, maxHeight);
  
  free(rgbBuffer);
}

// Draw simple frequency text (fast, small font)
void drawSimpleFrequencyFast(float frequency) {
  char freqStr[16];
  sprintf(freqStr, "%.1f FM", frequency);
  
  display.setTextColor(ILI9341_CYAN);
  display.setTextSize(6);
  
  // Calculate text bounds
  int16_t x1, y1;
  uint16_t w, h;
  display.getTextBounds(freqStr, 0, 0, &x1, &y1, &w, &h);
  
  // Center on screen
  int x = (display.width() - w) / 2;
  int y = (display.height() - h) / 2;
  
  display.setCursor(x, y);
  display.print(freqStr);
}

// Animate scrolling from current frequency to target frequency
// Steps by 0.1 MHz increments using simple font, then shows final arched display
void scrollToFrequency(float currentFreq, float targetFreq, int stepDelayMs = 10) {
  // Determine direction
  float step = (targetFreq > currentFreq) ? 0.1 : -0.1;
  
  // Calculate max text bounds once
  display.setTextSize(6);
  int16_t x1, y1;
  uint16_t w, h;
  display.getTextBounds("888.8 FM", 0, 0, &x1, &y1, &w, &h);
  
  int x = (display.width() - w) / 2;
  int y = (display.height() - h) / 2;
  
  // Clear screen once at start
  display.fillScreen(ILI9341_BLACK);
  
  // Animate with simple text
  float freq = currentFreq + step;
  while ((step > 0 && freq <= targetFreq) || (step < 0 && freq >= targetFreq)) {
    // Clear only text area
    display.fillRect(x - 2, y - 2, w + 4, h + 4, ILI9341_BLACK);
    drawSimpleFrequencyFast(freq);
    delay(stepDelayMs);
    freq += step;
  }
  
  // Draw final arched frequency
  display.fillScreen(ILI9341_BLACK);
  drawArchedFrequency(targetFreq);
}

void setup() {
  Serial.begin(9600);
  
  display.begin();
  display.setRotation(1);  // Landscape mode
  display.fillScreen(ILI9341_BLACK);
  
  Serial.println(F("ILI9341 initialized"));
  drawArchedFrequency(85.2);
}

void loop() {
  // Demonstrate scrolling animation
  static float frequencies[] = {87.5, 88.0, 92.3, 99.9, 100.0, 105.5, 108.0};
  static int freqIndex = 0;
  static float lastFreq = 87.5;
  
  // Show initial frequency
  if (freqIndex == 0) {
    display.fillScreen(ILI9341_BLACK);
    drawArchedFrequency(lastFreq);
    delay(1500);
  }
  
  // Move to next frequency
  freqIndex = (freqIndex + 1) % 7;
  float targetFreq = frequencies[freqIndex];
  
  unsigned long startTime = millis();
  scrollToFrequency(lastFreq, targetFreq, 10);  // 10ms per step
  unsigned long endTime = millis();
  
  Serial.print(F("Scrolled from "));
  Serial.print(lastFreq, 1);
  Serial.print(F(" to "));
  Serial.print(targetFreq, 1);
  Serial.print(F(" in "));
  Serial.print(endTime - startTime);
  Serial.println(F(" ms"));
  
  lastFreq = targetFreq;
  delay(1500);  // Pause at final frequency
}