// KARL'S ULTIMATE DIY DYNO — V3.69 (06.07.26) - Streamlined

#include <WiFi.h>
#include <esp_wifi.h>
#include <ESPAsyncWebServer.h>
#include <AsyncTCP.h>
#include <DNSServer.h>
#include <LittleFS.h>
#include <deque>
#include <vector>
#include <Preferences.h>
#include <ESPmDNS.h>
#include <freertos/FreeRTOS.h>
#include <freertos/task.h>
#include <freertos/queue.h>
#include <sstream>
#include <cmath>
#include <Wire.h>
#include <Adafruit_MCP9600.h>
#include <Adafruit_NAU7802.h>
#include <Update.h>
#include <algorithm>
#include <utility>
#include <Adafruit_BMP280.h>
#include <cstring>

// ─── Version Info ────────────────────────────────────────────────
#define FIRMWARE_VERSION "3.69"
#define BUILD_DATE "06.07.26"

// ─── ESP32-S3 Pin Configuration ──────────────────────────────────
const int HALL_PIN = 4;
const int DRUM_PIN = 5;
const int I2C_SDA = 1;
const int I2C_SCL = 2;
const int RESERVED_PINS[] = {26, 27, 28, 29, 30, 31, 32, 33};
const int RESERVED_COUNT = 8;

// ─── AP Credentials ──────────────────────────────────────────────
const char* ap_ssid = "karlsdyno";
const char* ap_password = "karlsdyno";

// ─── Hall debounce (µs) ──────────────────────────────────────────
int ENGINE_DEBOUNCE_US = 1400;
int DRUM_DEBOUNCE_US = 220;
int MEDIAN_SIZE = 9;
const int MAX_RUN_POINTS = 1500;
const int MAX_SAVED_RUNS = 10;
const int EGT_FILTER_SIZE = 5;

// ─── Update Rates (Hz) ───────────────────────────────────────────
const int SENSOR_UPDATE_MS = 50;
const int WS_UPDATE_MS = 50;
const int EGT_UPDATE_MS = 66;
const int GEAR_UPDATE_MS = 100;
const int BMP_UPDATE_MS = 500;

// ─── Safety Limits ──────────────────────────────────────────────
float MAX_SAFE_DRUM_RPM = 3000.0f;
const float MIN_RPM_THRESHOLD_BRAKE = 35.0f;

// ─── MCP9600 ────────────────────────────────────────────────────
const uint8_t MCP_EGT_ADDR = 0x67;
Adafruit_MCP9600 mcp_egt;
bool mcp_egt_ok = false;

bool apModeOnly = false;
int GPIO_TRIGGER_PIN = 17;

volatile bool gpioTriggerPressed = false;
unsigned long lastGpioDebounceTime = 0;
const unsigned long GPIO_DEBOUNCE_MS = 500;

// ─── EGT Tracker ────────────────────────────────────────────────
struct EGTTracker {
  float buffer[EGT_FILTER_SIZE];
  int index = 0;
  int count = 0;
  float filtered = 0.0f;
  float peak = 0.0f;
  float final = 0.0f;
  float ambient = 0.0f;
  
  void update(float raw, float amb) {
    buffer[index] = raw;
    index = (index + 1) % EGT_FILTER_SIZE;
    if (count < EGT_FILTER_SIZE) count++;
    float sum = 0;
    for (int i = 0; i < count; i++) sum += buffer[i];
    filtered = sum / count;
    if (filtered > peak) peak = filtered;
    final = filtered;
    if (amb > -50 && amb < 200) ambient = amb;
  }
  
  void reset() {
    peak = 0.0f;
    final = 0.0f;
    count = 0;
    index = 0;
    memset(buffer, 0, sizeof(buffer));
  }
  
  float getFiltered() { return filtered; }
  float getAmbient() { return ambient; }
};

EGTTracker egtTracker;

// ─── Smoothing Buffer ───────────────────────────────────────────
struct SmoothingBuffer {
  float buffer[9];
  int index = 0;
  int count = 0;
  
  SmoothingBuffer() { memset(buffer, 0, sizeof(buffer)); }
  
  float add(float value) {
    buffer[index] = value;
    index = (index + 1) % 9;
    if (count < 9) count++;
    float sum = 0;
    for (int i = 0; i < count; i++) sum += buffer[i];
    return sum / count;
  }
  
  void reset() {
    index = 0;
    count = 0;
    memset(buffer, 0, sizeof(buffer));
  }
};

// ─── OPTIMIZED: Move sensor task static variables to GLOBAL SCOPE ───
static float engine_rpm_filtered_g = 0.0f;
static float drum_rpm_filtered_g = 0.0f;
static float ultra_smooth_rpm_g = 0.0f;
static float ultra_smooth_drum_g = 0.0f;
static float last_drum_rpm_g = 0.0f;
static float accel_filtered_g = 0.0f;

// Inertia mode static variables
static float rpm_history[5] = {0};
static int rpm_hist_idx = 0;
static unsigned long last_smooth_time = 0;
static float smooth_accel = 0.0f;
static float drag_torque_g = 0.0f;
static float last_known_torque = 0.0f;
static unsigned long coast_start_time = 0;
static float coast_start_rpm = 0;
static float display_drag_g = 0.0f;

// ─── Global variables ───────────────────────────────────────────
AsyncWebServer server(80);
AsyncWebSocket ws("/ws");
DNSServer dnsServer;

Adafruit_NAU7802 nau;
bool nau_initialized = false;

Adafruit_BMP280 bmp;
bool bmp_initialized = false;

int rpmMedIdx = 0;
int drumMedIdx = 0;
std::vector<float> rpmMedian;
std::vector<float> drumMedian;

float RPM_FILTER_ALPHA = 0.18f;
float CALIBRATION_FACTOR = 0.00678;
float LEVER_ARM_METERS = 0.25;
float MAX_RPM = 16000;
float MAX_HP = 25;
float MAX_TORQUE = 25;
float PULSES_PER_REVOLUTION = 2.0f;
float PULSES_PER_REV_DRUM = 10.0f;
float BRAKE_TO_ENGINE_RATIO = 10.0f;   
float DRUM_MASS = 10.0f;
float DRUM_DIAM_METERS = 0.3f;
float AUTO_HYSTERESIS = 180.0f;
float GEAR_RATIO_MANUAL = 10.0f;
float current_gear_ratio = 10.0f; 
bool GEAR_RATIO_AUTO = true;
uint32_t runTimestamp = 0;
bool useImperial = false; 

String DYNO_MODE = "brake";
String DRUM_INERTIA_TYPE = "solid";
float DRUM_WALL_THICKNESS_CM = 2.0f;
float DRUM_INERTIA_CUSTOM = 0.5f;
String RPM_SOURCE = "spark";
float VIRTUAL_GEAR_RATIO = 10.0f;
float DRIVETRAIN_EFFICIENCY = 0.93f;

// ─── Gear Ratio Tracker ─────────────────────────────────────────
struct GearRatioTracker {
    float current_ratio = 10.0f;
    float confidence = 0.0f;
    float filtered_ratio = 10.0f;
    int stable_samples = 0;
    float last_valid = 10.0f;
    unsigned long last_update = 0;
    float ratio_history[10] = {0};
    int history_index = 0;

    void update(float engine_rpm, float drum_rpm) {
        if (engine_rpm > 800.0f && drum_rpm > 20.0f) {
            float instant_ratio = engine_rpm / drum_rpm;
            if (instant_ratio > 0.5f && instant_ratio < 50.0f) {
                ratio_history[history_index] = instant_ratio;
                history_index = (history_index + 1) % 10;
                float sum = 0, sum_sq = 0;
                int valid_samples = 0;
                for (int i = 0; i < 10; i++) {
                    if (ratio_history[i] > 0.5f) {
                        sum += ratio_history[i];
                        sum_sq += ratio_history[i] * ratio_history[i];
                        valid_samples++;
                    }
                }
                if (valid_samples >= 5) {
                    float mean = sum / valid_samples;
                    float variance = (sum_sq / valid_samples) - (mean * mean);
                    float stability = 1.0f - constrain(variance / (mean * mean), 0.0f, 0.5f);
                    if (stability > 0.98f && valid_samples >= 8) {
                        stable_samples++;
                        confidence = min(1.0f, confidence + 0.02f);
                    } else if (stability > 0.95f) {
                        stable_samples = max(0, stable_samples - 1);
                        confidence = max(0.0f, confidence - 0.005f);
                    } else {
                        stable_samples = max(0, stable_samples - 2);
                        confidence = max(0.0f, confidence - 0.01f);
                    }
                    float filter_strength = (confidence > 0.8f && stability > 0.98f) ? 0.01f : 
                                           (confidence < 0.3f) ? 0.1f : 0.03f;
                    filtered_ratio = filtered_ratio * (1.0f - filter_strength) + instant_ratio * filter_strength;
                    if (confidence > 0.6f) {
                        current_ratio = filtered_ratio;
                        last_valid = filtered_ratio;
                        last_update = millis();
                    }
                }
            }
        } else if (confidence > 0) {
            confidence = max(0.0f, confidence - 0.002f);
            if (millis() - last_update > 5000) filtered_ratio *= 0.999f;
        }
    }
    
