#include <stdio.h>
#include <string.h>

#include "pico/stdlib.h"
#include "pico/cyw43_arch.h"
#include "pico/sha256.h"
#include "hardware/flash.h"
#include "hardware/sync.h"
#include "hardware/structs/scb.h"

#include "lwip/pbuf.h"
#include "lwip/tcp.h"

// ---- OTA-Wire-Protokoll (siehe server/server.py) ----
#define OTA_MAGIC          0xC0FFEE01u
#define OTA_STATUS_NONE    0x00u
#define OTA_STATUS_UPDATE  0x01u

#define SERVER_IP   "192.168.161.25"
#define SERVER_PORT 9000

// ---- Flash-Layout ----
#define APP_BASE          0x10080000u
#define APP_FLASH_OFFSET  (APP_BASE - XIP_BASE)   // = 0x80000, fuer flash_range_*
#define APP_SLOT_SIZE     (3584u * 1024u)         // 3.5 MB, muss zum Server passen
#define SHA256_BYTES      32

#ifndef PICO_ID
#error "Bitte PICO_ID per -DPICO_ID=<n> setzen"
#endif

// ---------- Hilfen ----------

static void print_hash(const char *label, const uint8_t *hash) {
    printf("%s ", label);
    for (int i = 0; i < SHA256_BYTES; i++) {
        printf("%02x", hash[i]);
    }
    printf("\n");
}

static bool app_present(uint32_t base) {
    uint32_t sp = *(volatile uint32_t *)base;
    return sp > 0x20000000u && sp <= 0x20082000u;
}

__attribute__((noreturn))
static void jump_to_app(uint32_t base) {
    uint32_t sp    = *(volatile uint32_t *)base;
    uint32_t reset = *(volatile uint32_t *)(base + 4u);

    cyw43_arch_deinit();
    __asm volatile ("cpsid i" ::: "memory");

    scb_hw->vtor = base;
    __asm volatile ("dsb 0xf" ::: "memory");
    __asm volatile ("isb 0xf" ::: "memory");

    __asm volatile (
        "msr msp, %0\n\t"
        "bx  %1\n"
        :: "r" (sp), "r" (reset)
        : "memory"
    );
    __builtin_unreachable();
}

__attribute__((noreturn))
static void halt_with_led(void) {
    cyw43_arch_gpio_put(CYW43_WL_GPIO_LED_PIN, true);
    while (true) {
        tight_loop_contents();
    }
}

// ---------- Flash-Operationen mit IRQ-Maskierung ----------

static void flash_erase_sector(uint32_t flash_offset) {
    uint32_t irq = save_and_disable_interrupts();
    flash_range_erase(flash_offset, FLASH_SECTOR_SIZE);
    restore_interrupts(irq);
}

static void flash_program_page(uint32_t flash_offset, const uint8_t *data) {
    uint32_t irq = save_and_disable_interrupts();
    flash_range_program(flash_offset, data, FLASH_PAGE_SIZE);
    restore_interrupts(irq);
}

// ---------- App-Slot hashen ----------

static bool hash_app_slot(uint8_t out[SHA256_BYTES]) {
    pico_sha256_state_t st;
    if (pico_sha256_start_blocking(&st, SHA256_BIG_ENDIAN, true) != PICO_OK) {
        printf("pico_sha256_start fehlgeschlagen\n");
        return false;
    }
    pico_sha256_update_blocking(&st, (const uint8_t *)APP_BASE, APP_SLOT_SIZE);
    sha256_result_t result;
    pico_sha256_finish(&st, &result);
    memcpy(out, result.bytes, SHA256_BYTES);
    return true;
}

// ---------- OTA-Client ----------

typedef enum {
    OTA_PHASE_REQUEST,   // Request raus, warte auf Status-Byte
    OTA_PHASE_HEADER,    // Status=Update gelesen, sammle size + new_hash (36 Byte)
    OTA_PHASE_PAYLOAD,   // sammle payload (size Byte) — und schreibe ins Flash
    OTA_PHASE_DONE,
    OTA_PHASE_ERROR,
} ota_phase_t;

typedef struct {
    struct tcp_pcb  *pcb;
    ip_addr_t        remote;
    ota_phase_t      phase;
    uint8_t          request[40];

    // Update-Header
    uint32_t         expected_size;
    uint8_t          new_hash[SHA256_BYTES];
    uint32_t         header_have;
    uint8_t          header_buf[4 + SHA256_BYTES];

    // Payload + Flash
    uint32_t         payload_have;
    uint8_t          page_buf[FLASH_PAGE_SIZE];
    uint32_t         page_filled;

    volatile bool    done;
    bool             update_pending;
} ota_t;

