/*
  ESP32-S2/S3 Gotek Controller — RFID-tagged, Pin-controlled, mTLS backend

  Features:
    - Minimal web UI (/setup): Wi‑Fi SSID/password + backend URL (https)
    - Mutual TLS (client certificate) for backend communication
    - LOAD_PIN high => read RFID tag (64-bit hex), GET /api/volume?id=<id>, mount USB MSC (RW), pulse JA BTN2 (SELECT)
    - LOAD_PIN low  => pulse JA BTN3 (EJECT), unmount, POST /api/volume?id=<id> (save modified volume; server extracts image)
    - /api/remount  => force re-fetch (reads tag if present), mount, and SELECT
    - /api/status   => JSON (load pin, state, usb status/blocks, wifi, backend, lastId, lastNote)
    - TinyUSB MSC, LittleFS for /stick.img and /config.json

  Requirements:
    - Board: ESP32-S2 or ESP32-S3 (native USB)
    - Arduino-ESP32 core >= 3.0.0
    - LittleFS enabled in partition scheme
    - Wire JA BTN2 (SELECT) and BTN3 (EJECT) via NPN/optos (active-low)
    - RFID reader TX -> ESP32 RX pin (RFID_RX_PIN), level shift if reader is 5V

  Security:
    - Replace SERVER_CA_PEM, CLIENT_CRT_PEM, CLIENT_KEY_PEM with your actual CA, device certificate and private key.
*/

#include <WiFi.h>
#include <WebServer.h>
#include <ESPmDNS.h>
#include <HTTPClient.h>
#include <FS.h>
#include <LittleFS.h>
#include <ArduinoJson.h>
#include "tusb.h"

// ---------------- GPIO Configuration ----------------
static const int LOAD_PIN      = 21;  // Input: HIGH = load/mount, LOW = eject/save
static const int PIN_SELECT_JA = 4;   // Output to Gotek JA BTN2 (SELECT)  (active-low via NPN/opto)
static const int PIN_EJECT_JA  = 5;   // Output to Gotek JA BTN3 (EJECT)   (active-low via NPN/opto)
static const int RFID_RX_PIN   = 17;  // RX pin for RFID UART reader (reader TX -> this pin)
static const long RFID_BAUD    = 9600;

// ---------------- Timings ----------------
static const int PRESS_MS = 60;
static const int GAP_MS   = 120;
static const uint16_t BYTES_PER_SECTOR = 512;
static const size_t   MAX_VOLUME_BYTES = 8 * 1024 * 1024; // adjust to your partition size

// ---------------- Files ----------------
static const char* PATH_CFG = "/config.json";
static const char* PATH_VOL = "/stick.img";

// ---------------- mTLS Material (PLACEHOLDERS — replace with your real PEMs) ----------------
static const char SERVER_CA_PEM[] PROGMEM = R"(-----BEGIN CERTIFICATE-----
...your CA that signed the BACKEND SERVER certificate...
-----END CERTIFICATE-----)";
static const char CLIENT_CRT_PEM[] PROGMEM = R"(-----BEGIN CERTIFICATE-----
...device client certificate...
-----END CERTIFICATE-----)";
static const char CLIENT_KEY_PEM[] PROGMEM = R"(-----BEGIN PRIVATE KEY-----
...device private key...
-----END PRIVATE KEY-----)";

// ---------------- State / Config ----------------
struct AppCfg { String ssid, pass, backendUrl; } CFG;
WebServer server(80);
File mscFile;
volatile bool mscReady = false;
uint32_t mscBlocks = 0;

enum class LoadState { UNLOADED, LOADED };
LoadState state = LoadState::UNLOADED;

static String lastNote;
static String lastId;     // cache last RFID ID used (hex64)

HardwareSerial RFID(1);   // UART1 for the RFID reader

// ---------------- Helpers ----------------
static inline void line_set(int pin, bool high){ digitalWrite(pin, high?HIGH:LOW); }
static void pulse(int pin){ line_set(pin,false); delay(PRESS_MS); line_set(pin,true); delay(GAP_MS); }