    void reset() {
        current_ratio = 10.0f;
        confidence = 0.0f;
        filtered_ratio = 10.0f;
        stable_samples = 0;
        last_valid = 10.0f;
        last_update = 0;
        memset(ratio_history, 0, sizeof(ratio_history));
        history_index = 0;
    }
    
    float getRatio() { return (confidence > 0.5f) ? current_ratio : last_valid; }
    float getConfidence() { return confidence; }
};

GearRatioTracker gearTracker;

// ─── RPM Buffers ───
struct RPMBuffer {
    uint32_t intervals[24] = {0};
    uint8_t  idx = 0;
    uint32_t lastTime = 0;
};

RPMBuffer engineBuffer;
RPMBuffer drumBuffer;

// ─── Live Data ──────────────────────────────────────────────────
struct LiveData {
  float engine_rpm = 0;
  float drum_rpm = 0;
  float hp = 0;
  float engine_torque = 0;
  float drum_torque = 0;
  float peakHP = 0;
  float peakHP_RPM = 0;
  float peakTorque = 0;
  float peakTorque_RPM = 0;
  float air_density = 1.204;
  float ambient_temp = 20.0;
  float pressure = 1013.25;
  float load_kg = 0;
  float gear_ratio = 10.0;
  float drag_torque = 0; 
};
LiveData current;

float atmospheric_pressure = 1013.25;
float ambient_temp_c = 20.0;
float air_density = 1.204;
const float STD_AIR_DENSITY = 1.204;

volatile uint32_t drumPulseInterval = 0;
volatile uint32_t lastDrumPulseTime = 0;
volatile uint32_t lastPulseTime = 0;
volatile uint32_t pulseInterval = 0;

portMUX_TYPE rpmMux = portMUX_INITIALIZER_UNLOCKED;
static bool crossedStart = false;
static int gpioRunState = 0;

// ISR handlers
void IRAM_ATTR hallISR() {
  uint32_t now = micros();
  uint32_t dt = now - engineBuffer.lastTime;
  if (dt < ENGINE_DEBOUNCE_US) return;

  portENTER_CRITICAL_ISR(&rpmMux);
  engineBuffer.intervals[engineBuffer.idx] = dt;
  engineBuffer.idx = (engineBuffer.idx + 1) % 24;
  engineBuffer.lastTime = now;
  portEXIT_CRITICAL_ISR(&rpmMux);
}

void IRAM_ATTR drumISR() {
  uint32_t now = micros();
  uint32_t dt = now - drumBuffer.lastTime;
  if (dt < DRUM_DEBOUNCE_US) return;

  portENTER_CRITICAL_ISR(&rpmMux);
  drumBuffer.intervals[drumBuffer.idx] = dt;
  drumBuffer.idx = (drumBuffer.idx + 1) % 24;
  drumBuffer.lastTime = now;
  portEXIT_CRITICAL_ISR(&rpmMux);
}

void IRAM_ATTR gpioTriggerISR() {
  unsigned long now = millis();
  if (now - lastGpioDebounceTime > GPIO_DEBOUNCE_MS) {
    gpioTriggerPressed = true;
    lastGpioDebounceTime = now;
  }
}

// Recording state
bool recording = false;
bool autoMode = false;
float autoStartRPM = 0;
float autoEndRPM = 0;

struct RunPoint { float rpm, hp, torque; };
std::vector<RunPoint> currentRun;

float runPeakHP = 0, runPeakHP_RPM = 0;
float runPeakTorque = 0, runPeakTorque_RPM = 0;

bool settingsPageActive = false;

Preferences preferences;

String sta_ssid = "";
String sta_password = "";
String currentCustomer = "";
String currentUnit = "";
String currentComments = "";

unsigned long recordingStartedMs = 0;

TaskHandle_t calibTask = NULL;
volatile float calibKnownWeight = 0.0f;
volatile bool calibRequest = false;

int32_t zeroOffset = 0;

SmoothingBuffer hpSmooth;
SmoothingBuffer torqueSmooth;

// Function prototypes
void broadcastSettings();
void sendRunComplete(float peakHP, float peakHP_RPM, float peakTorque, float peakTorque_RPM);
void saveFullRun();
void loadRunList();
void sendRunList();
void sendRunData(String filename);
void processGpioTrigger();

// ─── Fast median filter ─────────────────────────────────────────
float medianGeneric(float *arr, int size) {
  float temp[size];
  for (int i = 0; i < size; i++) temp[i] = arr[i];
  
  for (int i = 0; i < size - 1; i++) {
    for (int j = 0; j < size - i - 1; j++) {
      if (temp[j] > temp[j + 1]) {
        float swap = temp[j];
        temp[j] = temp[j + 1];
        temp[j + 1] = swap;
      }
    }
  }
  return temp[size / 2];
}

// ─── Inertia calculation ────────────────────────────────────────
float calculateDrumInertia() {
  float radius = DRUM_DIAM_METERS / 2.0f;
  if (DRUM_INERTIA_TYPE == "solid") return 0.5f * DRUM_MASS * radius * radius;
  if (DRUM_INERTIA_TYPE == "hollow") {
    float inner_r = radius - (DRUM_WALL_THICKNESS_CM / 100.0f);
    if (inner_r <= 0) inner_r = 0.01f;
    return 0.5f * DRUM_MASS * (radius * radius + inner_r * inner_r);
  }
  return DRUM_INERTIA_CUSTOM;
}

struct RunInfo {
  uint32_t timestamp;
  String filename;
  String customer;
  String vehicle;
  float peakHP;
  float peakTorque;
  float peakEGT;
  String unitSystem;
};

std::vector<RunInfo> savedRuns;

// ─── Save/Load runs ──────────────────────────────────────
void saveFullRun() {
  if (currentRun.empty()) return;
  
  uint32_t timestamp = (runTimestamp > 0) ? runTimestamp : millis();
  String filename = "/run_" + String(timestamp) + ".csv";
  
  File file = LittleFS.open(filename, "w");
  if (!file) return;
  file.println("RPM,HP,Torque");
  int step = (currentRun.size() > 900) ? (currentRun.size() / 900) : 1;
  for (size_t i = 0; i < currentRun.size(); i += step) {
    file.println(String(currentRun[i].rpm) + "," + String(currentRun[i].hp, 2) + "," + String(currentRun[i].torque, 2));
  }
  file.close();
  
  String unitSystem = useImperial ? "imperial" : "metric";
  
  File idxFile = LittleFS.open("/runs.idx", "a");
  if (idxFile) {
    idxFile.println(String(timestamp) + "," + filename + "," + currentCustomer + "," + currentUnit + "," +
                    String(runPeakHP, 2) + "," + String(runPeakTorque, 2) + "," + String(egtTracker.peak, 0) +
                    "," + unitSystem);
    idxFile.close();
  }
  
  runTimestamp = 0;
  
  loadRunList();
  while (savedRuns.size() > MAX_SAVED_RUNS) {
    LittleFS.remove(savedRuns.back().filename);
    savedRuns.pop_back();
  }
  File newIdx = LittleFS.open("/runs.idx.tmp", "w");
  for (const auto& run : savedRuns) {
    newIdx.println(String(run.timestamp) + "," + run.filename + "," + run.customer + "," + run.vehicle + "," +
                   String(run.peakHP, 2) + "," + String(run.peakTorque, 2) + "," + String(run.peakEGT, 0) +
                   "," + run.unitSystem);
  }
  newIdx.close();
  LittleFS.remove("/runs.idx");
  LittleFS.rename("/runs.idx.tmp", "/runs.idx");
}

void loadRunList() {
  savedRuns.clear();
  File idxFile = LittleFS.open("/runs.idx", "r");
  if (!idxFile) return;
  while (idxFile.available()) {
    String line = idxFile.readStringUntil('\n');
    line.trim();
    if (line.length() == 0) continue;
    int comma1 = line.indexOf(',');
    int comma2 = line.indexOf(',', comma1 + 1);
    int comma3 = line.indexOf(',', comma2 + 1);
    int comma4 = line.indexOf(',', comma3 + 1);
    int comma5 = line.indexOf(',', comma4 + 1);
    int comma6 = line.indexOf(',', comma5 + 1);
    int comma7 = line.indexOf(',', comma6 + 1);
    
    if (comma1 > 0 && comma2 > 0) {
      RunInfo run;
      run.timestamp = line.substring(0, comma1).toInt();
      run.filename = line.substring(comma1 + 1, comma2);
      run.customer = line.substring(comma2 + 1, comma3);
      run.vehicle = line.substring(comma3 + 1, comma4);
      run.peakHP = line.substring(comma4 + 1, comma5).toFloat();
      run.peakTorque = line.substring(comma5 + 1, comma6).toFloat();
      run.peakEGT = line.substring(comma6 + 1, comma7).toFloat();
      if (comma7 > 0) {
        run.unitSystem = line.substring(comma7 + 1);
        run.unitSystem.trim();
      } else {
        run.unitSystem = "metric";
      }
      savedRuns.push_back(run);
    }
  }
  idxFile.close();
}