static void ota_finish(ota_t *o) {
    if (o->pcb) {
        tcp_arg(o->pcb, NULL);
        tcp_recv(o->pcb, NULL);
        tcp_err(o->pcb, NULL);
        tcp_close(o->pcb);
        o->pcb = NULL;
    }
    o->done = true;
}

// Schreibt eine voll gefuellte page_buf an den Slot. Erase passiert genau dann,
// wenn das die erste Page eines neuen Sektors ist.
static void flush_page_to_flash(ota_t *o, uint32_t offset_in_slot) {
    uint32_t flash_off = APP_FLASH_OFFSET + offset_in_slot;
    if ((offset_in_slot % FLASH_SECTOR_SIZE) == 0) {
        flash_erase_sector(flash_off);
    }
    flash_program_page(flash_off, o->page_buf);
}

static void ota_handle_bytes(ota_t *o, const uint8_t *data, uint16_t len) {
    uint16_t i = 0;
    while (i < len && o->phase != OTA_PHASE_DONE && o->phase != OTA_PHASE_ERROR) {
        if (o->phase == OTA_PHASE_REQUEST) {
            uint8_t status = data[i++];
            printf("<- status=0x%02x\n", status);
            if (status == OTA_STATUS_NONE) {
                o->phase = OTA_PHASE_DONE;
                break;
            }
            if (status != OTA_STATUS_UPDATE) {
                printf("unerwartetes Statusbyte\n");
                o->phase = OTA_PHASE_ERROR;
                break;
            }
            o->update_pending = true;
            o->phase = OTA_PHASE_HEADER;
            o->header_have = 0;
        } else if (o->phase == OTA_PHASE_HEADER) {
            uint16_t take = (uint16_t)(sizeof(o->header_buf) - o->header_have);
            if (take > (uint16_t)(len - i)) take = (uint16_t)(len - i);
            memcpy(&o->header_buf[o->header_have], &data[i], take);
            o->header_have += take;
            i += take;
            if (o->header_have == sizeof(o->header_buf)) {
                memcpy(&o->expected_size, &o->header_buf[0], 4);
                memcpy(o->new_hash, &o->header_buf[4], SHA256_BYTES);
                printf("Update-Header: size=%u\n", (unsigned)o->expected_size);
                print_hash("new_hash:", o->new_hash);
                if (o->expected_size == 0 || o->expected_size > APP_SLOT_SIZE) {
                    printf("Size ausserhalb gueltigem Bereich\n");
                    o->phase = OTA_PHASE_ERROR;
                    break;
                }
                o->phase = OTA_PHASE_PAYLOAD;
                o->payload_have = 0;
                o->page_filled = 0;
            }
        } else if (o->phase == OTA_PHASE_PAYLOAD) {
            uint32_t rest_total = o->expected_size - o->payload_have;
            uint32_t rest_chunk = (uint32_t)(len - i);
            uint32_t free_in_page = FLASH_PAGE_SIZE - o->page_filled;
            uint32_t take = rest_chunk;
            if (take > rest_total)  take = rest_total;
            if (take > free_in_page) take = free_in_page;

            memcpy(&o->page_buf[o->page_filled], &data[i], take);
            o->page_filled += take;
            o->payload_have += take;
            i += (uint16_t)take;

            bool page_done = (o->page_filled == FLASH_PAGE_SIZE);
            bool last_byte = (o->payload_have == o->expected_size);

            if (page_done || last_byte) {
                uint32_t offset_in_slot = o->payload_have - o->page_filled;
                if (last_byte && o->page_filled < FLASH_PAGE_SIZE) {
                    // Letzte Page: mit 0xFF auffuellen, damit der Slot-Hash passt
                    memset(&o->page_buf[o->page_filled], 0xFF,
                           FLASH_PAGE_SIZE - o->page_filled);
                }
                flush_page_to_flash(o, offset_in_slot);

                // Fortschritt alle 64 KB als Lebenszeichen
                if (((o->payload_have & 0xFFFF) == 0) || last_byte) {
                    printf("  geschrieben: %u/%u Byte\n",
                           (unsigned)o->payload_have, (unsigned)o->expected_size);
                }

                o->page_filled = 0;
                if (last_byte) {
                    o->phase = OTA_PHASE_DONE;
                    break;
                }
            }
        }
    }
}

