#include "AppUpgrade.h"
#include "esp_log.h"
#include <MD5Builder.h>
#include <LittleFS.h>
#include <memory>
#include <algorithm>
#include "global.h"
#include "JsonConstrain.h"
#include "BLE_UpdateService.h"
#include <HTTPClient.h>
#include <Update.h>
#include <cstring>
#include <esp_task_wdt.h>
#include <esp_ota_ops.h>

static const char* TAG = "AppUpdater";
TaskHandle_t Update_Task_Handle = NULL;
TaskHandle_t versionCheckTask_Handle = NULL;
volatile bool g_UpdateCancelFlag = false; // cancellation flag
UpdateMode g_UpdateMode = UpdateMode::UPDATE_BOTH; // Default to updating both files and firmware
String updateUrl = "";
Version otaVersion;


AppUpdater::AppUpdater(fs::FS& fs, Version localVersion, const char* bucket, const char* manifestName, const char* appBin)
    : localVersion(localVersion), manifestName(manifestName), appName(appBin), fileSystem(fs)
{
    // Use dynamic buffer size based on available memory - much more conservative now
    size_t available_heap = ESP.getFreeHeap();
    size_t buffer_size = std::min<size_t>(BUFFER_SIZE, available_heap / 16); // Use at most 1/16 of free heap
    // Ensure minimum viable size
    if (buffer_size < 1024) buffer_size = 1024;  // Absolute minimum is 1KB
    
    downloadBuffer.reset(new uint8_t[buffer_size]);
    
    baseUrl = bucket ? String(bucket) : String(DEFAULT_MANIFEST_URL);
    // Ensure baseUrl ends with a single '/'
    if(!baseUrl.endsWith("/")) baseUrl += "/";
    ESP_LOGI(TAG, "AppUpdater initialized (local v%s) baseUrl=%s, buffer=%u bytes", 
             localVersion.toString().c_str(), baseUrl.c_str(), buffer_size);
}

void AppUpdater::setProgressCallback(void (*callback)( UpdateStatus status, int percentage, const char* message)) {
    progressCb = callback;
}

void AppUpdater::setUpdateMode(UpdateMode mode) {
    updateMode = mode;
    g_UpdateMode = mode; // Sync with global mode for firmware update task
}

UpdateMode AppUpdater::getUpdateMode() const {
    return updateMode;
}

void AppUpdater::updateProgress(UpdateStatus newStatus, int percentage, const char* message) {
    status = newStatus;
    if (progressCb) {
        progressCb(status, percentage, message);
    }
}

AppUpdater::ManifestCheckResult AppUpdater::checkManifest() {
    String url = buildUrl(manifestName);
    ESP_LOGD(TAG, "Fetching manifest from: %s", url.c_str());

    String payload;
    for(int attempt=0; attempt<HTTP_RETRY_COUNT; ++attempt){
        if(g_UpdateCancelFlag) return ManifestCheckResult::ERROR_FETCH_FAILED;
        HTTPClient http;
        http.begin(url);
        http.setTimeout(10);  // 10 second timeout
        http.setConnectTimeout(10000);  // 10 second connect timeout
        int httpCode = http.GET();
        if (httpCode == HTTP_CODE_OK) {
            payload = http.getString();
            http.end();
            break;
        }
        ESP_LOGW(TAG, "Manifest GET failed (attempt %d/%d): %d", attempt+1, HTTP_RETRY_COUNT, httpCode);
        http.end();
        if(attempt+1 < HTTP_RETRY_COUNT) vTaskDelay(pdMS_TO_TICKS(HTTP_RETRY_DELAY_MS));
    }
    if(payload.isEmpty()){
        ESP_LOGE(TAG, "Failed to fetch manifest after retries");
        return ManifestCheckResult::ERROR_FETCH_FAILED;
    }

    if(payload.length() > MAX_MANIFEST_SIZE){
        ESP_LOGE(TAG, "Manifest too large (%u bytes)", (unsigned)payload.length());
        return ManifestCheckResult::ERROR_TOO_LARGE;
    }

    // Parse JSON
    DeserializationError error = deserializeJson(jsonManifest, payload);
    ESP_LOGD(TAG, "Manifest deserialized");
    if (error) {
        ESP_LOGE(TAG, "Failed to parse manifest: %s", error.c_str());
        return ManifestCheckResult::ERROR_PARSE_FAILED;
    }

    // Check for files section
    jsonFilesArray = jsonManifest["files"];
    if (jsonFilesArray.isNull()) {
        ESP_LOGE(TAG, "No files section in manifest");
        return ManifestCheckResult::ERROR_NO_FILES_SECTION;
    }else{
        ESP_LOGD(TAG, "%d Files found", jsonFilesArray.size());
    }

    // Check for version section
    JsonObject jsonVersion = jsonManifest["version"];
    ESP_LOGD(TAG, "Version section found");
    if (jsonVersion.isNull()) {
        ESP_LOGE(TAG, "No version section in manifest");
        return ManifestCheckResult::ERROR_NO_VERSION;
    }

    // Get the remote version
    byte major = jsonVersion["major"] | 0;
    byte minor = jsonVersion["minor"] | 0;
    byte patch = jsonVersion["patch"] | 0;
    otaVersion = {major, minor, patch};

    //Version localVersion;
    //::sscanf(localVersion, "%d.%d.%d", &localVersion.major, &localVersion.minor, &localVersion.patch);
    
    // Check if an update is available
    updateAvailable = false;
    // Only mark update available if remote is strictly newer than local
    if (otaVersion <= localVersion) {
        ESP_LOGI(TAG, "No updates available: remote=%s, local=%s", 
                 otaVersion.toString().c_str(), localVersion.toString().c_str());
        return ManifestCheckResult::VERSION_CURRENT;
    }else{
        updateAvailable = true;
        ESP_LOGI(TAG, "Update available: remote=%s, local=%s", 
                 otaVersion.toString().c_str(), localVersion.toString().c_str());
    }

    //ESP_LOGD(TAG, "Manifest content: %s", payload.c_str());

    return ManifestCheckResult::UPDATE_AVAILABLE;
}

