// FactoryTestをWebからダウンロードして書き込む
// 参考：https://lang-ship.com/blog/work/m5stickc-web-update-ota/

#include <WiFi.h>
#include <Update.h>
#include <esp_partition.h>
#include <esp_ota_ops.h>
#include <SPIFFS.h>

// ボード判定用マクロ定義
#if defined(ARDUINO_M5STACK_STICKC_PLUS2)
#define M5STICKC_PLUS2
#elif defined(ARDUINO_M5STICK_C_PLUS) || defined(ARDUINO_M5STACK_STICKC_PLUS)
#define M5STICKC_PLUS
#else
#define M5STICKC_PLUS2 // デフォルトはPlus2とする
#endif

const char *ssid0 = "ics-ap"; // 802.11b/g (2.4GHz)only. 5GHz is not supported.
const char *password0 = "jikkenics";

WiFiClient client;
void execOTA(String host, int port, String bin);
void lcdprintln(String mes, uint16_t color, int y)
{
    if (y == 0)
        M5.Lcd.fillScreen(color);
    if (y < 999)
        M5.Lcd.setCursor(0, y, 4);
    M5.Lcd.setTextColor(BLACK, color);
    M5.Lcd.println(mes);
}
void start_WebOTA()
{
    // M5.begin();
    // M5.Lcd.setRotation(3);
    lcdprintln(" -- WebOTA --\n\n wifi connecting...", ORANGE, 0);

    WiFi.begin(ssid0, password0); // 接続開始
    int count = 50;
    M5.Speaker.tone(2000, 200);
    M5.delay(500);
    while (WiFi.status() != WL_CONNECTED)
    { // 接続中...
        M5.Speaker.tone(2000, 200);
        M5.delay(200);
        M5.delay(1000);
        M5.Lcd.print(".");
        count--;
        if (count < 1)
        {
            break;
        }
    }
    if (count > 0)
    {
        // 接続完了！！
        M5.Speaker.tone(4000, 1500);
        lcdprintln("  Wifi Connected!\n  ", GREEN, 0);
        String gotip = WiFi.localIP().toString(); // m5デバイスのIPアドレス
        M5.Lcd.println(gotip);

        String host = "cit.istlab.info";
        int port = 80;
#ifdef M5STICKC_PLUS2
        String bin = "/f/plus2/FactoryTest.ino.bin";
#endif
#ifdef M5STICKC_PLUS
        String bin = "/f/plus1/FactoryTest.ino.bin";
#endif
        execOTA(host, port, bin);
    }
    else
    {
        wifi_down();
    }
}