static err_t on_recv(void *arg, struct tcp_pcb *pcb, struct pbuf *p, err_t err) {
    ota_t *o = (ota_t *)arg;
    if (!p) {
        if (o->phase == OTA_PHASE_PAYLOAD && o->payload_have < o->expected_size) {
            printf("Verbindung zu frueh geschlossen (%u/%u)\n",
                   (unsigned)o->payload_have, (unsigned)o->expected_size);
            o->phase = OTA_PHASE_ERROR;
        } else if (o->phase != OTA_PHASE_DONE) {
            o->phase = OTA_PHASE_DONE;
        }
        ota_finish(o);
        return ERR_OK;
    }

    struct pbuf *q = p;
    while (q) {
        ota_handle_bytes(o, (const uint8_t *)q->payload, q->len);
        q = q->next;
    }
    tcp_recved(pcb, p->tot_len);
    pbuf_free(p);

    if (o->phase == OTA_PHASE_DONE || o->phase == OTA_PHASE_ERROR) {
        ota_finish(o);
    }
    return ERR_OK;
}

static void on_err(void *arg, err_t err) {
    ota_t *o = (ota_t *)arg;
    printf("TCP-Fehler: %d\n", err);
    o->pcb = NULL;
    o->phase = OTA_PHASE_ERROR;
    o->done = true;
}

static err_t on_connected(void *arg, struct tcp_pcb *pcb, err_t err) {
    ota_t *o = (ota_t *)arg;
    if (err != ERR_OK) {
        printf("connect-callback err=%d\n", err);
        o->phase = OTA_PHASE_ERROR;
        ota_finish(o);
        return err;
    }
    err_t w = tcp_write(pcb, o->request, sizeof(o->request), TCP_WRITE_FLAG_COPY);
    if (w != ERR_OK) {
        printf("tcp_write err=%d\n", w);
        o->phase = OTA_PHASE_ERROR;
        ota_finish(o);
        return w;
    }
    tcp_output(pcb);
    return ERR_OK;
}

// Statisch, damit das halbe KB nicht jedesmal auf dem Stack steht.
static ota_t g_ota;

static void run_ota_request(const uint8_t current_hash[SHA256_BYTES],
                            bool *update_was_received,
                            uint32_t *received_size,
                            uint8_t  out_new_hash[SHA256_BYTES],
                            bool *failed) {
    *update_was_received = false;
    *received_size = 0;
    *failed = false;

    ota_t *o = &g_ota;
    memset(o, 0, sizeof(*o));

    if (!ip4addr_aton(SERVER_IP, &o->remote)) {
        printf("ungueltige Server-IP\n");
        *failed = true;
        return;
    }

    uint32_t magic = OTA_MAGIC;
    uint32_t pid   = (uint32_t)PICO_ID;
    memcpy(&o->request[0], &magic, 4);
    memcpy(&o->request[4], &pid,   4);
    memcpy(&o->request[8], current_hash, SHA256_BYTES);

    o->pcb = tcp_new_ip_type(IP_GET_TYPE(&o->remote));
    if (!o->pcb) {
        printf("tcp_new fehlgeschlagen\n");
        *failed = true;
        return;
    }
    tcp_arg(o->pcb, o);
    tcp_recv(o->pcb, on_recv);
    tcp_err(o->pcb, on_err);
    o->phase = OTA_PHASE_REQUEST;

    printf("Verbinde TCP %s:%d\n", SERVER_IP, SERVER_PORT);
    cyw43_arch_lwip_begin();
    err_t err = tcp_connect(o->pcb, &o->remote, SERVER_PORT, on_connected);
    cyw43_arch_lwip_end();
    if (err != ERR_OK) {
        printf("tcp_connect err=%d\n", err);
        ota_finish(o);
        *failed = true;
        return;
    }

    // 60 s Gesamttimeout. Reicht fuer 260 KB Download inkl. ein paar Erase-Pausen.
    const uint32_t TIMEOUT_MS = 60000;
    uint32_t waited = 0;
    while (!o->done && waited < TIMEOUT_MS) {
        sleep_ms(50);
        waited += 50;
    }
    if (!o->done) {
        printf("OTA-Timeout, breche ab\n");
        ota_finish(o);
        *failed = true;
        return;
    }
    if (o->phase == OTA_PHASE_ERROR) {
        *failed = true;
        return;
    }
    *update_was_received = o->update_pending;
    if (o->update_pending) {
        *received_size = o->expected_size;
        memcpy(out_new_hash, o->new_hash, SHA256_BYTES);
    }
}

// ---------- Nach-Download: Tail-Erase + Hash-Verify ----------