bool AppUpdater::updateFile(const char* remotePath, const char* localPath, const char* expectedMd5) {
    //updateProgress(UpdateStatus::DOWNLOADING, 0, localPath);
    
    // Construct full URL
    String url = buildUrl(remotePath);
    ESP_LOGD(TAG, "Downloading: %s -> %s", url.c_str(), localPath);

    // Quick skip: if exists and size & MD5 match
    bool skip = false;
    if(fileSystem.exists(localPath)){
        String localMd5 = getLocalMD5(localPath);
        ESP_LOGI(TAG, "Local file exists: %s, MD5: %s, Expected: %s", localPath, localMd5.c_str(), expectedMd5);
        if(localMd5.equals(expectedMd5)) skip = true;
    } else {
        ESP_LOGI(TAG, "Local file does not exist: %s", localPath);
    }
    if(skip){
        ESP_LOGI(TAG, "File already up to date: %s", localPath);
        updateProgress(UpdateStatus::FILE_SKIPPED, 100, localPath);
        return true;
    }

    ESP_LOGI(TAG, "Need to download file: %s (local MD5 mismatch or file missing)", localPath);

    // Start the download
    HTTPClient http;
    int httpCode = -1;
    for(int attempt=0; attempt<HTTP_RETRY_COUNT; ++attempt){
        if(g_UpdateCancelFlag) return false;
        http.begin(url);
        httpCode = http.GET();
        if(httpCode == HTTP_CODE_OK) break;
        ESP_LOGW(TAG, "File GET failed (attempt %d/%d): %d", attempt+1, HTTP_RETRY_COUNT, httpCode);
        http.end();
        if(attempt+1 < HTTP_RETRY_COUNT) vTaskDelay(pdMS_TO_TICKS(HTTP_RETRY_DELAY_MS));
    }
    if (httpCode != HTTP_CODE_OK) {
        ESP_LOGE(TAG, "Download failed for %s: HTTP code %d", localPath, httpCode);
        updateProgress(UpdateStatus::ERROR, 0, String(String("Download failed: ") + localPath).c_str());
        http.end();
        return false;
    }

    // Get the stream and content length
    WiFiClient* stream = http.getStreamPtr();
    size_t contentLength = http.getSize();
    
    // Verify and save the file
    bool success = verifyAndSaveFile(stream, contentLength, localPath, remotePath, expectedMd5);
    http.end();
    if(!success){
        String errMsg = String(localPath) + " MD5 failed";
        updateProgress(UpdateStatus::ERROR, 0, errMsg.c_str());
        
        // For HTML/CSS/JS files, we might have saved them anyway (see verifyAndSaveFile logic)
        if (fileSystem.exists(localPath) && 
            (String(localPath).endsWith(".html") || 
             String(localPath).endsWith(".css") || 
             String(localPath).endsWith(".js"))) {
            ESP_LOGW(TAG, "Using %s despite MD5 mismatch (non-critical file)", localPath);
            // Return false to indicate verification failure but file is still usable
        }
    }else{
        updateProgress(UpdateStatus::FILE_SAVED, 100, localPath);
    }

    return success;
}

