#include <Arduino.h>
#include <WiFi.h>
#include <ESPmDNS.h>
#include <esp_task_wdt.h>
#include <ArduinoOTA.h>
#include "soc/soc.h"
#include "soc/rtc_cntl_reg.h"

#include <atomic>

#include "config.h"
#include "attenuator.h"
#include "web_server.h"
#include "display.h"
#include "app.h"
#include "usb_serial.h"

// --- Global instances ---
Attenuator attenuator;

// --- Display update tracking ---
static uint32_t lastDisplayUpdate = 0;
static const uint32_t DISPLAY_UPDATE_MS = 100;  // Update every 100ms

// --- LED State Machine ---
enum class LEDState {
    Off,           // Headless mode (no WiFi)
    SlowBlink,     // Connecting to WiFi
    Solid,         // Connected
    FastBlink      // Sweeping
};

static LEDState ledState = LEDState::Off;
static uint32_t lastLedToggle = 0;
static bool ledOn = false;

void setLEDState(LEDState state) {
    ledState = state;
    if (state == LEDState::Off) {
        digitalWrite(PIN_LED, LOW);
        ledOn = false;
    } else if (state == LEDState::Solid) {
        digitalWrite(PIN_LED, HIGH);
        ledOn = true;
    }
}

void updateLED() {
    uint32_t now = millis();
    uint32_t interval = 0;

    switch (ledState) {
        case LEDState::SlowBlink:
            interval = 500;
            break;
        case LEDState::FastBlink:
            interval = 100;
            break;
        default:
            return;  // Off or Solid -- nothing to toggle
    }

    if (now - lastLedToggle >= interval) {
        lastLedToggle = now;
        ledOn = !ledOn;
        digitalWrite(PIN_LED, ledOn ? HIGH : LOW);
    }
}

// --- Sweep Mode ---
// std::atomic for cross-core visibility (async_tcp task vs Arduino loop task)
static std::atomic<bool> sweepRunning{false};
static std::atomic<int8_t> sweepDirection{1};     // 1 = up, -1 = down
static std::atomic<uint32_t> sweepDwellMs{SWEEP_DWELL_MS_DEFAULT};
static uint32_t lastSweepStep = 0;  // Only accessed from loop() task

void startSweep(bool up, uint32_t dwellMs) {
    sweepDirection.store(up ? 1 : -1, std::memory_order_release);
    sweepDwellMs.store(constrain(dwellMs, SWEEP_DWELL_MS_MIN, SWEEP_DWELL_MS_MAX),
                       std::memory_order_release);
    lastSweepStep = millis();
    sweepRunning.store(true, std::memory_order_release);
    setLEDState(LEDState::FastBlink);
    Serial0.printf("[Sweep] Started, direction=%s, dwell=%u ms\n",
                  up ? "up" : "down", sweepDwellMs.load());
}

void stopSweep() {
    sweepRunning.store(false, std::memory_order_release);
    // Persist final position to NVS (skipped during sweep to avoid flash wear)
    attenuator.persistCurrent();
    setLEDState(LEDState::Solid);
    Serial0.println("[Sweep] Stopped");
}

bool isSweeping() {
    return sweepRunning.load(std::memory_order_acquire);
}

int8_t getSweepDirection() {
    return sweepDirection.load(std::memory_order_acquire);
}

uint32_t getSweepDwellMs() {
    return sweepDwellMs.load(std::memory_order_acquire);
}

void updateSweep() {
    if (!sweepRunning.load(std::memory_order_acquire)) return;

    uint32_t now = millis();
    uint32_t dwell = sweepDwellMs.load(std::memory_order_acquire);
    if (now - lastSweepStep < dwell) return;
    lastSweepStep = now;

    // Atomic read-modify-write: no TOCTOU race
    attenuator.advanceStep(sweepDirection.load(std::memory_order_acquire));
}

