// ESP32-S2/S3 native USB MSC RAM disk for Arduino-ESP32 3.3.0
// Uses USB.h / USBMSC.h (no Adafruit_TinyUSB)

#include <Arduino.h>
#include <Adafruit_NeoPixel.h>
#include "USB.h"
#include "USBMSC.h"
#include <WiFi.h>
#include <Preferences.h>

#define LED_ON  2
#define LED_TX  43
#define LED_RX  44
#define BUTTON  0
#define WS2812  48

// ---- Disk parameters ----
static const uint32_t DISK_SECTOR_SIZE  = 512;          // 512 bytes/sector
static const uint32_t DISK_SECTOR_COUNT = 1000 * 8;      // ~2 MB (4000 * 512)
static const uint32_t DISK_BYTE_SIZE    = DISK_SECTOR_COUNT * DISK_SECTOR_SIZE;

static uint8_t* msc_disk = nullptr;
static volatile bool media_ready = false;

USBMSC msc;
Adafruit_NeoPixel pixels(1, WS2812, NEO_GRB + NEO_KHZ800);

// Utility: bounds check for (lba, offset, len)
static inline bool in_range(uint32_t lba, uint32_t offset, uint32_t len) 
{
  uint64_t start = (uint64_t)lba * DISK_SECTOR_SIZE + offset;
  uint64_t end   = start + len;
  return end <= (uint64_t)DISK_BYTE_SIZE;
}

// ---------- Settings ----------
static const uint32_t SERIAL_BAUD = 115200;
static const uint32_t CONNECT_TIMEOUT_MS = 20000; // 20s


Preferences prefs;
String ssid = "";
String pass = "";
bool autoConnect = false;

void loadPrefs() 
{
  prefs.begin("wifi", true); // read-only
  ssid = prefs.getString("ssid", "Barriball - Automation");
  pass = prefs.getString("pass", "password123abc");
  autoConnect = prefs.getBool("auto", true);
  prefs.end();
}


void savePrefs() 
{
  prefs.begin("wifi", false);
  prefs.putString("ssid", ssid);
  prefs.putString("pass", pass);
  prefs.putBool("auto", autoConnect);
  prefs.end();
}


void forgetPrefs() 
{
  prefs.begin("wifi", false);
  prefs.remove("ssid");
  prefs.remove("pass");
  prefs.remove("auto");
  prefs.end();
  ssid = "";
  pass = "";
  autoConnect = false;
}


void showStatus() 
{
  wl_status_t st = WiFi.status();

  Serial.println("\n=== WiFi Status ===");
  Serial.printf("Mode: %s\n", WiFi.getMode() == WIFI_MODE_STA ? "STA" : "Other");
  Serial.printf("Saved SSID: %s\n", ssid.c_str());
  Serial.printf("Auto-connect: %s\n", autoConnect ? "ON" : "OFF");
  Serial.printf("Runtime status: %d\n", (int)st);

  if (st == WL_CONNECTED) 
  {
    Serial.printf("Connected SSID: %s\n", WiFi.SSID().c_str());
    Serial.printf("IP: %s\n", WiFi.localIP().toString().c_str());
    Serial.printf("RSSI: %d dBm\n", WiFi.RSSI());
  }

  Serial.println("===================\n");
}





void changeMediaStatus(bool mediaPresent)
{

  if (mediaPresent != media_ready)
  {
    media_ready = mediaPresent;

    msc.mediaPresent(mediaPresent);

    if (mediaPresent)
    {
       Serial.printf("Mount disk\n");
      // Set the size of the disk
    //  initialiseDisk();

      showGreenLed();
    }
    else
    {
       Serial.printf("Unmount Disk\n");       
    //  msc.end();
    
      showOrangeLed();
    }
  }
}

void showGreenLed()
{
  pixels.setPixelColor(0, pixels.Color(0, 150, 0));
  pixels.show();  
}

void showOrangeLed()
{
  pixels.setPixelColor(0, pixels.Color(150, 150, 0));
  pixels.show();  
}

void showRedLed()
{
  pixels.setPixelColor(0, pixels.Color(150, 0, 0));
  pixels.show();  
}

void showBlueLed()
{
  pixels.setPixelColor(0, pixels.Color(0, 0, 150));
  pixels.show();  
}

void usbEvent(void* event_handler_arg, esp_event_base_t event_base, int32_t event_id, void* event_data)
{

}

bool onStartStop(uint8_t power_condition, bool start, bool load_eject) 
{
  (void)power_condition;

  // start: true = mount/open; false = stop
  // load_eject: true when host is (eject) requesting unload

  Serial.printf("StartStop: pwr=%u start=%u load_eject=%u\n", power_condition, start, load_eject);


  if (load_eject && !start) 
  {
    Serial.printf("Eject");

    // Zero out storage area
    //memset(msc_disk, 0x00, DISK_BYTE_SIZE);

    // Host requested eject
    changeMediaStatus(false);

    return true;
  } 
  else if (load_eject && start) 
  {
    Serial.printf("Insert");

    // Re-insert
    changeMediaStatus(true);

    return true;
  } 
  else if (!load_eject && !start) 
  {
    Serial.printf("Stop");

    //showRedLed();
    // Stop (could flush if backing store were non-volatile)
    return true;
  } 
  else 
  {
    Serial.printf("Start");

    // Start
    changeMediaStatus(true);
    return true;
  }
}