bool AppUpdater::verifyAndSaveFile(WiFiClient* stream, size_t contentLength, const char* localPath, const char* remotePath, const char* expectedMd5) 
{
    const int MAX_RETRIES = 2;  // Maximum number of retries for MD5 failure
    
    for (int retry = 0; retry <= MAX_RETRIES; retry++) {
        if (retry > 0) {
            ESP_LOGW(TAG, "Retrying download of %s (attempt %d/%d)", localPath, retry, MAX_RETRIES);
            // Need to re-fetch the file for retry - use the REMOTE path, not local path
            String url = buildUrl(remotePath);
            HTTPClient http;
            http.begin(url);
            int httpCode = http.GET();
            if (httpCode != HTTP_CODE_OK) {
                ESP_LOGE(TAG, "Retry download failed: %d", httpCode);
                http.end();
                continue;  // Try next retry if available
            }
            stream = http.getStreamPtr();
            contentLength = http.getSize();
        }
        
        MD5Builder md5;
        md5.begin();
        size_t totalRead = 0;
        
        // Create temporary filename in the same directory as the target file
        String targetDir = String(localPath);
        int lastSlash = targetDir.lastIndexOf('/');
        String tempPath;
        
        if (lastSlash >= 0) {
            // Extract directory and filename
            String directory = targetDir.substring(0, lastSlash + 1);
            String filename = targetDir.substring(lastSlash + 1);
            tempPath = directory + "temp_" + filename + ".download";
        } else {
            // File is in root directory
            tempPath = "/temp_" + String(localPath) + ".download";
        }
        
        // Clean up any existing temp file first
        if (fileSystem.exists(tempPath.c_str())) {
            ESP_LOGW(TAG, "Removing existing temp file: %s", tempPath.c_str());
            fileSystem.remove(tempPath.c_str());
        }
        
        ESP_LOGI(TAG, "Using temp file path: %s for target: %s", tempPath.c_str(), localPath);
        
        // Open temporary file for writing (LittleFS will create directories automatically)
        File file = fileSystem.open(tempPath.c_str(), FILE_WRITE, true); // true = create if not exists
        if (!file) {
            ESP_LOGE(TAG, "Failed to open temporary file for writing: %s", tempPath.c_str());
            
            // Try to diagnose the issue
            ESP_LOGE(TAG, "LittleFS info - Used: %u bytes, Total: %u bytes", 
                     LittleFS.usedBytes(), LittleFS.totalBytes());
            
            // Check if we're out of space
            if (LittleFS.usedBytes() >= LittleFS.totalBytes() * 0.95) {
                ESP_LOGE(TAG, "LittleFS nearly full - may not have space for temp file");
            }
            
            return false;
        }
    
        //updateProgress(UpdateStatus::DOWNLOADING, 0, localPath);
    
        if (contentLength > 0) {
            // Single pass with known content length
            while (totalRead < contentLength) {
                if(g_UpdateCancelFlag){ file.close(); fileSystem.remove(tempPath.c_str()); return false; }
                size_t available = stream->available();
                if (available) {
                    size_t readLen = stream->readBytes(downloadBuffer.get(), std::min(available, size_t(BUFFER_SIZE)));
                    
                    // Write to temp file and update MD5
                    if (file.write(downloadBuffer.get(), readLen) != readLen) {
                        ESP_LOGE(TAG, "Failed to write to temporary file");
    
                        file.close();
                        fileSystem.remove(tempPath.c_str());
                        return false;
                    }
                    
                    md5.add(downloadBuffer.get(), readLen);
                    totalRead += readLen;
                    updateProgress(UpdateStatus::DOWNLOADING, (totalRead * 80) / contentLength , localPath);
                }
                yield();
            }
        } else {
            // Unknown content length: read until stream ends
            for (;;) {
                if(g_UpdateCancelFlag){ file.close(); fileSystem.remove(tempPath.c_str()); return false; }
                size_t readLen = stream->readBytes(downloadBuffer.get(), BUFFER_SIZE);
                if (readLen == 0) {
                    break;
                }
                if (file.write(downloadBuffer.get(), readLen) != readLen) {
                    ESP_LOGE(TAG, "Failed to write to temporary file");
                    file.close();
                    fileSystem.remove(tempPath.c_str());
                    return false;
                }
                md5.add(downloadBuffer.get(), readLen);
                totalRead += readLen;
                // Progress unknown; emit periodic heartbeats at 0%
                // For unknown size, send heartbeats every ~16KB
                if((totalRead & 0x3FFF) == 0){
                    updateProgress(UpdateStatus::DOWNLOADING, 0, localPath);
                }
                yield();
            }
        }
    
        file.close();
        md5.calculate();
        String calculatedMd5 = md5.toString();
    
        // Verify MD5 hash
        updateProgress(UpdateStatus::VERIFYING, 90, localPath);
        
        ESP_LOGI(TAG, "MD5 verification for %s: Expected='%s', Calculated='%s'", 
                localPath, expectedMd5, calculatedMd5.c_str());
        
        // Compare MD5 case-insensitively (in case there are case differences)
        String expectedMd5Lower = String(expectedMd5);
        expectedMd5Lower.toLowerCase();
        String calculatedMd5Lower = calculatedMd5;
        calculatedMd5Lower.toLowerCase();
        
        if (!calculatedMd5Lower.equals(expectedMd5Lower)) {
            ESP_LOGE(TAG, "MD5 mismatch for %s (attempt %d/%d). Expected: %s, Got: %s", 
                    localPath, retry+1, MAX_RETRIES+1, expectedMd5, calculatedMd5.c_str());
            ESP_LOGE(TAG, "Length comparison - Expected: %d chars, Got: %d chars", 
                    strlen(expectedMd5), calculatedMd5.length());
            fileSystem.remove(tempPath.c_str());
            
            if (retry < MAX_RETRIES) {
                // Will retry in next loop iteration
                continue;
            }
            
            // Special case for certain file types - allow them to be used even with MD5 mismatch
            // This is a fallback option for non-critical files like HTML pages
            bool isNonCriticalFile = false;
            if (String(localPath).endsWith(".html") || 
                String(localPath).endsWith(".css") || 
                String(localPath).endsWith(".js")) {
                isNonCriticalFile = true;
            }
            
            if (isNonCriticalFile) {
                ESP_LOGW(TAG, "Using file %s despite MD5 mismatch (non-critical file)", localPath);
                // We'll still keep this file but report it as a verification failure
                // Ensure target directory exists before rename
                String dirPath = String(localPath);
                int targetLastSlash = dirPath.lastIndexOf('/');
                if (targetLastSlash > 0) {
                    dirPath = dirPath.substring(0, targetLastSlash);
                    if (!fileSystem.exists(dirPath.c_str())) {
                        ESP_LOGI(TAG, "Creating target directory: %s", dirPath.c_str());
                        String dummyFile = dirPath + "/.dummy";
                        File dummy = fileSystem.open(dummyFile.c_str(), FILE_WRITE, true);
                        if (dummy) {
                            dummy.print("temp");
                            dummy.close();
                            fileSystem.remove(dummyFile.c_str());
                        }
                    }
                }
                
                // Rename the temp file to the final location
                if (fileSystem.exists(localPath)) {
                    fileSystem.remove(localPath);
                }
                if (!fileSystem.rename(tempPath.c_str(), localPath)) {
                    ESP_LOGE(TAG, "Failed to rename temporary file for non-critical use: %s -> %s", 
                             tempPath.c_str(), localPath);
                    fileSystem.remove(tempPath.c_str());
                    return false;
                }
                // Return false to indicate verification failure, but the file will still be used
                return false;
            }
            
            return false;
        }
    
        updateProgress(UpdateStatus::VERIFYING, 95, localPath);
    
        // Ensure target directory exists before rename
        String dirPath = String(localPath);
        int targetLastSlash = dirPath.lastIndexOf('/');
        if (targetLastSlash > 0) {
            dirPath = dirPath.substring(0, targetLastSlash);
            if (!fileSystem.exists(dirPath.c_str())) {
                ESP_LOGI(TAG, "Creating target directory: %s", dirPath.c_str());
                String dummyFile = dirPath + "/.dummy";
                File dummy = fileSystem.open(dummyFile.c_str(), FILE_WRITE, true);
                if (dummy) {
                    dummy.print("temp");
                    dummy.close();
                    fileSystem.remove(dummyFile.c_str());
                }
            }
        }
        
        // Replace original file with verified temp file
        if (fileSystem.exists(localPath)) {
            fileSystem.remove(localPath);
        }
        if (!fileSystem.rename(tempPath.c_str(), localPath)) {
            ESP_LOGE(TAG, "Failed to rename temporary file: %s -> %s", tempPath.c_str(), localPath);
            fileSystem.remove(tempPath.c_str());
            return false;
        }
    
        updateProgress(UpdateStatus::VERIFYING, 100, localPath);
        return true;
    }
    
    return false;  // All retries failed
}