// Utility to extract header value from headers
String getHeaderValue(String header, String headerName)
{
    return header.substring(strlen(headerName.c_str()));
}
// OTA Logic
void execOTA(String host, int port, String bin)
{
    Serial.println("");
    Serial.println("=== execOTA Start ===");
    Serial.println("Host: " + String(host));
    Serial.println("Port: " + String(port));
    Serial.println("Bin: " + String(bin));
    Serial.println("Connected to Wi-Fi");
    
    long contentLength = 0;
    bool isValidContentType = false;
    lcdprintln("  Downloading...", YELLOW, 0);
    Serial.println("Attempting to connect to server...");
    if (client.connect(host.c_str(), port))
    {
        Serial.println("Successfully connected to server!");
        Serial.println("Fetching Bin: " + String(bin));
        
        String httpRequest = String("GET ") + bin + " HTTP/1.1\r\n" +
                            "Host: " + host + "\r\n" +
                            "Cache-Control: no-cache\r\n" +
                            "Connection: close\r\n\r\n";
        Serial.println("HTTP Request:");
        Serial.println(httpRequest);
        
        client.print(httpRequest);
        Serial.println("Waiting for server response...");
        unsigned long timeout = millis();
        while (client.available() == 0)
        {
            if (millis() - timeout > 5000)
            {
                Serial.println("Client Timeout after 5 seconds!");
                client.stop();
                return;
            }
            delay(100);
        }
        Serial.println("Server response received!");
        Serial.println("Reading HTTP response headers...");
        int headerCount = 0;
        while (client.available())
        {
            String line = client.readStringUntil('\n');
            line.trim();
            headerCount++;
            Serial.println("Header " + String(headerCount) + ": " + line);
            
            if (!line.length())
            {
                Serial.println("End of headers detected");
                break;
            }
            if (line.startsWith("HTTP/1.1"))
            {
                if (line.indexOf("200") < 0)
                {
                    Serial.println("ERROR: Got a non 200 status code: " + line);
                    Serial.println("Exiting OTA Update.");
                    break;
                }
                else
                {
                    Serial.println("SUCCESS: Got HTTP 200 OK");
                }
            }
            if (line.startsWith("Content-Length: "))
            {
                String lengthStr = getHeaderValue(line, "Content-Length: ");
                contentLength = atol(lengthStr.c_str());
                Serial.println("Content-Length detected: " + lengthStr + " -> " + String(contentLength) + " bytes");
            }
            if (line.startsWith("Content-Type: "))
            {
                String contentType = getHeaderValue(line, "Content-Type: ");
                Serial.println("Content-Type detected: " + contentType);
                if (contentType == "application/octet-stream")
                {
                    isValidContentType = true;
                    Serial.println("Valid content type confirmed!");
                }
            }
        }
    }
    else
    {
        Serial.println("ERROR: Connection to " + String(host) + ":" + String(port) + " failed!");
        Serial.println("Please check network setup and server availability");
        return;
    }
    
    Serial.println("=== HTTP Response Analysis ===");
    Serial.println("Final contentLength: " + String(contentLength));
    Serial.println("Final isValidContentType: " + String(isValidContentType));
    Serial.println("Condition check: contentLength && isValidContentType = " + String(contentLength && isValidContentType));
    if (contentLength && isValidContentType)
    {
        Serial.println("=== Starting OTA Update Process ===");
        
        // メモリ情報の表示
        Serial.println("=== Memory Information ===");
        Serial.println("Required size: " + String(contentLength) + " bytes (" + String(contentLength/1024) + " KB)");
        Serial.println("Free heap: " + String(ESP.getFreeHeap()) + " bytes");
        Serial.println("Total heap: " + String(ESP.getHeapSize()) + " bytes");
        Serial.println("Free sketch space: " + String(ESP.getFreeSketchSpace()) + " bytes (" + String(ESP.getFreeSketchSpace()/1024) + " KB)");
        Serial.println("Sketch size: " + String(ESP.getSketchSize()) + " bytes (" + String(ESP.getSketchSize()/1024) + " KB)");
        Serial.println("Flash chip size: " + String(ESP.getFlashChipSize()) + " bytes (" + String(ESP.getFlashChipSize()/1024/1024) + " MB)");
        
        // パーティション情報の表示
        Serial.println("=== Partition Information ===");
        Serial.println("Flash chip mode: " + String(ESP.getFlashChipMode()));
        Serial.println("Flash chip speed: " + String(ESP.getFlashChipSpeed()) + " Hz");
        
        // OTAパーティションの存在確認
        const esp_partition_t* ota_0 = esp_partition_find_first(ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_APP_OTA_0, NULL);
        const esp_partition_t* ota_1 = esp_partition_find_first(ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_APP_OTA_1, NULL);
        const esp_partition_t* factory = esp_partition_find_first(ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_APP_FACTORY, NULL);
        
        Serial.println("Factory partition: " + String(factory ? "Found" : "NOT FOUND"));
        if (factory) {
            Serial.println("  Factory size: " + String(factory->size) + " bytes (" + String(factory->size/1024) + " KB)");
        }
        
        Serial.println("OTA_0 partition: " + String(ota_0 ? "Found" : "NOT FOUND"));
        if (ota_0) {
            Serial.println("  OTA_0 size: " + String(ota_0->size) + " bytes (" + String(ota_0->size/1024) + " KB)");
        }
        
        Serial.println("OTA_1 partition: " + String(ota_1 ? "Found" : "NOT FOUND"));
        if (ota_1) {
            Serial.println("  OTA_1 size: " + String(ota_1->size) + " bytes (" + String(ota_1->size/1024) + " KB)");
        }
        
        if (!ota_0 && !ota_1) {
            Serial.println("ERROR: No OTA partitions found!");
            Serial.println("This board may not be configured for OTA updates.");
            Serial.println("Need to use a partition table with OTA partitions.");
            lcdprintln("ERROR: No OTA\npartitions found!", RED, 0);
            client.flush();
            return;
        }
        
        // OTA方式の選択
        if (ota_0 && !ota_1) {
            Serial.println("WARNING: Only OTA_0 partition found. Using single partition OTA mode.");
            // 現在のパーティションを取得
            const esp_partition_t* current = esp_ota_get_running_partition();
            Serial.println("Current running partition: " + String(current ? current->label : "unknown"));
            
            // OTA_0パーティションに直接書き込み（上書きモード）
            Serial.println("Using OTA_0 partition for direct overwrite");
        } else if (ota_0 && ota_1) {
            Serial.println("Both OTA partitions found. Using standard dual partition OTA.");
        }
        
        // OTA用の空き領域をチェック
        size_t otaSize = ESP.getFreeSketchSpace();
        Serial.println("Available OTA space: " + String(otaSize) + " bytes (" + String(otaSize/1024) + " KB)");
        
        if (contentLength > otaSize) {
            Serial.println("ERROR: Firmware too large for available OTA space!");
            Serial.println("Required: " + String(contentLength) + " bytes, Available: " + String(otaSize) + " bytes");
            Serial.println("Shortage: " + String(contentLength - otaSize) + " bytes");
            lcdprintln("ERROR: Not enough\nOTA space!", RED, 0);
            client.flush();
            return;
        }
        
        bool canBegin = false;
        
        if (ota_0 && !ota_1) {
            // 単一OTAパーティション構成：直接パーティション書き込み方式
            Serial.println("Single partition detected - Using direct partition write method");
            Serial.println("WARNING: This will overwrite the current firmware!");
            
            // 現在のパーティションを確認
            const esp_partition_t* current = esp_ota_get_running_partition();
            if (current && strcmp(current->label, "app0") == 0) {
                Serial.println("Currently running on app0 - will overwrite app0");
                
                // 直接パーティション書き込みフラグを設定
                canBegin = true;  // 直接書き込みモードを示すフラグ
                Serial.println("Direct partition write mode enabled");
            } else {
                Serial.println("ERROR: Running on unexpected partition");
                canBegin = false;
            }
        } else {
            // 標準デュアルパーティション構成
            Serial.println("Attempting standard dual partition OTA...");
            canBegin = Update.begin(contentLength);
        }
        
        Serial.println("OTA initialization result: " + String(canBegin));
        
        if (!canBegin) {
            Serial.println("Update.begin() failed!");
            Serial.println("Error code: " + String(Update.getError()));
            
            // エラーコードの詳細表示
            uint8_t error = Update.getError();
            switch(error) {
                case UPDATE_ERROR_OK: Serial.println("No error"); break;
                case UPDATE_ERROR_WRITE: Serial.println("Flash write failed"); break;
                case UPDATE_ERROR_ERASE: Serial.println("Flash erase failed"); break;
                case UPDATE_ERROR_READ: Serial.println("Flash read failed"); break;
                case UPDATE_ERROR_SPACE: Serial.println("Not enough space"); break;
                case UPDATE_ERROR_SIZE: Serial.println("Bad size given"); break;
                case UPDATE_ERROR_STREAM: Serial.println("Stream read timeout"); break;
                case UPDATE_ERROR_MD5: Serial.println("MD5 check failed"); break;
                case UPDATE_ERROR_MAGIC_BYTE: Serial.println("Wrong magic byte"); break;
                case UPDATE_ERROR_ACTIVATE: Serial.println("Could not activate"); break;
                case UPDATE_ERROR_NO_PARTITION: Serial.println("Partition could not be found"); break;
                case UPDATE_ERROR_BAD_ARGUMENT: Serial.println("Bad argument"); break;
                case UPDATE_ERROR_ABORT: Serial.println("Aborted"); break;
                default: Serial.println("Unknown error: " + String(error)); break;
            }
            Serial.println("Free space check passed but Update.begin() still failed");
        }
        if (canBegin)
        {
            lcdprintln("  Downloaded.\n Upgrading...(wait 2-3 min.)", GREENYELLOW, 999);
            Serial.println("Begin OTA. This may take 2 - 5 mins to complete. Things might be quite for a while.. Patience!");
            
            size_t written = 0;
            bool writeSuccess = false;
            
            if (ota_0 && !ota_1) {
                // 低レベルOTA API方式（最終手段）
                Serial.println("=== Low-Level OTA API Mode ===");
                Serial.println("Using esp_ota_ops API directly...");
                
                // 次回起動パーティションとして現在のパーティションを設定
                const esp_partition_t* update_partition = esp_ota_get_next_update_partition(NULL);
                if (update_partition == NULL) {
                    Serial.println("ERROR: esp_ota_get_next_update_partition failed");
                    // フォールバック：現在のパーティションを使用
                    update_partition = ota_0;
                }
                
                Serial.println("Update partition label: " + String(update_partition->label));
                Serial.println("Update partition size: " + String(update_partition->size) + " bytes");
                Serial.println("Update partition address: 0x" + String(update_partition->address, HEX));
                
                if (contentLength > update_partition->size) {
                    Serial.println("ERROR: Firmware too large for partition");
                    Serial.println("Required: " + String(contentLength) + ", Available: " + String(update_partition->size));
                    writeSuccess = false;
                } else {
                    // OTA操作開始
                    esp_ota_handle_t ota_handle = 0;
                    esp_err_t err = esp_ota_begin(update_partition, contentLength, &ota_handle);
                    if (err != ESP_OK) {
                        Serial.println("ERROR: esp_ota_begin failed: " + String(esp_err_to_name(err)));
                        
                        // 詳細なエラー診断
                        if (err == ESP_ERR_OTA_PARTITION_CONFLICT) {
                            Serial.println("");
                            Serial.println("=== OTA PARTITION CONFLICT DIAGNOSIS ===");
                            Serial.println("ISSUE: Cannot update running partition (app0)");
                            Serial.println("CAUSE: Single partition configuration detected");
                            Serial.println("REQUIREMENT: Dual OTA partition setup needed");
                            Serial.println("");
                            Serial.println("SOLUTION OPTIONS:");
                            Serial.println("1. Change Arduino IDE partition scheme:");
                            Serial.println("   Tools -> Partition Scheme -> 'Default 4MB with OTA'");
                            Serial.println("2. Or use 'Minimal SPIFFS (1.9MB APP with OTA/190KB SPIFFS)'");
                            Serial.println("3. Then recompile and upload the firmware");
                            Serial.println("");
                            Serial.println("CURRENT SETUP:");
                            Serial.println("  - app0: " + String(ota_0->size/1024) + "KB (RUNNING - cannot overwrite)");
                            Serial.println("  - app1: NOT FOUND (NEEDED for OTA)");
                            Serial.println("");
                            Serial.println("REQUIRED SETUP:");
                            Serial.println("  - app0: ~1.4MB (current firmware)");
                            Serial.println("  - app1: ~1.4MB (new firmware destination)");
                            Serial.println("  - OTA switches between app0 <-> app1");
                            
                            lcdprintln("ERROR: Need dual\\nOTA partitions!\\n\\nSee Serial for\\ndetailed solution", RED, 0);
                        } else {
                            Serial.println("Other OTA error: " + String(err));
                        }
                        
                        writeSuccess = false;
                    } else {
                        Serial.println("Low-level OTA began successfully");
                        Serial.println("OTA Handle: " + String((uint32_t)ota_handle, HEX));
                        
                        // ファームウェアデータを直接書き込み
                        uint8_t buffer[1024];
                        size_t totalWritten = 0;
                        
                        Serial.println("Starting low-level OTA write...");
                        while (client.available() > 0 && totalWritten < contentLength) {
                            size_t bytesToRead = min((size_t)client.available(), sizeof(buffer));
                            bytesToRead = min(bytesToRead, (size_t)(contentLength - totalWritten));
                            
                            size_t bytesRead = client.readBytes(buffer, bytesToRead);
                            if (bytesRead > 0) {
                                esp_err_t write_err = esp_ota_write(ota_handle, buffer, bytesRead);
                                if (write_err == ESP_OK) {
                                    totalWritten += bytesRead;
                                    
                                    // 進行状況表示
                                    if (totalWritten % (64 * 1024) == 0 || totalWritten == contentLength) {
                                        Serial.println("Written: " + String(totalWritten) + "/" + String(contentLength) + " bytes (" + String(totalWritten * 100 / contentLength) + "%)");
                                    }
                                } else {
                                    Serial.println("ERROR: esp_ota_write failed: " + String(esp_err_to_name(write_err)));
                                    break;
                                }
                            } else {
                                Serial.println("ERROR: Failed to read from stream");
                                break;
                            }
                            
                            // ウォッチドッグタイマーリセット
                            if (totalWritten % (32 * 1024) == 0) {
                                yield();
                            }
                        }
                        
                        written = totalWritten;
                        Serial.println("Low-level write completed: " + String(written) + " bytes");
                        
                        // OTA終了処理
                        if (written == contentLength) {
                            esp_err_t end_err = esp_ota_end(ota_handle);
                            if (end_err == ESP_OK) {
                                Serial.println("esp_ota_end successful");
                                
                                // ブートパーティション設定
                                esp_err_t set_boot_err = esp_ota_set_boot_partition(update_partition);
                                if (set_boot_err == ESP_OK) {
                                    Serial.println("Boot partition set successfully");
                                    writeSuccess = true;
                                } else {
                                    Serial.println("ERROR: esp_ota_set_boot_partition failed: " + String(esp_err_to_name(set_boot_err)));
                                    writeSuccess = false;
                                }
                            } else {
                                Serial.println("ERROR: esp_ota_end failed: " + String(esp_err_to_name(end_err)));
                                writeSuccess = false;
                            }
                        } else {
                            Serial.println("ERROR: Incomplete write, aborting OTA");
                            esp_ota_abort(ota_handle);
                            writeSuccess = false;
                        }
                    }
                }
            } else {
                // 標準デュアルパーティションOTA方式
                Serial.println("Using standard Update.writeStream()");
                written = Update.writeStream(client);
                writeSuccess = (written == contentLength);
            }
            if (written == contentLength)
            {
                Serial.println("Written : " + String(written) + " successfully");
            }
            else
            {
                Serial.println("Written only : " + String(written) + "/" + String(contentLength) + ". Retry?");
            }
            
            bool otaSuccess = false;
            if (writeSuccess) {
                if (ota_0 && !ota_1) {
                    // 低レベルOTA API方式の場合（既にブートパーティション設定済み）
                    Serial.println("Low-level OTA API completed successfully!");
                    otaSuccess = true;
                } else {
                    // 通常の標準OTA方式の場合
                    if (Update.end()) {
                        otaSuccess = Update.isFinished();
                        if (!otaSuccess) {
                            Serial.println("Update not finished? Something went wrong!");
                        }
                    } else {
                        Serial.println("Error Occurred. Error #: " + String(Update.getError()));
                    }
                }
            }
            
            if (otaSuccess)
            {
                M5.Speaker.tone(4000, 2000);
                lcdprintln("\\n\\n  OTA Complete!\\n Rebooting now...", GREENYELLOW, 0);
                M5.delay(2000);
                Serial.println("OTA Update completed successfully!");
                Serial.println("Rebooting device...");
                esp_restart(); // 再起動
            }
            else
            {
                if (ota_0 && !ota_1) {
                    Serial.println("Low-level OTA API failed");
                    lcdprintln("\\n\\n  OTA Failed!\\n Low-level error", RED, 0);
                } else {
                    Serial.println("Standard OTA update failed");
                    lcdprintln("\\n\\n  OTA Failed!\\n Try again", RED, 0);
                }
            }
        }
        else
        {
            Serial.println("Not enough space to begin OTA");
            client.flush();
        }
    }
    else
    {
        Serial.println("ERROR: Cannot start OTA update!");
        if (!contentLength) {
            Serial.println("Reason: No content length received");
        }
        if (!isValidContentType) {
            Serial.println("Reason: Invalid content type");
        }
        Serial.println("There was no valid content in the response");
        client.flush();
    }
    Serial.println("=== execOTA End ===");
}