#include "Audio_Player.h"
#include <ESP_I2S.h>
#include <algorithm>

I2SClass i2s;
TaskHandle_t audioTaskHandle = NULL;
File audioFile;
volatile bool isAudioPlaying = false;
char currentPlayPath[128];

uint16_t Wav_Count = 0;
char Wav_Names[100][100];

void Search_Wavs(const char* directory) {
  Wav_Count = Folder_retrieval(directory, ".wav", Wav_Names, 100);
}


// Parse WAV header and skip meta chunks (like LIST) to find the "fmt " and "data" segments
bool parseWavHeader(File& file, uint32_t& sampleRate, uint16_t& channels, uint16_t& bitsPerSample, uint32_t& dataStartOffset, uint32_t& dataSize) {
  char id[4];
  
  // Read RIFF header
  file.seek(0);
  if (file.readBytes(id, 4) != 4 || memcmp(id, "RIFF", 4) != 0) {
    printf("WAV Header Error: RIFF signature not found!\r\n");
    return false;
  }
  
  file.seek(8);
  if (file.readBytes(id, 4) != 4 || memcmp(id, "WAVE", 4) != 0) {
    printf("WAV Header Error: WAVE signature not found!\r\n");
    return false;
  }
  
  uint32_t offset = 12;
  bool foundFmt = false;
  bool foundData = false;
  
  while (offset < file.size() && (!foundFmt || !foundData)) {
    file.seek(offset);
    if (file.readBytes(id, 4) != 4) break;
    
    uint32_t chunkSize;
    if (file.read((uint8_t*)&chunkSize, 4) != 4) break;
    
    offset += 8; // move past chunk ID and Size field
    
    if (memcmp(id, "fmt ", 4) == 0) {
      uint16_t audioFormat;
      file.read((uint8_t*)&audioFormat, 2);
      if (audioFormat != 1) { // 1 = PCM (uncompressed)
        printf("WAV Header Error: Only uncompressed PCM WAV is supported!\r\n");
        return false;
      }
      file.read((uint8_t*)&channels, 2);
      file.read((uint8_t*)&sampleRate, 4);
      file.seek(offset + 8); // skip byte rate and block align
      file.read((uint8_t*)&bitsPerSample, 2);
      foundFmt = true;
    } 
    else if (memcmp(id, "data", 4) == 0) {
      dataStartOffset = offset;
      dataSize = chunkSize;
      foundData = true;
    }
    
    offset += chunkSize;
  }
  
  return foundFmt && foundData;
}

void audioPlayTask(void* parameter) {
  uint32_t sampleRate = 44100;
  uint16_t channels = 2;
  uint16_t bitsPerSample = 16;
  uint32_t dataStartOffset = 44;
  uint32_t dataSize = 0;

  audioFile = SD_MMC.open(currentPlayPath, "r");
  if (!audioFile) {
    printf("Failed to open audio file: %s\r\n", currentPlayPath);
    isAudioPlaying = false;
    audioTaskHandle = NULL;
    vTaskDelete(NULL);
    return;
  }

  if (!parseWavHeader(audioFile, sampleRate, channels, bitsPerSample, dataStartOffset, dataSize)) {
    printf("Failed to parse WAV header\r\n");
    audioFile.close();
    isAudioPlaying = false;
    audioTaskHandle = NULL;
    vTaskDelete(NULL);
    return;
  }

  printf("WAV File: Rate=%d Hz, Channels=%d, Bits=%d, Size=%d bytes\r\n", 
         sampleRate, channels, bitsPerSample, dataSize);

  // Configure I2S based on WAV parameters
  i2s_data_bit_width_t bitWidth = I2S_DATA_BIT_WIDTH_16BIT;
  if (bitsPerSample == 8) bitWidth = I2S_DATA_BIT_WIDTH_8BIT;
  else if (bitsPerSample == 24) bitWidth = I2S_DATA_BIT_WIDTH_24BIT;
  else if (bitsPerSample == 32) bitWidth = I2S_DATA_BIT_WIDTH_32BIT;

  i2s_slot_mode_t slotMode = (channels == 1) ? I2S_SLOT_MODE_MONO : I2S_SLOT_MODE_STEREO;

  // Set standard pins
  i2s.setPins(I2S_BCLK, I2S_LRCK, I2S_DIN, -1, -1);

  // Stop any active driver configuration
  i2s.end();

  if (!i2s.begin(I2S_MODE_STD, sampleRate, bitWidth, slotMode)) {
    printf("Failed to initialize I2S!\r\n");
    audioFile.close();
    isAudioPlaying = false;
    audioTaskHandle = NULL;
    vTaskDelete(NULL);
    return;
  }

  // Seek past the parsed headers to raw data
  audioFile.seek(dataStartOffset);

  const size_t bufferSize = 1024;
  uint8_t buffer[bufferSize];
  uint32_t bytesRead = 0;
  uint32_t totalBytesWritten = 0;

  while (isAudioPlaying && totalBytesWritten < dataSize) {
    size_t toRead = std::min(bufferSize, (size_t)(dataSize - totalBytesWritten));
    bytesRead = audioFile.read(buffer, toRead);
    if (bytesRead == 0) {
      break; 
    }

    size_t written = i2s.write(buffer, bytesRead);
    if (written == 0) {
      vTaskDelay(pdMS_TO_TICKS(1));
      continue;
    }
    totalBytesWritten += written;
  }

  printf("Playback done. Written %d bytes.\r\n", totalBytesWritten);
  audioFile.close();
  i2s.end();
  isAudioPlaying = false;
  audioTaskHandle = NULL;
  vTaskDelete(NULL);
}

void playWav(const char* filePath) {
  if (isAudioPlaying) {
    stopWav();
  }
  strncpy(currentPlayPath, filePath, sizeof(currentPlayPath) - 1);
  currentPlayPath[sizeof(currentPlayPath) - 1] = '\0';
  isAudioPlaying = true;
  
  xTaskCreatePinnedToCore(
    audioPlayTask,
    "audioPlayTask",
    4096,
    NULL,
    5,
    &audioTaskHandle,
    0
  );
}

void stopWav() {
  if (isAudioPlaying) {
    isAudioPlaying = false;
    while (audioTaskHandle != NULL) {
      vTaskDelay(pdMS_TO_TICKS(5));
    }
  }
}

bool isWavPlaying() {
  return isAudioPlaying;
}