String AppUpdater::getLocalMD5(const char* filePath){
    File file = fileSystem.open(filePath, "r");
    if(!file){
        ESP_LOGE(TAG, "Error opening %s...", filePath);
        return String();
    }

    MD5Builder md5Builder;
    md5Builder.begin();
    size_t fileSize = file.size();
    size_t totalRead = 0;
    size_t readLen = 0;
    while (totalRead < fileSize) {
        readLen = file.readBytes(reinterpret_cast<char*>(downloadBuffer.get()), std::min(fileSize - totalRead, size_t(BUFFER_SIZE)));
        md5Builder.add(downloadBuffer.get(), readLen);
        totalRead += readLen;
    }

    md5Builder.calculate();
    file.close();
    return md5Builder.toString();
}

bool AppUpdater::updateFilesArray() {
    int successCount = 0;
    int totalFiles = jsonFilesArray.size();
    int failedCount = 0;
    ESP_LOGI(TAG, "Found %d files in manifest", totalFiles);

    // Iterate over each file entry in the manifest
    for (JsonObject file : jsonFilesArray) {
        const char* remotePath = file["path"];
        const char* localPath = remotePath;
        // If path begins with "data/" or "/data/" strip only the "data" portion, retaining the leading slash
        if (localPath) {
            if (strncmp(localPath, "data/", 5) == 0) {
                localPath += 4; // points to '/'
            } else if (strncmp(localPath, "/data/", 6) == 0) {
                localPath += 5; // points to '/'
            }
        }
        const char* expectedMd5 = file["md5"];

        // Skip invalid entries
        if (!remotePath || !localPath || !expectedMd5) {
            ESP_LOGE(TAG, "Invalid file entry in manifest");
            continue;
        }

        // Attempt to update the file
        if (updateFile(remotePath, localPath, expectedMd5)) {
            successCount++;
        } else {
            failedCount++;
            ESP_LOGE(TAG, "Failed to update file: %s", localPath);
            // Continue with remaining files instead of stopping
        }
    }

    ESP_LOGI(TAG, "Manifest update complete: %d/%d files updated, %d failed", 
             successCount, totalFiles, failedCount);
    
    // Consider the update successful if most files updated correctly
    // Allow up to 20% of files to fail before considering the entire update failed
    return (failedCount <= totalFiles * 0.2);
}