void sendRunList() {
  String json = "{\"type\":\"runs\",\"data\":[";
  for (size_t i = 0; i < savedRuns.size(); i++) {
    if (i > 0) json += ",";
    json += "{\"timestamp\":" + String(savedRuns[i].timestamp) + ",\"filename\":\"" + savedRuns[i].filename + "\"" +
            ",\"customer\":\"" + savedRuns[i].customer + "\",\"vehicle\":\"" + savedRuns[i].vehicle + "\"" +
            ",\"hp\":" + String(savedRuns[i].peakHP, 1) + ",\"torque\":" + String(savedRuns[i].peakTorque, 1) +
            ",\"egt\":" + String(savedRuns[i].peakEGT, 0) +
            ",\"unitSystem\":\"" + savedRuns[i].unitSystem + "\"}";
  }
  json += "]}";
  ws.textAll(json);
}

void sendRunData(String filename) {
  File file = LittleFS.open(filename, "r");
  if (!file) return;
  
  String unitSystem = "metric";
  for (const auto& run : savedRuns) {
    if (run.filename == filename) {
      unitSystem = run.unitSystem;
      break;
    }
  }
  
  String json = "{\"type\":\"run_data\",\"filename\":\"" + filename + "\",\"unitSystem\":\"" + unitSystem + "\",\"points\":[";
  bool first = true;
  file.readStringUntil('\n');
  while (file.available()) {
    String line = file.readStringUntil('\n');
    line.trim();
    if (line.length() == 0) continue;
    int comma1 = line.indexOf(',');
    int comma2 = line.indexOf(',', comma1 + 1);
    if (comma1 > 0 && comma2 > 0) {
      if (!first) json += ",";
      float rpm = line.substring(0, comma1).toFloat();
      float hp = line.substring(comma1 + 1, comma2).toFloat();
      float torque = line.substring(comma2 + 1).toFloat();
      json += "{\"r\":" + String(rpm, 0) + ",\"h\":" + String(hp, 2) + ",\"t\":" + String(torque, 2) + "}";
      first = false;
    }
  }
  json += "]}";
  file.close();
  ws.textAll(json);
}

void deleteRun(String filename, int index) {
  LittleFS.remove(filename);
  if (index >= 0 && index < (int)savedRuns.size()) savedRuns.erase(savedRuns.begin() + index);
  File idxFile = LittleFS.open("/runs.idx", "w");
  for (const auto& run : savedRuns) {
    idxFile.println(String(run.timestamp) + "," + run.filename + "," + run.customer + "," + run.vehicle + "," +
                    String(run.peakHP, 2) + "," + String(run.peakTorque, 2) + "," + String(run.peakEGT, 0) +
                    "," + run.unitSystem);
  }
  idxFile.close();
}

// ─── Broadcast Settings ─────────────────────────────────────────
void broadcastSettings() {
  String json = "{";
  json += "\"type\":\"settings\",";
  json += "\"version\":\"" + String(FIRMWARE_VERSION) + "\",";
  json += "\"buildDate\":\"" + String(BUILD_DATE) + "\",";
  json += "\"maxRPM\":" + String(MAX_RPM, 0) + ",";
  json += "\"maxHP\":" + String(MAX_HP, 1) + ",";
  json += "\"maxTorque\":" + String(MAX_TORQUE, 1) + ",";
  json += "\"armCm\":" + String(LEVER_ARM_METERS * 100, 1) + ",";
  json += "\"calFactor\":" + String(CALIBRATION_FACTOR, 6) + ",";
  json += "\"pulsesPerRev\":" + String(PULSES_PER_REVOLUTION, 1) + ",";
  json += "\"rpmFilter\":" + String(RPM_FILTER_ALPHA, 2) + ",";
  json += "\"mode\":\"" + DYNO_MODE + "\",";
  json += "\"rpmSource\":\"" + RPM_SOURCE + "\",";
  json += "\"gearRatioMode\":\"" + String(GEAR_RATIO_AUTO ? "auto" : "manual") + "\",";
  json += "\"gearRatioManual\":" + String(GEAR_RATIO_MANUAL, 2) + ",";
  json += "\"brakeToEngineRatio\":" + String(BRAKE_TO_ENGINE_RATIO, 2) + ",";
  json += "\"drumMass\":" + String(DRUM_MASS, 1) + ",";
  json += "\"drumCm\":" + String(DRUM_DIAM_METERS * 100, 1) + ",";
  json += "\"pulsesPerRevDrum\":" + String(PULSES_PER_REV_DRUM, 1) + ",";
  json += "\"drumInertiaType\":\"" + DRUM_INERTIA_TYPE + "\",";
  json += "\"drumWallCm\":" + String(DRUM_WALL_THICKNESS_CM, 1) + ",";
  json += "\"drumInertiaCustom\":" + String(DRUM_INERTIA_CUSTOM, 3) + ",";
  json += "\"drivetrainEff\":" + String(DRIVETRAIN_EFFICIENCY * 100, 1) + ",";
  json += "\"maxSafeDrumRpm\":" + String(MAX_SAFE_DRUM_RPM, 0) + ",";
  json += "\"engDebounce\":" + String(ENGINE_DEBOUNCE_US) + ",";
  json += "\"drumDebounce\":" + String(DRUM_DEBOUNCE_US) + ",";
  json += "\"medianSize\":" + String(MEDIAN_SIZE) + ",";
  json += "\"gpioTriggerPin\":" + String(GPIO_TRIGGER_PIN);
  json += "}";
  ws.textAll(json);
}

// ─── JSON SENDER ─────────────────────────────────────────────────
void notifyClients(const LiveData& data, bool isRecording) {
  static char jsonBuffer[512];
  
  int len = snprintf(jsonBuffer, sizeof(jsonBuffer),
    "{\"er\":%.0f,\"dr\":%.1f,\"hp\":%.2f,\"et\":%.2f,\"dt\":%.2f,\"l\":%.2f,\"gr\":%.2f,\"rc\":%.2f,\"ad\":%.3f,\"at\":%.1f,\"pr\":%.1f",
    data.engine_rpm, data.drum_rpm, data.hp, data.engine_torque, data.drum_torque,
    data.load_kg, data.gear_ratio, gearTracker.getConfidence(), data.air_density,
    data.ambient_temp, data.pressure);
  
  if (isRecording) {
    len += snprintf(jsonBuffer + len, sizeof(jsonBuffer) - len, ",\"rec\":1");
  }
  
  if (isRecording) {
    len += snprintf(jsonBuffer + len, sizeof(jsonBuffer) - len, ",\"ph\":%.2f,\"prpm\":%.0f,\"pt\":%.2f,\"ptrpm\":%.0f",
      runPeakHP, runPeakHP_RPM, runPeakTorque, runPeakTorque_RPM);
  }
  
  static int egt_counter = 0;
  if (mcp_egt_ok && ++egt_counter >= 3) {
    egt_counter = 0;
    len += snprintf(jsonBuffer + len, sizeof(jsonBuffer) - len, ",\"eg\":%.1f", egtTracker.getFiltered());
  }
  
  snprintf(jsonBuffer + len, sizeof(jsonBuffer) - len, ",\"drag\":%.2f}", data.drag_torque);
  
  ws.textAll(jsonBuffer);
}

void sendRunComplete(float peakHP, float peakHP_RPM, float peakTorque, float peakTorque_RPM) {
  saveFullRun();
  String json = "{\"type\":\"run_complete\",\"peakHP\":" + String(peakHP, 2) +
                ",\"peakHP_RPM\":" + String(peakHP_RPM, 0) +
                ",\"peakTorque\":" + String(peakTorque, 2) +
                ",\"peakTorque_RPM\":" + String(peakTorque_RPM, 0) +
                ",\"peakEGT\":" + String(egtTracker.peak, 0) +
                ",\"finalEGT\":" + String(egtTracker.final, 0) +
                ",\"ambient\":" + String(egtTracker.ambient, 1) + "}";
  ws.textAll(json);
  egtTracker.reset();
  sendRunList();
}

String getSensorStatusJSON() {
  return "{\"nau7802\":" + String(nau_initialized ? "true" : "false") +
         ",\"mcp9600\":" + String(mcp_egt_ok ? "true" : "false") +
         ",\"bmp280\":" + String(bmp_initialized ? "true" : "false") + "}";
}