static inline bool load_pin(){ return digitalRead(LOAD_PIN)==HIGH; }

static const char* wifi_state() {
  wl_status_t s = WiFi.status();
  switch (s) {
    case WL_CONNECTED: return "connected";
    case WL_IDLE_STATUS: return "idle";
    case WL_NO_SSID_AVAIL: return "no_ssid";
    case WL_CONNECT_FAILED: return "connect_failed";
    case WL_CONNECTION_LOST: return "lost";
    case WL_DISCONNECTED: return "disconnected";
    default: return "unknown";
  }
}

// ---------------- USB MSC (TinyUSB) ----------------
//extern "C" {
  //int32_t tud_msc_read10_cb (uint8_t, uint32_t lba, uint32_t off, void* buf, uint32_t len){
  //  if(!mscFile) return -1; uint64_t a=(uint64_t)lba*BYTES_PER_SECTOR + off;
  //  if(!mscFile.seek(a)) return -1; return (int32_t)mscFile.read((uint8_t*)buf, len);
 // }
  //int32_t tud_msc_write10_cb(uint8_t, uint32_t lba, uint32_t off, uint8_t const* buf, uint32_t len){
  //  if(!mscFile) return -1; uint64_t a=(uint64_t)lba*BYTES_PER_SECTOR + off;
  //  if(!mscFile.seek(a)) return -1; return (int32_t)mscFile.write((const uint8_t*)buf, len);
 // }
  //bool    tud_msc_is_writable_cb(uint8_t){ return true; }
  //bool    tud_msc_test_unit_ready_cb(uint8_t){ return mscReady && mscFile; }
  //void    tud_msc_capacity_cb(uint8_t, uint32_t* bc, uint16_t* bs){ *bc=mscBlocks; *bs=BYTES_PER_SECTOR; }
  //bool    tud_msc_start_stop_cb(uint8_t, uint8_t, bool start, bool load_eject)
  //{
  //  mscReady = (start && !load_eject && mscFile); return true;
  //}
//}
static void usb_init()
{ 
  tinyusb_config_t cfg={}; 
  tusb_init(&cfg); 
  mscReady=false; 
}

static bool vol_open(){
  if(mscFile) mscFile.close();
  if(!LittleFS.exists(PATH_VOL)){ mscBlocks=0; return false; }
  // open RW; if fails, RO
  mscFile = LittleFS.open(PATH_VOL, "r+");
  if(!mscFile){
    mscFile = LittleFS.open(PATH_VOL, "r");
    if(!mscFile) return false;
  }
  uint64_t sz=mscFile.size(); if(sz<BYTES_PER_SECTOR){ mscFile.close(); return false; }
  mscBlocks = (uint32_t)(sz / BYTES_PER_SECTOR);
  return true;
}
static void set_ready(bool r){ mscReady = r && mscFile; }

// ---------------- Config I/O ----------------
static bool cfg_load(){
  if(!LittleFS.exists(PATH_CFG)) return false;
  File f=LittleFS.open(PATH_CFG,"r"); if(!f) return false;
  DynamicJsonDocument d(1024); if(deserializeJson(d,f)){ f.close(); return false; } f.close();
  CFG.ssid=(const char*)(d["ssid"]|""); CFG.pass=(const char*)(d["pass"]|""); CFG.backendUrl=(const char*)(d["backendUrl"]|"");
  return true;
}
static bool cfg_save(){
  DynamicJsonDocument d(1024); d["ssid"]=CFG.ssid; d["pass"]=CFG.pass; d["backendUrl"]=CFG.backendUrl;
  File f=LittleFS.open(PATH_CFG,"w"); if(!f) return false; bool ok=serializeJson(d,f)>0; f.close(); return ok;
}

// ---------------- Wi‑Fi ----------------
static void wifi_connect_or_ap(){
  if(CFG.ssid.length()){
    WiFi.mode(WIFI_STA); WiFi.begin(CFG.ssid.c_str(), CFG.pass.c_str());
    unsigned long t0=millis(); while(WiFi.status()!=WL_CONNECTED && millis()-t0<12000) delay(200);
  }
  if(WiFi.status()==WL_CONNECTED){ if(MDNS.begin("gotek")) MDNS.addService("http","tcp",80); }
  else { WiFi.mode(WIFI_AP); WiFi.softAP("GotekSetup","flashfloppy"); }
}