bool AppUpdater::updateApp() {
    updateProgress(UpdateStatus::MESSAGE, 0, "Starting firmware update");

    // Check for firmware section in manifest
    if (!jsonManifest["firmware"].is<JsonObject>() || !jsonManifest["firmware"]["md5"].is<const char*>()) {
        ESP_LOGE(TAG, "Invalid firmware section in manifest");
        updateProgress(UpdateStatus::ERROR, 0, "Firmware: Invalid firmware section in manifest");
        return false;
    }

    // Get the firmware MD5 hash and URL
    const char* expectedMd5 = jsonManifest["firmware"]["md5"];
    String firmwareUrl = buildUrl(appName);
    
    // First, try a HEAD request to verify server connectivity and file availability
    HTTPClient httpHead;
    bool fileExists = false;
    httpHead.begin(firmwareUrl);
    httpHead.setTimeout(10);
    int headCode = httpHead.sendRequest("HEAD");
    if (headCode == HTTP_CODE_OK) {
        ESP_LOGI(TAG, "Firmware file exists on server, size: %d bytes", httpHead.getSize());
        fileExists = true;
    } else {
        ESP_LOGW(TAG, "HEAD request failed: %d, proceeding with direct download attempt", headCode);
    }
    httpHead.end();
    
    // Download the firmware with progressive timeouts
    HTTPClient http;
    int httpCode = -1;
    for(int attempt=0; attempt<HTTP_RETRY_COUNT; ++attempt){
        if(g_UpdateCancelFlag) return false;
        
        // Progressive timeout increase with each attempt
        int timeout = 30 + (attempt * 15); // Start with 30s, add 15s each retry (30, 45, 60, 75, 90)
        int connectTimeout = 15000 + (attempt * 5000); // Start with 15s, add 5s each retry
        
        http.begin(firmwareUrl);
        http.setTimeout(timeout);
        http.setConnectTimeout(connectTimeout);
        
        ESP_LOGI(TAG, "Downloading firmware from %s (attempt %d/%d, timeout=%ds)", 
                 firmwareUrl.c_str(), attempt+1, HTTP_RETRY_COUNT, timeout);
                 
        updateProgress(UpdateStatus::DOWNLOADING, 0, 
                      String("Firmware download attempt " + String(attempt+1) + "/" + 
                             String(HTTP_RETRY_COUNT)).c_str());
        
        httpCode = http.GET();
        if(httpCode == HTTP_CODE_OK) break;
        
        ESP_LOGW(TAG, "Firmware GET failed (attempt %d/%d): %d", attempt+1, HTTP_RETRY_COUNT, httpCode);
        http.end();
        
        // Exponential backoff - start with HTTP_RETRY_DELAY_MS and double each time
        int delayMs = HTTP_RETRY_DELAY_MS * (1 << attempt); // 1000, 2000, 4000, 8000, 16000 ms
        if(attempt+1 < HTTP_RETRY_COUNT) {
            ESP_LOGI(TAG, "Waiting %d ms before next attempt", delayMs);
            vTaskDelay(pdMS_TO_TICKS(delayMs));
        }
    }
    if (httpCode != HTTP_CODE_OK) {
        ESP_LOGE(TAG, "Firmware download failed: %d", httpCode);
        updateProgress(UpdateStatus::ERROR, 0, "Firmware: Firmware download failed");
        return false;
    }

    // Check available space
    size_t firmwareSize = http.getSize();
    
    // Pre-validation checks
    if (firmwareSize < 1024*100) {  // Sanity check - 100KB minimum expected
        ESP_LOGE(TAG, "Firmware size too small (%d bytes), possible corruption", firmwareSize);
        updateProgress(UpdateStatus::ERROR, 0, "Firmware: Invalid firmware size");
        http.end();
        return false;
    }
    
    // Verify OTA partitions
    const esp_partition_t* running = esp_ota_get_running_partition();
    const esp_partition_t* update = esp_ota_get_next_update_partition(NULL);
    if (!running || !update) {
        ESP_LOGE(TAG, "Failed to get partitions");
        updateProgress(UpdateStatus::ERROR, 0, "Firmware: Failed to get partitions");
        http.end();
        return false;
    }
    ESP_LOGI(TAG, "Running partition: %s, update target: %s", 
             running->label, update->label);
    
    // Check for sufficient free memory - much more lenient now with chunked downloads
    size_t freeHeap = ESP.getFreeHeap();
    size_t minRequiredHeap = 40 * 1024; // Only require 40KB minimum
    
    ESP_LOGI(TAG, "Free heap: %d bytes, firmware size: %d bytes", freeHeap, firmwareSize);
    
    if (freeHeap < minRequiredHeap) {
        ESP_LOGE(TAG, "Critically low memory: %d bytes (minimum required: %d bytes)", 
                freeHeap, minRequiredHeap);
        updateProgress(UpdateStatus::ERROR, 0, "Firmware: Critically low memory");
        http.end();
        return false;
    }
    
    // Check if we can roll back (just for logging, not a blocker)
    if (!Update.canRollBack()) {
        ESP_LOGW(TAG, "No valid firmware to roll back to");
    }
    
    if (!Update.begin(firmwareSize > 0 ? firmwareSize : UPDATE_SIZE_UNKNOWN)) {
        ESP_LOGE(TAG, "Firmware: Not enough space for update");
        updateProgress(UpdateStatus::ERROR, 0, "Firmware: Not enough space for update");
        http.end();
        return false;
    }

    // Set up MD5 checking
    MD5Builder md5;
    md5.begin();
    
    // Prepare for download with improved resilience
    WiFiClient* stream = http.getStreamPtr();
    unsigned long lastProgressTime = millis(); // Track time since last progress
    unsigned long connectStartTime = millis(); // Track total download time
    unsigned long lastWatchdogKick = millis(); // Track watchdog kicks
    
    // Calculate smaller buffer size based on heap
    size_t bufferSize = std::min(size_t(4096), ESP.getFreeHeap() / 16); // Use at most 1/16 of free heap, max 4KB
    if (bufferSize < 1024) bufferSize = 1024; // Minimum 1KB
    
    ESP_LOGI(TAG, "Using download buffer size: %u bytes (free heap: %u bytes)", 
             bufferSize, ESP.getFreeHeap());
    
    // Create a buffer just for this operation if needed
    std::unique_ptr<uint8_t[]> localBuffer;
    uint8_t* downloadBuf = nullptr;
    
    if (bufferSize <= BUFFER_SIZE) {
        // Use the existing downloadBuffer
        downloadBuf = downloadBuffer.get();
    } else {
        // Create a larger buffer for this specific download
        localBuffer.reset(new uint8_t[bufferSize]);
        downloadBuf = localBuffer.get();
    }
    
    updateProgress(UpdateStatus::DOWNLOADING, 0, "Firmware download started");
    
    // Download and verify firmware
    if (firmwareSize > 0) {
        size_t remaining = firmwareSize;
        size_t totalReceived = 0;
        size_t failedReads = 0;  // Count consecutive read failures
        const size_t MAX_FAILED_READS = 10; // Allow up to 10 consecutive failures
        
        while (remaining > 0 && failedReads < MAX_FAILED_READS) {
            // Check for cancellation
            if(g_UpdateCancelFlag) { 
                ESP_LOGE(TAG, "Update cancelled by user");
                Update.abort(); 
                http.end(); 
                return false; 
            }
            
            // Check WiFi status every 5 seconds
            if (millis() - lastWatchdogKick > 5000) {
                if (WiFi.status() != WL_CONNECTED) {
                    ESP_LOGE(TAG, "WiFi disconnected during download");
                    Update.abort(); 
                    http.end(); 
                    return false;
                }
                
                // Reset watchdog to prevent timeout
                esp_task_wdt_reset();
                lastWatchdogKick = millis();
            }
            
            // Calculate optimum chunk size to balance speed vs reliability
            size_t chunk = std::min(remaining, bufferSize);
            size_t available = stream->available();
            
            if (available > 0) {
                // Read what's available (up to our chunk size)
                size_t read = stream->readBytes(downloadBuf, std::min(available, chunk));
                
                if (read > 0) {
                    // Reset failure counter on successful read
                    failedReads = 0;
                    
                    // Update MD5 and write firmware
                    md5.add(downloadBuf, read);
                    if (Update.write(downloadBuf, read) != read) {
                        ESP_LOGE(TAG, "Write failed");
                        Update.abort();
                        http.end();
                        return false;
                    }
                    
                    remaining -= read;
                    totalReceived += read;
                    lastProgressTime = millis(); // Reset progress timer
                    
                    // Update progress every ~2% or at least every 2 seconds
                    static int lastPercent = 0;
                    int percent = (totalReceived * 100) / firmwareSize;
                    if (percent != lastPercent || (millis() - lastProgressTime > 2000)) {
                        updateProgress(UpdateStatus::DOWNLOADING, percent, "Firmware");
                        lastPercent = percent;
                    }
                } else {
                    // Handle zero bytes read despite data being available
                    failedReads++;
                    ESP_LOGW(TAG, "Read returned 0 bytes despite data available (failure %d/%d)", 
                              failedReads, MAX_FAILED_READS);
                    delay(100); // Short delay before retry
                }
            } else {
                // No data currently available
                if (millis() - lastProgressTime > 60000) {
                    // No progress for 60 seconds - too long
                    ESP_LOGE(TAG, "Download timed out - no progress for 60 seconds");
                    Update.abort();
                    http.end();
                    return false;
                }
                
                // Send periodic progress updates to keep client informed
                if (millis() - lastWatchdogKick > 5000) {
                    int percent = (totalReceived * 100) / firmwareSize;
                    updateProgress(UpdateStatus::DOWNLOADING, percent, 
                                  String("Firmware: " + String(percent) + "% - waiting for data...").c_str());
                }
                
                delay(100); // Short delay to prevent CPU hogging
            }
        }
        
        // Check if we failed due to too many consecutive read failures
        if (failedReads >= MAX_FAILED_READS) {
            ESP_LOGE(TAG, "Too many consecutive read failures");
            Update.abort();
            http.end();
            return false;
        }
    } else {
        // Unknown size: stream until end (less common case)
        ESP_LOGW(TAG, "Firmware size unknown, streaming until end");
        size_t totalReceived = 0;
        size_t emptyReads = 0;
        const size_t MAX_EMPTY_READS = 20;  // Allow up to 20 empty reads before considering done
        
        while (emptyReads < MAX_EMPTY_READS) {
            if(g_UpdateCancelFlag){ Update.abort(); http.end(); return false; }
            
            // Reset watchdog periodically
            if (millis() - lastWatchdogKick > 5000) {
                esp_task_wdt_reset();
                lastWatchdogKick = millis();
            }
            
            size_t available = stream->available();
            if (available > 0) {
                size_t read = stream->readBytes(downloadBuf, std::min(available, bufferSize));
                if (read > 0) {
                    emptyReads = 0;  // Reset empty read counter
                    md5.add(downloadBuf, read);
                    if (Update.write(downloadBuf, read) != read) {
                        ESP_LOGE(TAG, "Write failed");
                        Update.abort();
                        http.end();
                        return false;
                    }
                    totalReceived += read;
                    lastProgressTime = millis();
                    
                    // Just update with received byte count since we don't know total
                    updateProgress(UpdateStatus::DOWNLOADING, 0, 
                                 String("Firmware: " + String(totalReceived / 1024) + "KB received").c_str());
                } else {
                    emptyReads++;
                    delay(100);
                }
            } else {
                // No data available
                if (totalReceived > 0 && millis() - lastProgressTime > 30000) {
                    // If we've received data but nothing for 30s, probably done
                    ESP_LOGI(TAG, "No data for 30s after receiving %d bytes, assuming download complete", totalReceived);
                    break;
                }
                
                emptyReads++;
                delay(100);
            }
        }
        
        if (totalReceived < 100*1024) {  // Less than 100KB received
            ESP_LOGE(TAG, "Downloaded firmware too small (%d bytes)", totalReceived);
            Update.abort();
            http.end();
            return false;
        }
    }
    
    // Verify MD5
    md5.calculate();
    String calculatedMd5 = md5.toString();
    updateProgress(UpdateStatus::VERIFYING, 95, "firmware");
    if (!calculatedMd5.equals(expectedMd5)) {
        ESP_LOGE(TAG, "MD5 mismatch. Expected: %s, Got: %s", expectedMd5, calculatedMd5.c_str());
        updateProgress(UpdateStatus::MD5_FAILED, 0, "Firmware: MD5 mismatch");
        Update.abort();
        http.end();
        return false;
    }

    // Finish update
    if (!Update.end()) {
        ESP_LOGE(TAG, "Update end failed: %d", Update.getError());
        updateProgress(UpdateStatus::ERROR, 0, "Firmware: Update failed");
        
        // Try to roll back if possible
        if (Update.hasError() && Update.canRollBack()) {
            ESP_LOGI(TAG, "Rolling back to previous version");
            Update.rollBack();
            // Don't restart here, let the main task do it
        }
        
        http.end();
        return false;
    }

    http.end();
    updateProgress(UpdateStatus::COMPLETE, 100, "Firmware: Complete");
    return true;
}