void calibrationTask(void *pvParameters) {
  for (;;) {
    if (calibRequest) {
      calibRequest = false;
      ws.textAll("{\"type\":\"info\",\"msg\":\"Calibrating... hold weight still for 5 seconds\"}");
      vTaskDelay(pdMS_TO_TICKS(2000));
      long raw_sum = 0;
      int valid_count = 0;
      const unsigned long durationMs = 5000UL;
      unsigned long startTime = millis();
      while (millis() - startTime < durationMs) {
        if (nau_initialized && nau.available()) {
          long current = nau.read() - zeroOffset;
          if (labs(current) > 100000) {
            if (valid_count >= 10) raw_sum += current;
            valid_count++;
          }
        }
        vTaskDelay(pdMS_TO_TICKS(50));
      }
      int averaged_samples = valid_count - 10;
      if (averaged_samples < 5) {
        ws.textAll("{\"type\":\"info\",\"msg\":\"Not enough stable readings. Try again.\"}");
      } else {
        double avgRaw = (double)raw_sum / averaged_samples;
        CALIBRATION_FACTOR = (calibKnownWeight * 1000.0f) / avgRaw;
        preferences.begin("dyno", false);
        preferences.putFloat("cal_factor", CALIBRATION_FACTOR);
        preferences.end();
        broadcastSettings();
        ws.textAll("{\"type\":\"info\",\"msg\":\"Calibration done! New factor: " + String(CALIBRATION_FACTOR, 6) + "\"}");
      }
    }
    vTaskDelay(pdMS_TO_TICKS(50));
  }
}

void processGpioTrigger() {
  if (gpioTriggerPressed) {
    gpioTriggerPressed = false;
    
    if (gpioRunState == 0) {
      gpioRunState = 1;
      ws.textAll("{\"type\":\"show_run_dialog\"}");
    } else if (gpioRunState == 1) {
      gpioRunState = 2;
      ws.textAll("{\"type\":\"gpio_start_run\"}");
    } else if (gpioRunState == 2) {
      gpioRunState = 0;
      if (recording) {
        recording = false;
        autoMode = false;
        crossedStart = false;
        recordingStartedMs = 0;
        sendRunComplete(runPeakHP, runPeakHP_RPM, runPeakTorque, runPeakTorque_RPM);
      }
      ws.textAll("{\"type\":\"info\",\"msg\":\"Run finished via GPIO\"}");
    }
  }
}