// ---------------- RFID (read tag -> uint64) ----------------
HardwareSerial RFID(1);
static bool rfid_read_tag(uint64_t &outId, uint32_t timeout_ms=3000){
  uint32_t t0 = millis();
  String line; line.reserve(24);
  while(millis()-t0 < timeout_ms){
    while(RFID.available()){
      int c = RFID.read();
      if (c == '\r' || c == '\n') {
        // parse hex out of line
        String hex; hex.reserve(16);
        for (size_t i=0;i<line.length();++i){
          char ch = line[i];
          if (isxdigit((unsigned char)ch)) hex += (char)toupper(ch);
        }
        if (hex.length() == 0 || hex.length() > 16){ line = ""; break; } // discard invalid
        uint64_t v = 0;
        for (size_t i=0;i<hex.length();++i){
          char h = hex[i]; uint8_t nib = (h <= '9') ? (h - '0') : (10 + (h - 'A'));
          v = (v << 4) | nib;
        }
        outId = v; return true;
      } else {
        line += (char)c;
        if (line.length() > 32) line.remove(0, line.length()-32);
      }
    }
    delay(2);
  }
  return false;
}
static String hex64(uint64_t v){
  char buf[17]; char* p = &buf[16]; *p = '\0';
  if (v == 0){ buf[0]='0'; buf[1]='\0'; return String(buf); }
  while (v && p > buf){
    uint8_t nib = v & 0xF; *--p = (nib < 10) ? ('0'+nib) : ('a'+(nib-10)); v >>= 4;
  }
  return String(p);
}

// ---------------- Backend I/O (HTTPS mTLS) ----------------
static const char SERVER_CA_PEM[] PROGMEM = R"(-----BEGIN CERTIFICATE-----
...your CA that signed the BACKEND SERVER certificate...
-----END CERTIFICATE-----)";
static const char CLIENT_CRT_PEM[] PROGMEM = R"(-----BEGIN CERTIFICATE-----
...device client certificate...
-----END CERTIFICATE-----)";
static const char CLIENT_KEY_PEM[] PROGMEM = R"(-----BEGIN PRIVATE KEY-----
...device private key...
-----END PRIVATE KEY-----)";

static bool https_fetch_volume(const String& id){
  if(CFG.backendUrl.isEmpty() || id.isEmpty()){ lastNote="no backend/id"; return false; }
  String url = CFG.backendUrl; if(url.endsWith("/")) url.remove(url.length()-1);
  url += "/api/volume?id="; url += id;

  WiFiClientSecure client;
  client.setCACert(SERVER_CA_PEM);
  client.setCertificate(CLIENT_CRT_PEM);
  client.setPrivateKey(CLIENT_KEY_PEM);

  HTTPClient https; if(!https.begin(client, url)){ lastNote="begin fail"; return false; }
  int code=https.GET(); if(code!=HTTP_CODE_OK){ https.end(); lastNote="HTTP "+String(code); return false; }

  int len=https.getSize(); if(len>0 && (size_t)len>MAX_VOLUME_BYTES){ https.end(); lastNote="too big"; return false; }
  File out=LittleFS.open(PATH_VOL,"w"); if(!out){ https.end(); lastNote="open vol fail"; return false; }

  WiFiClient* s=https.getStreamPtr(); size_t total=0; uint8_t buf[4096];
  while(https.connected()){
    size_t avail=s->available(); if(!avail){ if(!https.connected()) break; delay(5); continue; }
    size_t n = avail>sizeof(buf)?sizeof(buf):avail; int got=s->readBytes((char*)buf,n);
    if(got<=0) break; if(total+got>MAX_VOLUME_BYTES){ out.close(); LittleFS.remove(PATH_VOL); https.end(); lastNote="overflow"; return false; }
    out.write(buf,got); total+=got;
  }
  out.close(); https.end(); lastNote="downloaded";
  return true;
}