bool AppUpdater::IsUpdateAvailable(){
    return updateAvailable;
}

String AppUpdater::buildUrl(const char* path) const {
    if(!path || !*path) return baseUrl; // just base
    String p(path);
    // If already absolute URL, pass through
    if(p.startsWith("http://") || p.startsWith("https://")) return p;
    // Strip leading slashes to avoid double
    while(p.startsWith("/")) p.remove(0,1);
    // Ensure baseUrl has single trailing slash
    String b = baseUrl;
    if(!b.endsWith("/")) b += "/";
    return b + p;
}


AsyncEventSource* eventProgress = nullptr;
void startFirmwareUpdateTask(AsyncEventSource* evProg) {
    eventProgress = evProg;
    if(Update_Task_Handle) {
        ESP_LOGW(TAG, "Firmware update task already running");
        return;
    }
    // Create task with higher priority (3) and optimized stack size
    xTaskCreate(firmwareUpdateTask, "FirmwareUpdate", 1024*6, NULL, 3, &Update_Task_Handle);
}

void firmwareUpdateTask(void* parameter) {
    static const char* TAG = "UpdateTask";

    // Initialize watchdog timer for the update task
    esp_task_wdt_init(60, true); // 60 second timeout, panic on timeout
    esp_task_wdt_add(NULL);      // Add current task to watchdog

    try {
        loadUpdateJson();
        
        esp_task_wdt_reset();    // Reset watchdog timer after JSON loading

        // Initialize updater with smart pointer
        std::unique_ptr<AppUpdater> updater(new AppUpdater(
            LittleFS, localVersion, updateUrl.c_str(), "manifest.json", "firmware.bin"));
        updater->setProgressCallback(updateProgress);
        
        ESP_LOGI(TAG, "Starting update check from: %s", updateUrl.c_str());

        // Check and perform updates
        auto manifestResult = updater->checkManifest();
        
        if (manifestResult != AppUpdater::ManifestCheckResult::UPDATE_AVAILABLE) {
            // Handle different error cases
            std::string errorMsg;
            switch (manifestResult) {
                case AppUpdater::ManifestCheckResult::ERROR_FETCH_FAILED:
                    errorMsg = "Failed to fetch manifest";
                    break;
                case AppUpdater::ManifestCheckResult::ERROR_TOO_LARGE:
                    errorMsg = "Manifest file too large";
                    break;
                case AppUpdater::ManifestCheckResult::ERROR_PARSE_FAILED:
                    errorMsg = "Failed to parse manifest";
                    break;
                case AppUpdater::ManifestCheckResult::ERROR_NO_FILES_SECTION:
                    errorMsg = "Manifest missing files section";
                    break;
                case AppUpdater::ManifestCheckResult::ERROR_NO_VERSION:
                    errorMsg = "Manifest missing version section";
                    break;
                case AppUpdater::ManifestCheckResult::VERSION_CURRENT:
                    errorMsg = "Current version is up to date";
                    // This is not actually an error
                    ESP_LOGI(TAG, "No update needed: %s", errorMsg.c_str());
                    updateProgress(AppUpdater::UpdateStatus::MESSAGE, 0, errorMsg.c_str());
                    // Don't throw, just exit gracefully
                    break;
                default:
                    errorMsg = "Unknown manifest check error";
            }
            
            if (manifestResult != AppUpdater::ManifestCheckResult::VERSION_CURRENT) {
                ESP_LOGE(TAG, "Manifest check failed: %s", errorMsg.c_str());
                updateProgress(AppUpdater::UpdateStatus::ERROR, 0, errorMsg.c_str());
            }
        }

        if (updater->IsUpdateAvailable()) {
            bool filesUpdated = true;
            bool firmwareUpdated = false; // Initialize to false - only set to true if firmware is actually updated
            
            // Update files based on update mode
            if (g_UpdateMode == UpdateMode::UPDATE_FILES_ONLY || g_UpdateMode == UpdateMode::UPDATE_BOTH) {
                ESP_LOGI(TAG, "Update mode includes files, updating files...");
                filesUpdated = updater->updateFilesArray();
                if (!filesUpdated) { 
                    ESP_LOGW(TAG, "Some files failed to update");
                    if (g_UpdateMode == UpdateMode::UPDATE_FILES_ONLY) {
                        ESP_LOGE(TAG, "Files-only update failed");
                        updateProgress(AppUpdater::UpdateStatus::ERROR, 0, "Failed to update files");
                        // Skip to cleanup since this is files-only mode and it failed
                        goto cleanup;
                    } else {
                        ESP_LOGW(TAG, "File update failed, but continuing with firmware update");
                    }
                }
            } else {
                ESP_LOGI(TAG, "Skipping file updates (mode: firmware only)");
            }

            // Update firmware based on update mode
            if (g_UpdateMode == UpdateMode::UPDATE_FIRMWARE_ONLY || g_UpdateMode == UpdateMode::UPDATE_BOTH) {
                ESP_LOGI(TAG, "Update mode includes firmware, updating firmware...");
                firmwareUpdated = updater->updateApp();
                if (!firmwareUpdated) { 
                    ESP_LOGE(TAG, "Failed to update firmware");
                    updateProgress(AppUpdater::UpdateStatus::ERROR, 0, "Failed to update firmware");
                    // Skip to cleanup since firmware update failed
                    goto cleanup;
                }
            } else {
                ESP_LOGI(TAG, "Skipping firmware update (mode: files only)");
            }
            
            // Determine if we need to restart
            bool needsRestart = (g_UpdateMode == UpdateMode::UPDATE_FIRMWARE_ONLY || g_UpdateMode == UpdateMode::UPDATE_BOTH) && firmwareUpdated;
            
            if (needsRestart) {
                ESP_LOGI(TAG, "Firmware update successful, restarting...");
                sendUpdateMessage("Restarting... ", true, 100);
                vTaskDelay(2000);
                ESP.restart();
            } else {
                ESP_LOGI(TAG, "Update completed successfully (no restart required)");
                updateProgress(AppUpdater::UpdateStatus::COMPLETE, 100, "Update completed successfully");
            }
        }
        
    } catch (const std::exception& e) {
        ESP_LOGE(TAG, "Update failed with exception: %s", e.what());
        updateProgress(AppUpdater::UpdateStatus::ERROR, 0, e.what());
    } catch (...) {
        ESP_LOGE(TAG, "Update failed with unknown exception");
        updateProgress(AppUpdater::UpdateStatus::ERROR, 0, "Unknown error during update");
    }
    
cleanup:
    // Clean up watchdog before exit
    esp_task_wdt_delete(NULL);
    
    Update_Task_Handle = NULL;
    vTaskDelete(NULL);
}