// --- WiFi TX Power Control ---
static wifi_power_t wifiTxPower = WIFI_TX_POWER_DBM;

// Valid wifi_power_t levels (quarter-dBm units)
static const int VALID_WIFI_POWERS[] = {8, 20, 28, 34, 44, 52, 60, 68, 74, 76, 78};
static const float VALID_WIFI_DBMS[] = {2.0, 5.0, 7.0, 8.5, 11.0, 13.0, 15.0, 17.0, 18.5, 19.0, 19.5};
static const int NUM_WIFI_POWER_LEVELS = 11;

bool isValidWifiPower(int raw) {
    for (int i = 0; i < NUM_WIFI_POWER_LEVELS; i++) {
        if (VALID_WIFI_POWERS[i] == raw) return true;
    }
    return false;
}

void setWiFiTxPower(wifi_power_t power) {
    wifiTxPower = power;
    WiFi.setTxPower(power);
    Serial0.printf("[WiFi] TX power set to %d (quarter dBm units)\n", power);
}

wifi_power_t getWiFiTxPower() {
    return wifiTxPower;
}

// Convert wifi_power_t enum to approximate dBm float
float wifiPowerToDbm(wifi_power_t power) {
    // wifi_power_t values are in quarter-dBm units
    return static_cast<int>(power) / 4.0f;
}

const int* getValidWifiPowers() { return VALID_WIFI_POWERS; }
const float* getValidWifiDbms() { return VALID_WIFI_DBMS; }
int getNumWifiPowerLevels() { return NUM_WIFI_POWER_LEVELS; }

// --- WiFi Connection ---
bool connectWiFi() {
    Serial0.printf("[WiFi] Connecting to %s...\n", WIFI_SSID);
    // Keep LED off during connection to reduce power draw (prevents brownout)
    setLEDState(LEDState::Off);

    WiFi.mode(WIFI_STA);
    WiFi.setHostname(FW_HOSTNAME);
    WiFi.setTxPower(wifiTxPower);
    Serial0.printf("[WiFi] TX power: %.1f dBm\n", wifiPowerToDbm(wifiTxPower));
    WiFi.begin(WIFI_SSID, WIFI_PASS);

    uint32_t startAttempt = millis();
    while (WiFi.status() != WL_CONNECTED) {
        updateLED();
        delay(100);
        esp_task_wdt_reset();

        if (millis() - startAttempt > WIFI_TIMEOUT_MS) {
            Serial0.println("[WiFi] Connection timeout, continuing without network");
            setLEDState(LEDState::Off);
            return false;
        }
    }

    Serial0.printf("[WiFi] Connected! IP: %s\n", WiFi.localIP().toString().c_str());
    setLEDState(LEDState::Solid);
    return true;
}

// --- mDNS Setup ---
void setupMDNS() {
    if (MDNS.begin(FW_HOSTNAME)) {
        MDNS.addService("http", "tcp", WEB_PORT);
        Serial0.printf("[mDNS] Registered as %s.local\n", FW_HOSTNAME);
    } else {
        Serial0.println("[mDNS] Failed to start");
    }
}

// --- OTA Setup ---
static bool otaEnabled = false;