static bool https_post_volume(const String& id){
  if(!LittleFS.exists(PATH_VOL) || id.isEmpty()){ lastNote="no vol/id"; return false; }
  String url = CFG.backendUrl; if(url.endsWith("/")) url.remove(url.length()-1);
  url += "/api/volume?id="; url += id;

  WiFiClientSecure client;
  //client.setCACert(SERVER_CA_PEM);
  //client.setCertificate(CLIENT_CRT_PEM);
  //client.setPrivateKey(CLIENT_KEY_PEM);

  HTTPClient https; if(!https.begin(client, url)){ lastNote="begin fail"; return false; }
  https.addHeader("Content-Type","application/octet-stream");
  File in = LittleFS.open(PATH_VOL,"r"); if(!in){ https.end(); lastNote="open vol fail"; return false; }
  int code = https.sendRequest("POST",&in,in.size());
  in.close(); https.end();
  lastNote = (code==HTTP_CODE_OK || code==HTTP_CODE_ACCEPTED) ? "uploaded OK" : ("upload fail "+String(code));
  return (code==HTTP_CODE_OK || code==HTTP_CODE_ACCEPTED);
}

// ---------------- Web UI (/setup, /api/config, /api/reboot, /api/status, /api/remount) ----------------
static String html_setup(){
  String h; h.reserve(2600);
  h += F("<!doctype html><meta charset=utf-8><meta name=viewport content='width=device-width,initial-scale=1'>"
         "<title>Gotek Wi‑Fi Setup</title><style>body{font-family:system-ui;margin:2rem;max-width:640px}"
         "label{display:block;margin:.75rem 0 .25rem}input,button{width:100%;padding:.6rem;border:1px solid #ccc;border-radius:10px}"
         ".row{display:flex;gap:.75rem;margin-top:.75rem}.row>*{flex:1}</style>"
         "<h1>Wi‑Fi & Backend</h1>"
         "<label>SSID<input id=s value='"); h+=CFG.ssid; h+=F("'></label>"
         "<label>Password<input id=p type=password value='"); h+=CFG.pass; h+=F("'></label>"
         "<label>Backend URL (https://host:port)<input id=u value='"); h+=CFG.backendUrl; h+=F("'></label>"
         "<div class=row><button id=save>Save</button><button id=reboot>Reboot</button></div>"
         "<script>save.onclick=async()=>{const b=JSON.stringify({ssid:s.value,pass:p.value,backendUrl:u.value});"
         "const r=await fetch('/api/config',{method:'POST',headers:{'Content-Type':'application/json'},body:b});"
         "alert((await r.json()).ok?'Saved':'Save failed');};"
         "reboot.onclick=async()=>{await fetch('/api/reboot',{method:'POST'});};</script>");
  return h;
}
static void h_setup(){ server.send(200,"text/html; charset=utf-8", html_setup()); }

static void h_cfg(){
  if(!server.hasArg("plain")){ server.send(400,"application/json","{\"error\":\"no body\"}"); return; }
  DynamicJsonDocument d(1024); if(deserializeJson(d,server.arg("plain"))){ server.send(400,"application/json","{\"error\":\"bad json\"}"); return; }
  CFG.ssid=(const char*)(d["ssid"]|""); CFG.pass=(const char*)(d["pass"]|""); CFG.backendUrl=(const char*)(d["backendUrl"]|"");
  bool ok=cfg_save(); server.send(ok?200:500,"application/json", ok?"{\"ok\":true}":"{\"error\":\"save failed\"}");
}
static void h_reboot(){ server.send(200,"application/json","{\"ok\":true}"); delay(200); ESP.restart(); }