void startVersionCheckTask() {
    if(versionCheckTask_Handle != NULL) {
        ESP_LOGW(TAG, "Version Check Tak already running");
        return;
    }
    // Create task with higher priority (3) and optimized stack size
    xTaskCreate(versionCheckTask, "VersionCheckTask", 1024*6, NULL, 3, &versionCheckTask_Handle);
}

void versionCheckTask(void* parameter){
    if(updateUrl == ""){
        loadUpdateJson();
    }
    AppUpdater updater(LittleFS, localVersion, updateUrl.c_str(), "manifest.json", "firmware.bin");
    
    auto manifestResult = updater.checkManifest();
    
    if (manifestResult == AppUpdater::ManifestCheckResult::UPDATE_AVAILABLE || 
        manifestResult == AppUpdater::ManifestCheckResult::VERSION_CURRENT) {
        otaVersion = updater.otaVersion; // capture remote
        ESP_LOGI(TAG, "Version check: remote=%s", otaVersion.toString().c_str());
    } else {
        ESP_LOGE(TAG, "Version check: manifest check failed with code %d", static_cast<int>(manifestResult));
    }
    
    versionCheckTask_Handle = NULL;
    vTaskDelete(NULL);
}

void loadUpdateJson(void) {
    try {
        ESP_LOGD(TAG, "loadUpdateJaon function...");
        if(updateUrl == "") {
            String updateJsonPath = "/system/update.json";

            // Read and parse update.json
            File file = LittleFS.open(updateJsonPath);
            if (!file) { 
                throw std::runtime_error("Failed to open update.json"); 
            }

            JsonDocument doc;
            DeserializationError error = deserializeJson(doc, file);
            file.close();
            
            if (error) { throw std::runtime_error("Failed to parse update.json"); }

            // Get update configuration
            JsonObject jObj = doc.as<JsonObject>();
            String folderName = jsonConstrainString(TAG, jObj, "folder", "latest/");
            String baseUrl = jsonConstrainString(TAG, jObj, "baseurl", "https://s3-minio.boothwizard.com/boothifier/");
            updateUrl = baseUrl + folderName;

            ESP_LOGD(TAG, "updateUrl: %s", updateUrl.c_str());
        }
    } catch (const std::exception& e) {
        ESP_LOGE(TAG, "Update failed: %s", e.what());
    }
}