void onWsEvent(AsyncWebSocket *server, AsyncWebSocketClient *client, AwsEventType type, void *arg, uint8_t *data, size_t len) {
  if (type == WS_EVT_CONNECT) {
    client->text("{\"type\":\"sensor_status\",\"data\":" + getSensorStatusJSON() + "}");
  }
  if (type == WS_EVT_DATA) {
    String msg = "";
    for (size_t i = 0; i < len; i++) msg += (char)data[i];
    msg.trim();

    if (msg == "TARE") {
      ws.textAll("{\"type\":\"info\",\"msg\":\"Taring...\"}");
      vTaskDelay(pdMS_TO_TICKS(1000));
      long sum = 0;
      int count = 0;
      for (int i = 0; i < 16; i++) {
        if (nau_initialized && nau.available()) {
          sum += nau.read();
          count++;
          vTaskDelay(pdMS_TO_TICKS(5));
        }
      }
      if (count > 0) {
        zeroOffset = sum / count;
        ws.textAll("{\"type\":\"info\",\"msg\":\"Tare complete\"}");
      } else {
        ws.textAll("{\"type\":\"info\",\"msg\":\"Tare failed\"}");
      }
      return;
    }

    if (msg == "RESET_RATIO") {
      gearTracker.reset();
      ws.textAll("{\"type\":\"info\",\"msg\":\"Gear ratio tracking reset\"}");
      return;
    }

    if (msg.startsWith("SET_CUSTOMER:")) {
      String data = msg.substring(13);
      int pipe1 = data.indexOf('|');
      int pipe2 = data.indexOf('|', pipe1 + 1);
      if (pipe1 > 0 && pipe2 > 0) {
        currentCustomer = data.substring(0, pipe1);
        currentUnit = data.substring(pipe1 + 1, pipe2);
        currentComments = data.substring(pipe2 + 1);
      }
      return;
    }
    
    if (msg.startsWith("SET_RUN_TIMESTAMP:")) {
      runTimestamp = msg.substring(18).toInt();
      Serial.printf("Run timestamp set: %u\n", runTimestamp);
      return;
    }

    if (msg.startsWith("CALIBRATE:")) {
      float knownWeightKg = msg.substring(10).toFloat();
      if (knownWeightKg <= 0) {
        ws.textAll("{\"type\":\"info\",\"msg\":\"Enter valid weight > 0\"}");
        return;
      }
      calibKnownWeight = knownWeightKg;
      calibRequest = true;
      ws.textAll("{\"type\":\"info\",\"msg\":\"Calibration starting...\"}");
      return;
    }

    if (msg.startsWith("SET_SINGLE:")) {
      String part = msg.substring(11);
      int eq = part.indexOf('=');
      if (eq == -1) return;
      String key = part.substring(0, eq);
      String valStr = part.substring(eq + 1);
      float val = valStr.toFloat();

      if (key == "maxRPM" && val >= 5000 && val <= 50000) { MAX_RPM = val; preferences.putFloat("max_rpm", MAX_RPM); }
      else if (key == "maxHP" && val >= 5 && val <= 1000) { MAX_HP = val; preferences.putFloat("max_hp", MAX_HP); }
      else if (key == "maxTorque" && val >= 5 && val <= 1000) { MAX_TORQUE = val; preferences.putFloat("max_torque", MAX_TORQUE); }
      else if (key == "armCm" && val >= 5 && val <= 100) { LEVER_ARM_METERS = val / 100.0f; preferences.putFloat("lever_arm", LEVER_ARM_METERS); }
      else if (key == "pulsesPerRev" && val >= 0.5f && val <= 10.0f) { PULSES_PER_REVOLUTION = val; preferences.putFloat("pulses_per_rev", val); }
      else if (key == "brakeToEngineRatio" && val >= 0.1f && val <= 50.0f) { BRAKE_TO_ENGINE_RATIO = val; preferences.putFloat("brake_eng_ratio", val); }
      else if (key == "rpmFilter" && val >= 0.1f && val <= 0.9f) { RPM_FILTER_ALPHA = val; preferences.putFloat("rpm_filter", val); }
      else if (key == "autoHyst" && val >= 50 && val <= 500) { AUTO_HYSTERESIS = val; preferences.putFloat("auto_hyst", val); }
      else if (key == "dynoMode") { if (valStr == "brake" || valStr == "inertia") { DYNO_MODE = valStr; preferences.putString("dyno_mode", DYNO_MODE); } }
      else if (key == "drumMass" && val >= 1.0f && val <= 300.0f) { DRUM_MASS = val; preferences.putFloat("drum_mass", DRUM_MASS); }
      else if (key == "drumCm" && val >= 10.0f && val <= 300.0f) { DRUM_DIAM_METERS = val / 100.0f; preferences.putFloat("drum_diam", DRUM_DIAM_METERS); }
      else if (key == "pulsesPerRevDrum" && val >= 0.5f && val <= 20.0f) { PULSES_PER_REV_DRUM = val; preferences.putFloat("pulses_drum", PULSES_PER_REV_DRUM); }
      else if (key == "drumInertiaType") { if (valStr == "solid" || valStr == "hollow" || valStr == "custom") { DRUM_INERTIA_TYPE = valStr; preferences.putString("drum_inertia_type", DRUM_INERTIA_TYPE); } }
      else if (key == "drumWallCm" && val >= 0.5f && val <= 20.0f) { DRUM_WALL_THICKNESS_CM = val; preferences.putFloat("drum_wall_cm", DRUM_WALL_THICKNESS_CM); }
      else if (key == "drumInertiaCustom" && val >= 0.01f && val <= 100.0f) { DRUM_INERTIA_CUSTOM = val; preferences.putFloat("drum_inertia_custom", DRUM_INERTIA_CUSTOM); }
      else if (key == "drivetrainEff" && val >= 70.0f && val <= 99.0f) { DRIVETRAIN_EFFICIENCY = val / 100.0f; preferences.putFloat("drivetrain_eff", DRIVETRAIN_EFFICIENCY); }
      else if (key == "maxSafeDrumRpm" && val >= 500 && val <= 25000) { MAX_SAFE_DRUM_RPM = val; preferences.putFloat("max_safe_drum_rpm", MAX_SAFE_DRUM_RPM); }
      else if (key == "rpmSource") { if (valStr == "spark" || valStr == "brake" || valStr == "drum") { RPM_SOURCE = valStr; preferences.putString("rpm_source", RPM_SOURCE); } }
      else if (key == "gearRatioMode") { GEAR_RATIO_AUTO = (valStr == "auto"); preferences.putBool("gear_ratio_auto", GEAR_RATIO_AUTO); }
      else if (key == "gearRatioManual" && val >= 0.5f && val <= 100.0f) { GEAR_RATIO_MANUAL = val; preferences.putFloat("gear_ratio_manual", GEAR_RATIO_MANUAL); }
      else if (key == "engDebounce" && val >= 100 && val <= 5000) { 
        ENGINE_DEBOUNCE_US = (int)val; 
        preferences.putInt("eng_debounce", ENGINE_DEBOUNCE_US); 
      }
      else if (key == "drumDebounce" && val >= 50 && val <= 2000) { 
        DRUM_DEBOUNCE_US = (int)val; 
        preferences.putInt("drum_debounce", DRUM_DEBOUNCE_US); 
      }
      else if (key == "medianSize" && val >= 3 && val <= 9 && ((int)val % 2 == 1)) { 
        MEDIAN_SIZE = (int)val; 
        preferences.putInt("median_size", MEDIAN_SIZE);
        rpmMedian.resize(MEDIAN_SIZE);
        drumMedian.resize(MEDIAN_SIZE);
        std::fill(rpmMedian.begin(), rpmMedian.end(), 0.0f);
        std::fill(drumMedian.begin(), drumMedian.end(), 0.0f); 
      }

      broadcastSettings();
      return;
    }

    if (msg.startsWith("WIFI_CLIENT:")) {
      String payload = msg.substring(12);
      int split = payload.indexOf(':');
      if (split != -1) {
        sta_ssid = payload.substring(0, split);
        sta_password = payload.substring(split + 1);
        sta_ssid.trim();
        sta_password.trim();
        preferences.begin("dyno", false);
        preferences.putString("sta_ssid", sta_ssid);
        preferences.putString("sta_password", sta_password);
        preferences.end();
        ws.textAll("{\"type\":\"info\",\"msg\":\"WiFi saved. Rebooting...\"}");
        vTaskDelay(pdMS_TO_TICKS(1000));
        ESP.restart();
      }
      return;
    }

    if (msg == "RESET_DEFAULT") {
      preferences.begin("dyno", false);
      preferences.clear();
      preferences.end();
      ws.textAll("{\"type\":\"info\",\"msg\":\"WiFi cleared. Rebooting to AP...\"}");
      vTaskDelay(pdMS_TO_TICKS(1500));
      ESP.restart();
      return;
    }

    if (msg == "RESET_ALL_SETTINGS") {
      preferences.clear();
      LittleFS.format();
      ws.textAll("{\"type\":\"info\",\"msg\":\"All settings reset. Rebooting...\"}");
      vTaskDelay(pdMS_TO_TICKS(1000));
      ESP.restart();
      return;
    }

    if (msg.startsWith("AUTO_RUN:")) {
      int colon1 = msg.indexOf(':', 9);
      if (colon1 != -1) {
        autoStartRPM = msg.substring(9, colon1).toFloat();
        autoEndRPM = msg.substring(colon1 + 1).toFloat();
        if (autoStartRPM > 0 && autoEndRPM > autoStartRPM) {
          autoMode = true;
          crossedStart = false;
          ws.textAll("{\"type\":\"info\",\"msg\":\"Auto Run activated\"}");
        }
      }
      return;
    }

    if (msg == "GET_SETTINGS") { broadcastSettings(); return; }
    if (msg == "GET_SENSOR_STATUS") { client->text("{\"type\":\"sensor_status\",\"data\":" + getSensorStatusJSON() + "}"); return; }
    if (msg == "GET_RUNS") { sendRunList(); return; }
    if (msg.startsWith("GET_RUN_DATA:")) { sendRunData(msg.substring(13)); return; }
    if (msg.startsWith("DELETE_RUN:")) { int idx = msg.substring(11).toInt(); if (idx >= 0 && idx < (int)savedRuns.size()) { deleteRun(savedRuns[idx].filename, idx); sendRunList(); ws.textAll("{\"type\":\"info\",\"msg\":\"Run deleted\"}"); } return; }

    if (msg == "SET_UNITS_IMPERIAL") { useImperial = true; return; }
    if (msg == "SET_UNITS_METRIC") { useImperial = false; return; } 
    if (msg == "PAGE_SETTINGS") { settingsPageActive = true; return; }
    if (msg == "PAGE_MAIN") { settingsPageActive = false; return; }

    if (msg == "CANCEL_AUTO_RUN") {
      gpioRunState = 0;
      if (autoMode || recording) {
        autoMode = false;
        crossedStart = false;
        if (recording) {
          recording = false;
          recordingStartedMs = 0;
          sendRunComplete(runPeakHP, runPeakHP_RPM, runPeakTorque, runPeakTorque_RPM);
        }
        currentRun.clear();
        runPeakHP = 0; runPeakHP_RPM = 0;
        runPeakTorque = 0; runPeakTorque_RPM = 0;
        egtTracker.reset();
        hpSmooth.reset();
        torqueSmooth.reset();
        ws.textAll("{\"type\":\"info\",\"msg\":\"Run cancelled\"}");
      }
      return;
    }

    if (msg == "MANUAL_START") {
      if (!recording) {
        recording = true;
        autoMode = false;
        crossedStart = false;
        currentRun.clear();
        runPeakHP = 0;
        runPeakHP_RPM = 0;
        runPeakTorque = 0;
        runPeakTorque_RPM = 0;
        egtTracker.reset();
        hpSmooth.reset();
        torqueSmooth.reset();
        recordingStartedMs = millis();
        ws.textAll("{\"type\":\"info\",\"msg\":\"Manual recording started\"}");
      }
      return;
    }

    if (msg == "MANUAL_STOP") {
      gpioRunState = 0;
      if (recording) {
        recording = false;
        autoMode = false;
        crossedStart = false;
        recordingStartedMs = 0;
        sendRunComplete(runPeakHP, runPeakHP_RPM, runPeakTorque, runPeakTorque_RPM);
      }
      return;
    }

    if (msg.startsWith("SET_GPIO_PIN:")) {
      int pin = msg.substring(13).toInt();
      bool isReserved = false;
      for (int i = 0; i < RESERVED_COUNT; i++) if (pin == RESERVED_PINS[i]) { isReserved = true; break; }
      if (pin >= 0 && pin <= 39 && !isReserved && pin != HALL_PIN && pin != DRUM_PIN && pin != I2C_SDA && pin != I2C_SCL) {
        detachInterrupt(digitalPinToInterrupt(GPIO_TRIGGER_PIN));
        GPIO_TRIGGER_PIN = pin;
        pinMode(GPIO_TRIGGER_PIN, INPUT_PULLUP);
        attachInterrupt(digitalPinToInterrupt(GPIO_TRIGGER_PIN), gpioTriggerISR, FALLING);
        preferences.begin("dyno", false);
        preferences.putInt("gpio_trigger_pin", GPIO_TRIGGER_PIN);
        preferences.end();
        ws.textAll("{\"type\":\"info\",\"msg\":\"GPIO pin set to " + String(pin) + "\"}");
      }
      return;
    }

    if (msg == "GET_GPIO_STATUS") {
      client->text("{\"type\":\"gpio_status\",\"pin\":" + String(GPIO_TRIGGER_PIN) + ",\"enabled\":true}");
      return;
    }

    if (msg == "TEST_GPIO_TRIGGER") {
      gpioTriggerPressed = true;
      ws.textAll("{\"type\":\"info\",\"msg\":\"GPIO trigger test\"}");
      return;
    }
  }
}

class CaptiveRequestHandler : public AsyncWebHandler {
public:
  bool canHandle(AsyncWebServerRequest *request) { return true; }
  void handleRequest(AsyncWebServerRequest *request) { request->redirect("/"); }
};