static void h_status(){
  char ipbuf[32] = {0};
  if (WiFi.status()==WL_CONNECTED) {
    IPAddress ip = WiFi.localIP();
    snprintf(ipbuf, sizeof(ipbuf), "%u.%u.%u.%u", ip[0], ip[1], ip[2], ip[3]);
  }
  String s = "{";
  s += "\"load_pin\":";  s += (load_pin() ? "true":"false"); s += ",";
  s += "\"state\":\"";   s += (state==LoadState::LOADED? "loaded":"unloaded"); s += "\",";
  s += "\"usb_ready\":"; s += (mscReady? "true":"false"); s += ",";
  s += "\"blocks\":";    s += String(mscBlocks); s += ",";
  s += "\"backend\":\""; s += CFG.backendUrl; s += "\",";
  s += "\"wifi\":{\"status\":\""; s += wifi_state(); s += "\",\"rssi\":";
  s += (WiFi.status()==WL_CONNECTED ? String(WiFi.RSSI()) : "null");
  s += ",\"ip\":\""; s += ipbuf; s += "\"},";
  s += "\"lastId\":\""; s += lastId; s += "\",";
  s += "\"note\":\""; s += lastNote; s += "\"}";
  server.send(200,"application/json",s);
}

// Force fetch + mount + SELECT (reads tag if available)
static void h_remount(){
  set_ready(false);
  if(mscFile) mscFile.flush();
  if(mscFile) mscFile.close();

  String idParam = "";
  uint64_t tag=0;
  if (rfid_read_tag(tag, 2000)) idParam = hex64(tag);
  else if (lastId.length()) idParam = lastId;

  bool ok = (idParam.length() && https_fetch_volume(idParam) && vol_open());
  if(ok){
    lastId = idParam;
    set_ready(true);
    pulse(PIN_SELECT_JA);
    state = LoadState::LOADED;
    server.send(200,"application/json","{\"ok\":true}");
  } else {
    server.send(500,"application/json","{\"error\":\"fetch/mount failed\"}");
  }
}

// ---------------- State transitions ----------------
static void enter_loaded(){
  // Read tag
  uint64_t tag=0;
  //if(!rfid_read_tag(tag, 3000)){ lastNote="no tag"; return; }
  String id = hex64(tag);
  if(!https_fetch_volume(id)) return;
  if(!vol_open()) return;
  lastId = id;
  set_ready(true);
  pulse(PIN_SELECT_JA); // SELECT
  state = LoadState::LOADED;
}

static void enter_unloaded(){
  pulse(PIN_EJECT_JA);  // EJECT
  set_ready(false);
  if(mscFile) mscFile.flush();
  if(mscFile) mscFile.close();
  if(lastId.length()) https_post_volume(lastId);
  state = LoadState::UNLOADED;
}

// ---------------- Setup / Loop ----------------
void setup(){
  Serial.begin(115200);
  LittleFS.begin(true);

  pinMode(LOAD_PIN,      INPUT); // use external pulldown if needed
  pinMode(PIN_SELECT_JA, OUTPUT); digitalWrite(PIN_SELECT_JA, HIGH); // idle-high
  pinMode(PIN_EJECT_JA,  OUTPUT); digitalWrite(PIN_EJECT_JA,  HIGH);

  // RFID UART (RX only)
  RFID.begin(RFID_BAUD, SERIAL_8N1, RFID_RX_PIN, -1);

  cfg_load();
  wifi_connect_or_ap();
  usb_init();

  // initial state by pin
  if(load_pin()) enter_loaded(); else enter_unloaded();

  // HTTP routes
  server.on("/",            HTTP_GET,  h_setup);
  server.on("/setup",       HTTP_GET,  h_setup);
  server.on("/api/config",  HTTP_POST, h_cfg);
  server.on("/api/reboot",  HTTP_POST, h_reboot);
  server.on("/api/status",  HTTP_GET,  h_status);
  server.on("/api/remount", HTTP_POST, h_remount);
  server.begin();

  if(MDNS.begin("gotek")) MDNS.addService("http","tcp",80);
}

void loop(){
  server.handleClient();
  tud_task();

  // Debounced edge detection on LOAD pin
  static bool last = load_pin();
  bool now = load_pin();
  if(now != last){
    delay(30);
    now = load_pin();
    if(now != last){
      last = now;
      if(now && state != LoadState::LOADED)      enter_loaded();
      else if(!now && state != LoadState::UNLOADED) enter_unloaded();
    }
  }
}