void updateProgress(AppUpdater::UpdateStatus newStatus, int percentage, const char* message = nullptr) {

    char buffer[128];
    const char* msg;
    bool isComplete = false;

    const char* safeMsg = message ? message : "";
    switch (newStatus) {
        case AppUpdater::UpdateStatus::IDLE:
            snprintf(buffer, sizeof(buffer), "Update idle");
            msg = buffer;
            break;
        case AppUpdater::UpdateStatus::MESSAGE:
            msg = message ? message : "";
            break;
        case AppUpdater::UpdateStatus::DOWNLOADING:
            snprintf(buffer, sizeof(buffer), "%s: Download progress: %d%%", safeMsg, percentage);
            msg = buffer;
            break;
        case AppUpdater::UpdateStatus::VERIFYING:
            snprintf(buffer, sizeof(buffer), "%s: Verifying update: %d%%", safeMsg, percentage);
            msg = buffer;
            break;
        case AppUpdater::UpdateStatus::FILE_SKIPPED:
            snprintf(buffer, sizeof(buffer), "%s: File Skipped, up to date", safeMsg);
            msg = buffer;
            break;
        case AppUpdater::UpdateStatus::FILE_SAVED:
            snprintf(buffer, sizeof(buffer), "%s: File Saved", safeMsg);
            msg = buffer;
            break;
        case AppUpdater::UpdateStatus::MD5_FAILED:
            snprintf(buffer, sizeof(buffer), "%s: MD5 Verification Failed", safeMsg);
            msg = buffer;
            break;
        case AppUpdater::UpdateStatus::COMPLETE:
            snprintf(buffer, sizeof(buffer), "Firmware Update Complete!!!");
            msg = buffer;
            isComplete = true;
            break;
        case AppUpdater::UpdateStatus::ERROR:
            snprintf(buffer, sizeof(buffer), "Error!: %s", safeMsg);
            msg = buffer;
            break;
        default:
            snprintf(buffer, sizeof(buffer), "Unknown update status: %d", (int)newStatus);
            msg = buffer;
            break;
    }

    ESP_LOGI(TAG, "%s", msg);
    sendUpdateMessage(msg, isComplete, percentage);
}

void sendUpdateMessage(const char* message, bool complete, int progress = -1) {

    if(eventProgress && eventProgress->count() > 0) {
        // This is for the web client and not the BLE client
        JsonDocument jsonDoc;
        jsonDoc["message"] = message;
        jsonDoc["complete"] = complete;
        jsonDoc["progress"] = progress;
        String strMessage;
        serializeJson(jsonDoc, strMessage);
        eventProgress->send(strMessage.c_str(), "update", millis());
    }
    else{
        ESP_LOGW(TAG, "No clients connected to event source");
    }

    bleUpgrade_send_message(message);
}

// Convenience functions for setting update mode
void setGlobalUpdateMode(UpdateMode mode) {
    g_UpdateMode = mode;
    ESP_LOGI(TAG, "Global update mode set to: %s", 
             mode == UpdateMode::UPDATE_FILES_ONLY ? "UPDATE_FILES_ONLY" :
             mode == UpdateMode::UPDATE_FIRMWARE_ONLY ? "UPDATE_FIRMWARE_ONLY" : "UPDATE_BOTH");
}

UpdateMode getGlobalUpdateMode() {
    return g_UpdateMode;
}

void setUpdateModeFilesOnly() {
    setGlobalUpdateMode(UpdateMode::UPDATE_FILES_ONLY);
}

void setUpdateModeFirmwareOnly() {
    setGlobalUpdateMode(UpdateMode::UPDATE_FIRMWARE_ONLY);
}

void setUpdateModeBoth() {
    setGlobalUpdateMode(UpdateMode::UPDATE_BOTH);
}