void enableOTA() {
    if (otaEnabled) return;

    ArduinoOTA.setHostname(FW_HOSTNAME);
#ifdef OTA_PASSWORD
    ArduinoOTA.setPassword(OTA_PASSWORD);
    Serial0.println("[OTA] Password authentication enabled");
#else
    Serial0.println("[OTA] WARNING: No password set (define OTA_PASSWORD in platformio_local.ini)");
#endif
    ArduinoOTA.onStart([]() {
        stopSweep();
        setLEDState(LEDState::FastBlink);
        Serial0.println("[OTA] Update starting...");
    });
    ArduinoOTA.onEnd([]() {
        Serial0.println("[OTA] Update complete, rebooting...");
    });
    ArduinoOTA.onProgress([](unsigned int progress, unsigned int total) {
        esp_task_wdt_reset();  // Keep watchdog happy during long OTA
        Serial0.printf("[OTA] Progress: %u%%\r", (progress / (total / 100)));
    });
    ArduinoOTA.onError([](ota_error_t error) {
        Serial0.printf("[OTA] Error %u: ", error);
        if (error == OTA_AUTH_ERROR) Serial0.println("Auth failed");
        else if (error == OTA_BEGIN_ERROR) Serial0.println("Begin failed");
        else if (error == OTA_CONNECT_ERROR) Serial0.println("Connect failed");
        else if (error == OTA_RECEIVE_ERROR) Serial0.println("Receive failed");
        else if (error == OTA_END_ERROR) Serial0.println("End failed");
        setLEDState(LEDState::Solid);
    });
    ArduinoOTA.begin();
    otaEnabled = true;
    Serial0.println("[OTA] Enabled");
}

bool isOTAEnabled() {
    return otaEnabled;
}

// --- Setup ---
void setup() {
    // Disable brownout detector - USB power can sag during WiFi TX
    // Trade-off: if supply drops below 3.0V, MCU may execute with corrupted SRAM
    // rather than cleanly resetting. Acceptable for bench-powered USB device,
    // NOT acceptable for battery or unstable power source.
    WRITE_PERI_REG(RTC_CNTL_BROWN_OUT_REG, 0);

    // UART0 on ESP32-S3-DevKitC-1 (CH343 bridge): TX=GPIO43, RX=GPIO44
    // These are the default pins for Serial0, no need to specify
    Serial0.begin(115200);
    delay(100);  // Brief delay for UART to stabilize

    Serial0.println();
    Serial0.println("================================");
    Serial0.printf("HMC472A Attenuator Controller %s\n", FW_VERSION);
    Serial0.println("================================");

    // Initialize LED
    pinMode(PIN_LED, OUTPUT);
    setLEDState(LEDState::Off);

    // Initialize OLED display
    if (initDisplay()) {
        showSplash(FW_VERSION);
        delay(1500);  // Show splash briefly
    }

    // Initialize watchdog (ESP-IDF v5.x API)
    esp_task_wdt_config_t wdt_config = {
        .timeout_ms = WDT_TIMEOUT_S * 1000,
        .idle_core_mask = (1 << portNUM_PROCESSORS) - 1,  // All cores
        .trigger_panic = true
    };
    esp_task_wdt_init(&wdt_config);
    esp_task_wdt_add(NULL);
    Serial0.printf("[WDT] Initialized, timeout=%d s\n", WDT_TIMEOUT_S);

    // Initialize attenuator (restores from NVS)
    attenuator.begin();

    // USB CDC serial command interface (native USB OTG port)
    setupUSBSerial(attenuator);

    // Power stabilization delay before WiFi
    // USB power supplies need time to stabilize under load
    Serial0.println("[Power] Waiting for supply to stabilize...");
    delay(1000);

    // Connect to WiFi
    bool wifiConnected = connectWiFi();

    if (wifiConnected) {
        // Start mDNS
        setupMDNS();

        // Start web server
        setupWebServer(attenuator);
    }

    Serial0.println("[Setup] Complete");
    Serial0.println();
}

// --- Main Loop ---
void loop() {
    esp_task_wdt_reset();
    updateLED();
    updateSweep();
    handleUSBSerial();

    if (otaEnabled) {
        ArduinoOTA.handle();
    }

    // Update display periodically
    uint32_t now = millis();
    if (now - lastDisplayUpdate >= DISPLAY_UPDATE_MS) {
        lastDisplayUpdate = now;
        AttenuatorState state = attenuator.getSnapshot();
        updateDisplay(
            state.db,
            state.step,
            WiFi.RSSI(),
            isSweeping(),
            WiFi.status() == WL_CONNECTED
        );
    }

    delay(1);  // Yield to other tasks
}