// ─── Sensor Task ─────────────────────────────────────────────────
void sensorTask(void * parameter) {
  static bool egt_first_read = true;
  static unsigned long last_ws_update = 0;
  static unsigned long last_egt_read = 0;
  static unsigned long last_gear_calc = 0;
  static unsigned long last_bmp_read = 0;
  static unsigned long last_accel_time = 0;

  for (;;) {
    unsigned long now = millis();

    if (!settingsPageActive) {

      // ==================== LOAD CELL ====================
      float grams = 0.0f;
      if (nau_initialized && nau.available()) {
        long raw = nau.read();
        grams = (float)(raw - zeroOffset) * CALIBRATION_FACTOR;
      }
      float kg = constrain(grams / 1000.0f, 0.0f, 200.0f);
      current.load_kg = kg;
      float force_N = kg * 9.81f;

      // ==================== BMP280 ====================
      if (bmp_initialized && (now - last_bmp_read >= BMP_UPDATE_MS)) {
        atmospheric_pressure = bmp.readPressure() / 100.0F;
        ambient_temp_c = bmp.readTemperature();
        float temp_kelvin = ambient_temp_c + 273.15f;
        if (temp_kelvin > 0 && atmospheric_pressure > 0) {
          air_density = (atmospheric_pressure * 100.0f) / (287.05f * temp_kelvin);
          air_density = constrain(air_density, 1.0f, 1.3f);
        }
        current.air_density = air_density;
        current.ambient_temp = ambient_temp_c;
        current.pressure = atmospheric_pressure;
        last_bmp_read = now;
      }

      // ==================== ENGINE RPM ====================
      uint32_t sum = 0; int valid = 0;
      portENTER_CRITICAL(&rpmMux);
      for (int i = 0; i < 24; i++) {
        if (engineBuffer.intervals[i] > 80) {
          sum += engineBuffer.intervals[i];
          valid++;
        }
      }
      uint32_t lastEngineTime = engineBuffer.lastTime;
      portEXIT_CRITICAL(&rpmMux);

      uint32_t now_us = micros();
      uint32_t timeSinceLast = now_us - lastEngineTime;

      float pulsesPerRevSafe = max(PULSES_PER_REVOLUTION, 0.5f);
      float engine_rpm_raw = 0.0f;

      if (timeSinceLast > 800000UL) {
        engine_rpm_raw = 0.0f;
      } else if (valid >= 4) {
        float avgInterval = (float)sum / valid;
        engine_rpm_raw = 60000000.0f / (avgInterval * pulsesPerRevSafe);
      }

      rpmMedian[rpmMedIdx] = engine_rpm_raw;
      rpmMedIdx = (rpmMedIdx + 1) % MEDIAN_SIZE;
      float median_rpm = medianGeneric(rpmMedian.data(), MEDIAN_SIZE);

      float alpha = (median_rpm > 5000) ? 0.28f : RPM_FILTER_ALPHA;
      engine_rpm_filtered_g = alpha * median_rpm + (1.0f - alpha) * engine_rpm_filtered_g;

      ultra_smooth_rpm_g = 0.09f * engine_rpm_filtered_g + 0.91f * ultra_smooth_rpm_g;
      float engine_rpm = constrain(ultra_smooth_rpm_g, 0.0f, MAX_RPM * 1.08f);

      // ==================== DRUM RPM ====================
      sum = 0; valid = 0;
      portENTER_CRITICAL(&rpmMux);
      for (int i = 0; i < 24; i++) {
        if (drumBuffer.intervals[i] > 60) {
          sum += drumBuffer.intervals[i];
          valid++;
        }
      }
      uint32_t lastDrumTime = drumBuffer.lastTime;
      portEXIT_CRITICAL(&rpmMux);

      uint32_t timeSinceLastDrum = now_us - lastDrumTime;
      float pulsesPerRevDrumSafe = max(PULSES_PER_REV_DRUM, 0.5f);
      float drum_rpm_raw = 0.0f;

      // Inertia mode: relaxed guards for low-RPM tracking
      if (DYNO_MODE == "inertia") {
        if (timeSinceLastDrum > 3000000UL) {
          drum_rpm_raw = 0.0f;
        } else if (valid >= 2) {
          float avgInterval = (float)sum / valid;
          drum_rpm_raw = 60000000.0f / (avgInterval * pulsesPerRevDrumSafe);
        }
        drumMedian[drumMedIdx] = drum_rpm_raw;
        drumMedIdx = (drumMedIdx + 1) % MEDIAN_SIZE;
        float median_drum = medianGeneric(drumMedian.data(), MEDIAN_SIZE);
        drum_rpm_filtered_g = 0.35f * median_drum + 0.65f * drum_rpm_filtered_g;
        ultra_smooth_drum_g = 0.20f * drum_rpm_filtered_g + 0.80f * ultra_smooth_drum_g;
      } else {
        // Brake mode: original tighter timing
        if (timeSinceLastDrum > 1000000UL) {
          drum_rpm_raw = 0.0f;
        } else if (valid >= 4) {
          float avgInterval = (float)sum / valid;
          drum_rpm_raw = 60000000.0f / (avgInterval * pulsesPerRevDrumSafe);
        }
        drumMedian[drumMedIdx] = drum_rpm_raw;
        drumMedIdx = (drumMedIdx + 1) % MEDIAN_SIZE;
        float median_drum = medianGeneric(drumMedian.data(), MEDIAN_SIZE);
        drum_rpm_filtered_g = 0.22f * median_drum + 0.78f * drum_rpm_filtered_g;
        ultra_smooth_drum_g = 0.10f * drum_rpm_filtered_g + 0.90f * ultra_smooth_drum_g;
      }

      float drum_rpm = constrain(ultra_smooth_drum_g, 0.0f, MAX_SAFE_DRUM_RPM);

      // ==================== GEAR RATIO TRACKING ====================
      if (now - last_gear_calc >= GEAR_UPDATE_MS) {
        gearTracker.update(engine_rpm, drum_rpm);
        last_gear_calc = now;
      }

      // Drum mode always uses manual ratio (no engine sensor to auto-track)
      if (RPM_SOURCE == "drum") {
        current_gear_ratio = GEAR_RATIO_MANUAL;
      } else if (GEAR_RATIO_AUTO) {
        if (DYNO_MODE == "brake") {
          current_gear_ratio = (gearTracker.getConfidence() > 0.6f) ? gearTracker.getRatio() : BRAKE_TO_ENGINE_RATIO;
        } else {
          current_gear_ratio = (gearTracker.getConfidence() > 0.4f) ? gearTracker.getRatio() : GEAR_RATIO_MANUAL;
        }
      } else {
        current_gear_ratio = GEAR_RATIO_MANUAL;
      }
      current.gear_ratio = current_gear_ratio;

      // ==================== POWER CALCULATION ====================
      float hp = 0.0f;
      float engine_torque = 0.0f;
      float drum_torque = 0.0f;
      float display_engine_rpm = engine_rpm;

      if (DYNO_MODE == "brake") {
        drum_torque = constrain(force_N * LEVER_ARM_METERS, 0.0f, 2000.0f);

        engine_torque = (current_gear_ratio > 0.1f) ? drum_torque / current_gear_ratio : drum_torque;
        engine_torque /= DRIVETRAIN_EFFICIENCY;
        engine_torque = constrain(engine_torque, 0.0f, 2000.0f);

        if (engine_rpm >= MIN_RPM_THRESHOLD_BRAKE && engine_rpm > 0) {
          hp = (engine_torque * engine_rpm) / 7121.0f;
        }

        if (RPM_SOURCE == "brake") {
          display_engine_rpm = drum_rpm * BRAKE_TO_ENGINE_RATIO;
        } else {
          display_engine_rpm = engine_rpm;
        }

      } else if (DYNO_MODE == "inertia") {
        float drum_inertia = calculateDrumInertia();
        
        float effective_drum_rpm = 0.0f;
        // Spark mode with working engine sensor: derive drum RPM from engine
        if (RPM_SOURCE == "spark" && current_gear_ratio > 0.1f && engine_rpm > 60.0f) {
            effective_drum_rpm = engine_rpm / current_gear_ratio;
        } else {
            // Drum mode (or spark with no engine sensor): use raw drum RPM
            effective_drum_rpm = drum_rpm;
        }
        
        if (effective_drum_rpm > MAX_SAFE_DRUM_RPM) effective_drum_rpm = MAX_SAFE_DRUM_RPM;
        
        // Display RPM: drum mode shows drum RPM, spark mode shows engine RPM
        float rpm_for_hp;
        if (RPM_SOURCE == "drum") {
            display_engine_rpm = effective_drum_rpm;
            rpm_for_hp = effective_drum_rpm;
        } else if (engine_rpm > 60.0f) {
            display_engine_rpm = engine_rpm;
            rpm_for_hp = engine_rpm;
        } else {
            display_engine_rpm = effective_drum_rpm * current_gear_ratio;
            rpm_for_hp = effective_drum_rpm * current_gear_ratio;
        }
        
        float angular_accel = 0.0f;
        
        rpm_history[rpm_hist_idx] = effective_drum_rpm;
        rpm_hist_idx = (rpm_hist_idx + 1) % 5;
        
        if (now - last_smooth_time >= 100) {
            float dt = (now - last_smooth_time) / 1000.0f;
            
            if (dt > 0.05f && dt < 0.5f) {
                float rpm_old = 0.0f, rpm_new = 0.0f;
                int old_count = 0, new_count = 0;
                
                for (int i = 0; i < 5; i++) {
                    if (i < 2) { rpm_old += rpm_history[i]; old_count++; }
                    if (i >= 3) { rpm_new += rpm_history[i]; new_count++; }
                }
                
                if (old_count > 0 && new_count > 0) {
                    rpm_old /= old_count;
                    rpm_new /= new_count;
                    
                    float rpm_delta = rpm_new - rpm_old;
                    float raw_accel = (rpm_delta / dt) * (2.0f * 3.14159f / 60.0f);
                    
                    float alpha_accel = 0.12f;
                    smooth_accel = alpha_accel * raw_accel + (1.0f - alpha_accel) * smooth_accel;
                    angular_accel = smooth_accel;
                }
            }
            last_smooth_time = now;
        }
        
        if (angular_accel < -0.5f && effective_drum_rpm > 500) {
            if (coast_start_time == 0) {
                coast_start_time = now;
                coast_start_rpm = effective_drum_rpm;
            }
            
            if (now - coast_start_time > 500) {
                drag_torque_g = drum_inertia * fabs(angular_accel);
                drag_torque_g = constrain(drag_torque_g, 0.0f, 100.0f);
            }
        } else {
            coast_start_time = 0;
            drag_torque_g *= 0.995f;
        }
        current.drag_torque = drag_torque_g;

        drum_torque = 0.0f;
        engine_torque = 0.0f;
        hp = 0.0f;
        
        if (angular_accel > 0.05f) {
            drum_torque = drum_inertia * angular_accel;
            drum_torque = constrain(drum_torque, 0.0f, 2000.0f);
            
            engine_torque = drum_torque * current_gear_ratio;
            engine_torque = constrain(engine_torque, 0.0f, 2000.0f);
            
            if (rpm_for_hp > 0) {
                hp = (engine_torque * rpm_for_hp) / 7121.0f;
            }
            
            last_known_torque = engine_torque;
        } else if (angular_accel < -0.3f && effective_drum_rpm > 500) {
            engine_torque = drag_torque_g * current_gear_ratio;
            engine_torque = constrain(engine_torque, 0.0f, 200.0f);
            hp = 0.0f;
        } else {
            engine_torque = last_known_torque * 0.98f;
            if (rpm_for_hp > 0) {
                hp = (engine_torque * rpm_for_hp) / 7121.0f;
            }
            if (hp < 0.02f) hp = 0.0f;
        }
        
        display_drag_g = display_drag_g * 0.95f + drag_torque_g * 0.05f;
        
        current.gear_ratio = current_gear_ratio;
        last_drum_rpm_g = effective_drum_rpm;
        drum_rpm = effective_drum_rpm;
        
        if (hp > MAX_HP * 5.0f) hp = 0.0f;
        if (engine_torque > MAX_TORQUE * 5.0f) engine_torque = 0.0f;
      }

      // Air density correction
      if (DYNO_MODE == "inertia") {
        float correction_factor = constrain(STD_AIR_DENSITY / air_density, 0.85f, 1.15f);
        hp *= correction_factor;
        engine_torque *= correction_factor;
        drum_torque *= correction_factor;
      }

      // Smoothing
      float hp_smooth = hpSmooth.add(hp);
      float torque_smooth = torqueSmooth.add(engine_torque);

      if (hp_smooth > MAX_HP * 2.0f) hp_smooth = MAX_HP * 2.0f;
      if (torque_smooth > MAX_TORQUE * 2.0f) torque_smooth = MAX_TORQUE * 2.0f;

      // EGT
      if (mcp_egt_ok && (now - last_egt_read >= EGT_UPDATE_MS || egt_first_read)) {
        float hot = mcp_egt.readThermocouple();
        float amb = mcp_egt.readAmbient();
        if (egt_first_read) {
          vTaskDelay(pdMS_TO_TICKS(10));
          hot = mcp_egt.readThermocouple();
          amb = mcp_egt.readAmbient();
          egt_first_read = false;
        }
        if (!isnan(hot) && !isnan(amb) && hot > -50 && hot < 1500) {
          egtTracker.update(hot, amb);
        }
        last_egt_read = now;
      }

      // Populate LiveData
      current.engine_rpm = display_engine_rpm;
      current.drum_rpm = drum_rpm;
      current.hp = hp_smooth;
      current.engine_torque = torque_smooth;
      current.drum_torque = drum_torque;
      current.peakHP = runPeakHP;
      current.peakHP_RPM = runPeakHP_RPM;
      current.peakTorque = runPeakTorque;
      current.peakTorque_RPM = runPeakTorque_RPM;

      // Auto-run & Recording logic
      float rec_rpm = (RPM_SOURCE == "drum") ? drum_rpm : display_engine_rpm;

      if (autoMode && !recording) {
        if (!crossedStart && rec_rpm >= autoStartRPM + AUTO_HYSTERESIS) crossedStart = true;
        if (crossedStart && rec_rpm >= autoStartRPM && rec_rpm < autoEndRPM) {
          recording = true;
          currentRun.clear();
          runPeakHP = 0; runPeakHP_RPM = 0;
          runPeakTorque = 0; runPeakTorque_RPM = 0;
          egtTracker.reset();
          hpSmooth.reset();
          torqueSmooth.reset();
          recordingStartedMs = millis();
          ws.textAll("{\"type\":\"info\",\"msg\":\"Auto Run started\"}");
          crossedStart = false;
        }
      }

      if (autoMode && recording && rec_rpm >= autoEndRPM) {
        recording = false; autoMode = false; crossedStart = false; recordingStartedMs = 0;
        sendRunComplete(runPeakHP, runPeakHP_RPM, runPeakTorque, runPeakTorque_RPM);
      }

      if (recording && (millis() - recordingStartedMs > 45000UL)) {
        recording = false; autoMode = false; crossedStart = false;
        ws.textAll("{\"type\":\"info\",\"msg\":\"Safety timeout\"}");
        sendRunComplete(runPeakHP, runPeakHP_RPM, runPeakTorque, runPeakTorque_RPM);
        recordingStartedMs = 0;
      }

      if (recording) {
        if (hp_smooth > runPeakHP) {
          runPeakHP = hp_smooth; runPeakHP_RPM = rec_rpm;
        }
        if (torque_smooth > runPeakTorque) {
          runPeakTorque = torque_smooth; runPeakTorque_RPM = rec_rpm;
        }
        
        static unsigned long lastRecordTime = 0;
        if (now - lastRecordTime >= 100) {
          if (currentRun.size() >= MAX_RUN_POINTS) currentRun.erase(currentRun.begin());
          currentRun.push_back({rec_rpm, hp_smooth, torque_smooth});
          lastRecordTime = now;
        }
      }

      // WebSocket update
      if (now - last_ws_update >= WS_UPDATE_MS) {
        notifyClients(current, recording);
        last_ws_update = now;
      }

      processGpioTrigger();
    }

    vTaskDelay(pdMS_TO_TICKS(SENSOR_UPDATE_MS));
  }
}