// Erased alle Sektoren ab `start_offset` (im Slot, relativ) bis Slot-Ende.
static void erase_tail(uint32_t start_offset) {
    // Auf Sektor-Grenze hoch runden.
    uint32_t aligned = (start_offset + FLASH_SECTOR_SIZE - 1) & ~(FLASH_SECTOR_SIZE - 1);
    uint32_t total_sectors = APP_SLOT_SIZE / FLASH_SECTOR_SIZE;
    uint32_t first_sector  = aligned / FLASH_SECTOR_SIZE;

    printf("Tail-Erase: Sektoren %u..%u (= %u KB)\n",
           (unsigned)first_sector, (unsigned)(total_sectors - 1),
           (unsigned)((total_sectors - first_sector) * FLASH_SECTOR_SIZE / 1024));

    for (uint32_t s = first_sector; s < total_sectors; s++) {
        flash_erase_sector(APP_FLASH_OFFSET + s * FLASH_SECTOR_SIZE);
        // Optisches Feedback: LED toggeln alle 32 Sektoren (~ alle 1.4 s).
        if ((s & 0x1F) == 0) {
            cyw43_arch_gpio_put(CYW43_WL_GPIO_LED_PIN, (s & 0x20) != 0);
        }
    }
    cyw43_arch_gpio_put(CYW43_WL_GPIO_LED_PIN, false);
    printf("Tail-Erase fertig\n");
}

// ---------- main ----------

int main(void) {
    stdio_init_all();
    sleep_ms(2500);

    printf("\n=== Bootloader Etappe C  (pico_id=%d) ===\n", (int)PICO_ID);

    if (cyw43_arch_init()) {
        printf("cyw43_arch_init fehlgeschlagen\n");
        halt_with_led();
    }
    cyw43_arch_enable_sta_mode();

    uint8_t current_hash[SHA256_BYTES];
    printf("Hashe App-Slot (%u Byte)...\n", (unsigned)APP_SLOT_SIZE);
    absolute_time_t t0 = get_absolute_time();
    if (!hash_app_slot(current_hash)) {
        halt_with_led();
    }
    int64_t took = absolute_time_diff_us(t0, get_absolute_time()) / 1000;
    printf("SHA256 fertig in %lld ms\n", took);
    print_hash("current:", current_hash);

    printf("Verbinde mit WLAN \"%s\"...\n", WIFI_SSID);
    if (cyw43_arch_wifi_connect_timeout_ms(
            WIFI_SSID, WIFI_PASSWORD, CYW43_AUTH_WPA2_AES_PSK, 15000) != 0) {
        printf("WLAN-Verbindung fehlgeschlagen — boote alte App\n");
        goto boot;
    }
    printf("WLAN OK\n");

    bool update_received = false;
    bool fail = false;
    uint32_t received_size = 0;
    uint8_t  expected_new_hash[SHA256_BYTES];
    run_ota_request(current_hash, &update_received, &received_size,
                    expected_new_hash, &fail);

    if (fail) {
        printf("OTA-Request fehlgeschlagen — boote alte App\n");
        goto boot;
    }
    if (!update_received) {
        printf("Server: kein Update\n");
        goto boot;
    }

    // Update wurde komplett ins Flash geschrieben. Jetzt Tail erasen
    // und den ganzen Slot neu hashen.
    erase_tail(received_size);

    uint8_t post_hash[SHA256_BYTES];
    printf("Verifiziere durch erneutes Hashen des Slots...\n");
    t0 = get_absolute_time();
    if (!hash_app_slot(post_hash)) {
        printf("Post-Hash fehlgeschlagen — bleibe stehen\n");
        halt_with_led();
    }
    took = absolute_time_diff_us(t0, get_absolute_time()) / 1000;
    printf("Post-Hash fertig in %lld ms\n", took);
    print_hash("got:     ", post_hash);
    print_hash("expected:", expected_new_hash);

    if (memcmp(post_hash, expected_new_hash, SHA256_BYTES) != 0) {
        printf("HASH-MISMATCH nach Flash-Write — bleibe stehen, springe nicht.\n");
        halt_with_led();
    }
    printf("Update verifiziert. \\o/\n");

boot:
    if (!app_present(APP_BASE)) {
        printf("Kein gueltiges App-Image bei 0x%08x — bleibe stehen\n",
               (unsigned)APP_BASE);
        halt_with_led();
    }
    printf("Springe in die App...\n");
    sleep_ms(100);
    jump_to_app(APP_BASE);
}