// 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>
#include <AuthApiClient.h>
#include "esp32-hal-psram.h"

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

const char* WIFI_SSID = "Aqua Cube IT";
const char* WIFI_PASS = "Iamazombie123";
const char* API_BASE = "https://netfloppy.com.puter.club";
const char* USERNAME = "test";
const char* PASSWORD = "password";
uint32_t CHIP_ID = 0;

// ---- Disk parameters ----
uint32_t DISK_SIZE_MB = 5;
uint32_t DISK_SECTOR_SIZE  = 512;          // 512 bytes/sector
uint32_t DISK_SECTOR_COUNT = DISK_SIZE_MB * 1024 * 1024 / DISK_SECTOR_SIZE;       
uint32_t DISK_BYTE_SIZE    = DISK_SECTOR_COUNT * DISK_SECTOR_SIZE;

//static uint8_t* msc_disk = nullptr;
static volatile bool media_ready = false;
                               
AuthApiClient _api(API_BASE);

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", WIFI_SSID);
  pass = prefs.getString("pass", WIFI_PASS);
  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 setupWifi() 
{

  forgetPrefs();
  loadPrefs();

  if (ssid == "")
  {
    ssid = WIFI_SSID;
    pass = WIFI_PASS;
    autoConnect = true;

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

    savePrefs();
  }

  Serial.printf("  Connecting to WIFI '%s'...\n", ssid.c_str());
  WiFi.begin(ssid.c_str(), pass.c_str());

  while (WiFi.status() != WL_CONNECTED) 
  { 
    delay(500);
    Serial.print("."); 
  }

  Serial.println("\n  WiFi connected.");

  showStatus();
}

void setupTime() 
{
  configTime(0, 0, "pool.ntp.org", "time.nist.gov");
  Serial.print("Syncing time");
  
  for (int i = 0; i < 20; i++) 
  { 
    if (time(nullptr) > 1700000000) break; 

    delay(500); 
    Serial.print("."); 
  }

  Serial.println();
}

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);

    String unitId = "0";
    String driveId = "0";

    _api.loadDiskSector(unitId, driveId, lba, offset, buffer, 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);
}

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

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

  delay(200);

  Serial.flush();

  if (psramInit() && psramFound()) 
  {
    if (psramAddToHeap()) 
    {
      Serial.println("PSRAM added to heap");
    }
  }

  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");
  

  for (int i = 0; i < 17; i = i + 8) 
  {
    CHIP_ID |= ((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(CHIP_ID);

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

  setupWifi();
  setupTime();  

  //_api.useCertBundle(false);
  //_api.setInsecure(true);

  if (_api.login(USERNAME, PASSWORD))
  {
      Serial.printf("  API Login Successfull\n");
  }
  else
  {
      Serial.printf("  API Login Failed\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());

  initialiseDisk();
}

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