void initialiseDisk()
{
  // Zero out storage area
  memset(msc_disk, 0x00, DISK_BYTE_SIZE);
  
  // Bring up native USB (CDC + MSC share the device)
  //USB.onEvent(usbEvent)
  USB.begin();

  // int32_t (*msc_read_cb)(uint32_t lba, uint32_t offset, void* buffer, uint32_t bufsize)
  msc.onRead([](uint32_t lba, uint32_t offset, void* buffer, uint32_t bufsize) -> int32_t 
  {
    if (!media_ready) 
    {
      //Serial.println("READ: media not ready");

      return 0;
    }

    if (!in_range(lba, offset, bufsize)) 
    {
      Serial.printf("READ OOB: lba=%lu off=%lu len=%lu\n", lba, offset, bufsize);
      return -1;
    }

    memcpy(buffer, msc_disk + (lba * DISK_SECTOR_SIZE) + offset, bufsize);
    //Serial.printf("READ lba=%lu off=%lu len=%lu\n", lba, offset, bufsize);
    return (int32_t)bufsize;
  });

  // int32_t (*msc_write_cb)(uint32_t lba, uint32_t offset, uint8_t* buffer, uint32_t bufsize)
  msc.onWrite([](uint32_t lba, uint32_t offset, uint8_t* buffer, uint32_t bufsize) -> int32_t 
  {
    if (!media_ready) 
    {
      //Serial.println("WRITE: media not ready");
      return -1;
    }

    if (!in_range(lba, offset, bufsize)) 
    {
      Serial.printf("WRITE OOB: lba=%lu off=%lu len=%lu\n", lba, offset, bufsize);
      return -1;
    }

    memcpy(msc_disk + (lba * DISK_SECTOR_SIZE) + offset, buffer, bufsize);
    //Serial.printf("WRITE lba=%lu off=%lu len=%lu\n", lba, offset, bufsize);
    return (int32_t)bufsize;
  });

  msc.onStartStop(onStartStop);

  // Set disks properties
  msc.vendorID("ACIT");
  msc.productID("NetFloppy");
  msc.productRevision("1.0");
  msc.isWritable(true);
  msc.mediaPresent(true);

  // Set the size of the disk
  if (!msc.begin(DISK_SECTOR_COUNT, DISK_SECTOR_SIZE)) 
  {
    Serial.println("Error: Could not initialise USB disk");
  }   

  changeMediaStatus(true);
  // No setUnitReady() / onReady() in this API; readiness handled inside callbacks.
}

void setup() 
{
  
  Serial.begin(SERIAL_BAUD);

  while (!Serial) 
  { 
    delay(10); 
  }

  delay(200);

  Serial.flush();

  pinMode(LED_ON, OUTPUT);  

  digitalWrite(LED_ON, LOW);

  pixels.begin(); 
  showBlueLed();

  Serial.printf("\n\n+=====================================+\n");
  Serial.printf("| NetFloppy                           |\n");
  Serial.printf("| Version 1.0                         |\n");
  Serial.printf("| Copyright Aqua Cube IT Limited 2025 |\n");
  Serial.printf("+=====================================+\n\n");

  uint32_t chipId = 0;

  for (int i = 0; i < 17; i = i + 8) 
  {
    chipId |= ((ESP.getEfuseMac() >> (40 - i)) & 0xff) << i;
  }

  Serial.printf("ESP32 Chip model = %s Rev %d\n", ESP.getChipModel(), ESP.getChipRevision());
  Serial.printf("This chip has %d cores\n", ESP.getChipCores());
  Serial.print("Chip ID: ");
  Serial.println(chipId);


  Serial.printf("Initialising...\n");
  Serial.printf("Allocating memory...\n");

  // Allocate ramdisk in PSRAM
  msc_disk = (uint8_t*) ps_malloc(DISK_BYTE_SIZE);

  if (!msc_disk) 
  {
    Serial.println("PSRAM alloc FAILED\n");
  }

  Serial.printf("Allocated %u bytes of %u - %u Free\n", (unsigned)DISK_BYTE_SIZE, ESP.getPsramSize(), ESP.getFreePsram());
  

  loadPrefs();

  if (ssid == "")
  {
    ssid = "Barriball - Automation";
    pass = "password123abc";
    autoConnect = true;

Serial.printf("SSID: '%s', Pass: '%s', Auto: %u\n", ssid.c_str(), pass.c_str(), autoConnect);

    savePrefs();
  }

  if (autoConnect && ssid.length()) 
  {
    Serial.printf("Connecting to WIFI '%s'...\n", ssid.c_str());
    //connectWiFi(ssid, pass);
    WiFi.begin(ssid.c_str(), pass.c_str());

    showStatus();
  }

  initialiseDisk();
}

void loop() 
{
  // Simple heartbeat
  //digitalWrite(LED_ON, !digitalRead(LED_ON));
  delay(10);  
}