//-------------------Setup---------------------
void setup() {
  Serial.begin(115200);
  while (!Serial && millis() < 3000) delay(10);

  Serial.println("\nKARLSDIYDYNO " + String(FIRMWARE_VERSION));

  // === WIFI FIRST ===
  WiFi.setSleep(false);
  WiFi.setTxPower(WIFI_POWER_15dBm);

  preferences.begin("dyno", true);
  sta_ssid = preferences.getString("sta_ssid", "");
  sta_password = preferences.getString("sta_password", "");
  preferences.end();

  apModeOnly = false;

  if (sta_ssid.length() > 3 && sta_password.length() > 3) {
    WiFi.mode(WIFI_STA);
    esp_wifi_set_max_tx_power(40);
    delay(100); 
    WiFi.begin(sta_ssid.c_str(), sta_password.c_str());

    Serial.print("Connecting to WiFi");
    unsigned long start = millis();
    while (WiFi.status() != WL_CONNECTED && millis() - start < 20000) {
      delay(500);
      Serial.print(".");
    }

    if (WiFi.status() == WL_CONNECTED) {
      Serial.println("\nSTA OK");
      Serial.print("IP: ");
      Serial.println(WiFi.localIP());
    } else {
      Serial.println("\nSTA failed, using AP mode");
      apModeOnly = true;
    }
  } else {
    Serial.println("No WiFi credentials, using AP mode");
    apModeOnly = true;
  }

  if (apModeOnly) {
    WiFi.mode(WIFI_AP);
    WiFi.softAP(ap_ssid, ap_password);
    Serial.print("AP Mode - IP: ");
    Serial.println(WiFi.softAPIP());
  }

  // === I2C & SENSORS ===
  Wire.begin(I2C_SDA, I2C_SCL);
  Wire.setClock(100000);

  if (!LittleFS.begin(true)) LittleFS.format();

  // BMP280
  bmp_initialized = bmp.begin(0x76) || bmp.begin(0x77);
  if (bmp_initialized) {
    bmp.setSampling(Adafruit_BMP280::MODE_NORMAL, Adafruit_BMP280::SAMPLING_X2,
                    Adafruit_BMP280::SAMPLING_X16, Adafruit_BMP280::FILTER_X16,
                    Adafruit_BMP280::STANDBY_MS_500);
  }

  // MCP9600 EGT
  mcp_egt_ok = mcp_egt.begin(MCP_EGT_ADDR);
  if (mcp_egt_ok) {
    mcp_egt.setThermocoupleType(MCP9600_TYPE_K);
    mcp_egt.setFilterCoefficient(3);
    mcp_egt.setADCresolution(MCP9600_ADCRESOLUTION_14);
  }

  // NAU7802
  nau_initialized = nau.begin();
  if (nau_initialized) {
    nau.setGain(NAU7802_GAIN_128);
    nau.setRate(NAU7802_RATE_320SPS);
    delay(2000); 
  }

  // Auto-tare at boot
  long sum = 0;
  int count = 0;
  for (int i = 0; i < 16; i++) {
    if (nau_initialized && nau.available()) {
      sum += nau.read();
      count++;
      delay(5);
    }
  }
  if (count > 0) {
    zeroOffset = sum / count;  
    current.load_kg = 0.0f;
  }

  // Load preferences
  preferences.begin("dyno", false);
  CALIBRATION_FACTOR     = preferences.getFloat("cal_factor", 0.000890);
  LEVER_ARM_METERS       = preferences.getFloat("lever_arm", 0.25);
  PULSES_PER_REVOLUTION  = preferences.getFloat("pulses_per_rev", 2.0f);
  BRAKE_TO_ENGINE_RATIO  = preferences.getFloat("brake_eng_ratio", 10.0f);
  RPM_FILTER_ALPHA       = preferences.getFloat("rpm_filter", 0.25f);
  AUTO_HYSTERESIS        = preferences.getFloat("auto_hyst", 180.0f);
  DYNO_MODE              = preferences.getString("dyno_mode", "brake");
  DRUM_MASS              = preferences.getFloat("drum_mass", 10.0f);
  DRUM_DIAM_METERS       = preferences.getFloat("drum_diam", 0.3f);
  PULSES_PER_REV_DRUM    = preferences.getFloat("pulses_drum", 10.0f);
  DRUM_INERTIA_TYPE      = preferences.getString("drum_inertia_type", "solid");
  DRUM_WALL_THICKNESS_CM = preferences.getFloat("drum_wall_cm", 2.0f);
  DRUM_INERTIA_CUSTOM    = preferences.getFloat("drum_inertia_custom", 0.5f);
  RPM_SOURCE             = preferences.getString("rpm_source", "spark");
  DRIVETRAIN_EFFICIENCY  = preferences.getFloat("drivetrain_eff", 0.93f);
  MAX_SAFE_DRUM_RPM      = preferences.getFloat("max_safe_drum_rpm", 3000.0f);
  MAX_RPM                = preferences.getFloat("max_rpm", 16000);
  MAX_HP                 = preferences.getFloat("max_hp", 25);
  MAX_TORQUE             = preferences.getFloat("max_torque", 25);
  GEAR_RATIO_MANUAL      = preferences.getFloat("gear_ratio_manual", 10.0f);
  GEAR_RATIO_AUTO        = preferences.getBool("gear_ratio_auto", true);
  GPIO_TRIGGER_PIN       = preferences.getInt("gpio_trigger_pin", 17);
  ENGINE_DEBOUNCE_US     = preferences.getInt("eng_debounce", 1400);
  DRUM_DEBOUNCE_US       = preferences.getInt("drum_debounce", 220);
  MEDIAN_SIZE            = preferences.getInt("median_size", 9);

  if (MEDIAN_SIZE < 3) MEDIAN_SIZE = 3;
  if (MEDIAN_SIZE > 9) MEDIAN_SIZE = 9;
  if (MEDIAN_SIZE % 2 == 0) MEDIAN_SIZE = 5;

  rpmMedian.resize(MEDIAN_SIZE);
  drumMedian.resize(MEDIAN_SIZE);
  std::fill(rpmMedian.begin(), rpmMedian.end(), 0.0f);
  std::fill(drumMedian.begin(), drumMedian.end(), 0.0f);

  preferences.end();

  // === WEB SERVER ===
  if (apModeOnly) {
    dnsServer.start(53, "*", WiFi.softAPIP());
  }
  server.addHandler(new CaptiveRequestHandler());

  server.on("/favicon.ico", HTTP_GET, [](AsyncWebServerRequest *request) { request->send(200, "text/plain", ""); });
  server.on("/version", HTTP_GET, [](AsyncWebServerRequest *request) { request->send(200, "text/plain", "KARLSDIYDYNO " + String(FIRMWARE_VERSION)); });
  server.serveStatic("/", LittleFS, "/").setDefaultFile("index.html");
  
  server.on("/settings", HTTP_GET, [](AsyncWebServerRequest *request) {
    if (LittleFS.exists("/settings.html")) request->send(LittleFS, "/settings.html", "text/html");
    else request->send(404);
  });

  server.on("/trend.html", HTTP_GET, [](AsyncWebServerRequest *request) {
    if (LittleFS.exists("/trend.html")) request->send(LittleFS, "/trend.html", "text/html");
    else request->send(404);
  });

  ws.onEvent(onWsEvent);
  server.addHandler(&ws);

  // OTA endpoints
  server.on("/update", HTTP_GET, [](AsyncWebServerRequest *request) {
    if (LittleFS.exists("/update.html")) request->send(LittleFS, "/update.html", "text/html");
    else request->send(404);
  });

  server.on("/update", HTTP_POST, [](AsyncWebServerRequest *request) {
    String msg = Update.hasError() ? "FAIL" : "OK";
    AsyncWebServerResponse *response = request->beginResponse(200, "text/plain", msg);
    response->addHeader("Connection", "close");
    request->send(response);
    if (!Update.hasError()) {
      delay(1000);
      ESP.restart();
    }
  }, [](AsyncWebServerRequest *request, String filename, size_t index, uint8_t *data, size_t len, bool final) {
    if (!index && !Update.begin(UPDATE_SIZE_UNKNOWN)) Update.printError(Serial);
    if (Update.write(data, len) != len) Update.printError(Serial);
    if (final && Update.end(true)) Serial.println("OTA OK");
  });

  server.on("/update_fs", HTTP_GET, [](AsyncWebServerRequest *request) {
    if (LittleFS.exists("/update_fs.html")) request->send(LittleFS, "/update_fs.html", "text/html");
    else request->send(404);
  });

  server.on("/api/sensors", HTTP_GET, [](AsyncWebServerRequest *request) {
    request->send(200, "application/json", getSensorStatusJSON());
  });

  server.on("/update_fs", HTTP_POST, [](AsyncWebServerRequest *request) {
    request->send(200, "text/plain", Update.hasError() ? "FAIL" : "OK");
  }, [](AsyncWebServerRequest *request, String filename, size_t index, uint8_t *data, size_t len, bool final) {
    static bool started = false;
    if (!index) started = Update.begin(UPDATE_SIZE_UNKNOWN, U_SPIFFS);
    if (started && len > 0) Update.write(data, len);
    if (final && started) { Update.end(true); started = false; }
  });

  server.begin();

  // === INTERRUPTS ===
  pinMode(HALL_PIN, INPUT_PULLUP);
  attachInterrupt(digitalPinToInterrupt(HALL_PIN), hallISR, FALLING);
  pinMode(DRUM_PIN, INPUT_PULLUP);
  attachInterrupt(digitalPinToInterrupt(DRUM_PIN), drumISR, FALLING);

  // GPIO trigger
  bool pinOk = (GPIO_TRIGGER_PIN >= 0 && GPIO_TRIGGER_PIN <= 39);
  if (pinOk) {
    for (int i = 0; i < RESERVED_COUNT; i++) if (GPIO_TRIGGER_PIN == RESERVED_PINS[i]) pinOk = false;
  }
  if (pinOk && GPIO_TRIGGER_PIN != HALL_PIN && GPIO_TRIGGER_PIN != DRUM_PIN && GPIO_TRIGGER_PIN != I2C_SDA && GPIO_TRIGGER_PIN != I2C_SCL) {
    pinMode(GPIO_TRIGGER_PIN, INPUT_PULLUP);
    attachInterrupt(digitalPinToInterrupt(GPIO_TRIGGER_PIN), gpioTriggerISR, FALLING);
  } else {
    GPIO_TRIGGER_PIN = -1;
  }

  xTaskCreatePinnedToCore(calibrationTask, "CalibTask", 8192, NULL, 1, &calibTask, 1);
  xTaskCreatePinnedToCore(sensorTask, "SensorTask", 12288, NULL, 2, NULL, 1); 

  Serial.println("READY - http://" + (apModeOnly ? WiFi.softAPIP() : WiFi.localIP()).toString());
  
  Serial.printf("Free heap: %d bytes\n", ESP.getFreeHeap());
  Serial.printf("Min free heap: %d bytes\n", ESP.getMinFreeHeap());
}

// ─── LOOP ────────────────────────────────────────────────────────
void loop() {
  static unsigned long lastWiFiCheck = 0;
  if (!apModeOnly && WiFi.status() != WL_CONNECTED) {
    if (millis() - lastWiFiCheck > 10000) {
      lastWiFiCheck = millis();
      Serial.println("WiFi disconnected, reconnecting...");
      WiFi.reconnect();
    }
  }
  
  if (apModeOnly) {
    dnsServer.processNextRequest();
  }
  
  ws.cleanupClients();
  vTaskDelay(pdMS_TO_TICKS(1));
}