// librerie
#include <fcntl.h>
#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "esp_log.h"
#include "mbedtls/sha256.h"

#include "esp_partition.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"

// costanti
#define MAX_SIZE 4194304 // dimensione massima in byte
#define MAX_TOTAL_BITS 32 // numero massimo di bit
#define MIN_TOTAL_CRPS 2 // numero minimo di CRP
#define SHA256_DIGEST_LENGTH 32 // lunghezza di un hash SHA-256 in byte
#define TAG "Attestazione"

// struttura per una challenge-response pair (CRP)
typedef struct {
    uint32_t challenge; // challenge
    uint32_t response; // risposta
} crp_type;

// struttura per una memoria
typedef struct {
    /* const */ esp_partition_t *partition; // DA RIVEDERE
    long size; // dimensione in byte
    int block_size; // dimensione di un blocco in byte
    int total_blocks; // numero di blocchi
} memory_type;

// struttura per una physically unclonable function (PUF)
typedef struct {
    int seed; // seme per un generatore di numeri casuali
    int total_response_bits; // numero di bit di una risposta
    float bit_error_rate; // bit error rate (BER) compreso tra 0 e 1
    crp_type *used_crps; // CRP utilizzate
    int total_used_crps; // numero di CRP utilizzate
} puf_type;

// funzioni per la gestione di una PUF
uint32_t add_noise(puf_type *puf, uint32_t response); // funzione per aggiungere del rumore alla risposta di una PUF
void initialize_puf(puf_type *puf, int seed, int total_response_bits, float bit_error_rate); // funzione per inizializzare una PUF
void print_puf(puf_type *puf); // funzione per stampare le informazioni di una PUF
uint32_t query_puf(puf_type *puf, uint32_t challenge); // funzione per interrogare una PUF

// funzioni per la gestione di una memoria
void initialize_memory(memory_type *memory, int block_size); // funzione per inizializzare una memoria
uint8_t *get_block(memory_type *memory, int block_index); // funzione per ottenere un blocco
void print_block(uint8_t *block, int size, char format); // funzione per stampare un blocco
void print_memory(memory_type *memory); // funzione per stampare le informazioni di una memoria

// funzioni per l'esecuzione dell'attestazione
uint8_t *calculate_block_hash(uint8_t *block, int size, uint8_t *previous_hash); // funzione per calcolare l'hash di un blocco
uint8_t *run_attestation(puf_type *puf, memory_type *memory, uint32_t challenge, int nonce, int security_parameter); // funzione per eseguire l'attestazione

// funzioni di utilità
uint32_t bits_to_unsigned_int(int *bits, int total_bits); // funzione per convertire dei bit in un intero senza segno
void print_sha256_hash(uint8_t hash[SHA256_DIGEST_LENGTH]); // funzione per stampare un hash SHA-256
int *unsigned_int_to_bits(uint32_t unsigned_integer, int total_bits); // funzione per convertire un intero senza segno in bit

// main
void app_main() {
    puf_type puf; // PUF
    memory_type memory; // memoria
    uint32_t challenge = 123; // challenge
    int nonce = 456; // nonce
    int security_parameter = 100; // security parameter
    uint8_t *final_hash; // hash finale

    initialize_puf(&puf, 42, 32, 0.1); // inizializzazione della PUF
    print_puf(&puf); // stampa delle informazioni della PUF

    printf("\n"); // stampa di una nuova riga

    initialize_memory(&memory, 1024); // inizializzazione della memoria
    print_memory(&memory); // stampa delle informazioni della memoria

    printf("\n"); // stampa di una nuova riga

    final_hash = run_attestation(&puf, &memory, challenge, nonce, security_parameter); // esecuzione dell'attestazione

    printf("\nHash finale: ");
    print_sha256_hash(final_hash); // stampa dell'hash finale

    free(final_hash); // liberazione della memoria allocata per l'hash finale
    free(puf.used_crps); // liberazione della memoria allocata per le CRP utilizzate

    /* if (close(memory.fd) == -1) { // se la funzione close non è andata a buon fine
        ESP_LOGE(TAG, "Errore durante la chiusura del file descriptor\n"); // stampa dell'errore
        exit(1); // uscita
    } */
}

// funzione per aggiungere del rumore alla risposta di una PUF
uint32_t add_noise(puf_type *puf, uint32_t response) {
    int *response_bits; // bit della risposta
    int total_bits_to_flip; // numero di bit da invertire
    int bit_to_flip; // bit da invertire
    uint32_t noisy_response; // risposta con rumore

    response_bits = unsigned_int_to_bits(response, puf->total_response_bits); // conversione della risposta

    total_bits_to_flip = puf->bit_error_rate * puf->total_response_bits; // calcolo del numero di bit da invertire

    for (int i = 0; i < total_bits_to_flip; i++) { // iterazione sul numero di bit da invertire
        bit_to_flip = rand() % puf->total_response_bits; // selezione di un bit
        response_bits[bit_to_flip] = (response_bits[bit_to_flip] + 1) % 2; // inversione di un bit
    }

    noisy_response = bits_to_unsigned_int(response_bits, puf->total_response_bits); // conversione della risposta con rumore

    free(response_bits); // liberazione della memoria allocata per i bit della risposta

    return noisy_response; // restituzione della risposta con rumore
}

// funzione per inizializzare una PUF
void initialize_puf(puf_type *puf, int seed, int total_response_bits, float bit_error_rate) {
    puf->seed = seed; // inizializzazione del seme
    puf->total_response_bits = total_response_bits; // inizializzazione del numero di bit di una risposta
    puf->bit_error_rate = bit_error_rate; // inizializzazione del BER
    puf->used_crps = NULL; // inizializzazione delle CRP utilizzate
    puf->total_used_crps = 0; // inizializzazione del numero di CRP utilizzate

    srand(puf->seed); // inizializzazione di un generatore di numeri casuali
}

// funzione per stampare le informazioni di una PUF
void print_puf(puf_type *puf) {
    printf("Informazioni della PUF:\n");
    printf("- Seme: %d\n", puf->seed); // stampa del seme
    printf("- Numero di bit di una risposta: %d\n", puf->total_response_bits); // stampa del numero di bit di una risposta
    printf("- Bit error rate: %.2f\n", puf->bit_error_rate); // stampa del BER
    printf("- Numero di CRP utilizzate: %d\n", puf->total_used_crps); // stampa del numero di CRP utilizzate
}

// funzione per interrogare una PUF
uint32_t query_puf(puf_type *puf, uint32_t challenge) {
    int used = 0; // flag per controllare se la challenge è già stata utilizzata
    int crp_index = -1; // indice di una CRP
    uint32_t clean_response; // risposta senza rumore
    uint64_t max_response; // risposta massima

    for (int i = 0; i < puf->total_used_crps && !used; i++) { // iterazione sul numero di CRP utilizzate
        if (puf->used_crps[i].challenge == challenge) { // se la challenge è già stata utilizzata
            used = 1; // modifica del flag
            crp_index = i; // memorizzazione dell'indice
        }
    }

    if (used) { // se la challenge è già stata utilizzata
        clean_response = puf->used_crps[crp_index].response; // ottenimento della risposta senza rumore
    } else { // se la challenge non è mai stata utilizzata
        max_response = ((uint64_t) 1 << puf->total_response_bits) - 1; // calcolo della risposta massima

        clean_response = rand() % (max_response + 1); // calcolo della risposta senza rumore

        puf->total_used_crps++; // incremento del numero di CRP utilizzate

        if (puf->total_used_crps == 1) { // se è la prima CRP utilizzata
            puf->used_crps = (crp_type *) malloc(sizeof(crp_type)); // allocazione della memoria per le CRP utilizzate

            if (puf->used_crps == NULL) { // se la funzione malloc non è andata a buon fine
                ESP_LOGE(TAG, "Errore durante l'allocazione della memoria per le CRP utilizzate\n"); // stampa dell'errore
                exit(1); // uscita
            }
        } else { // se non è la prima CRP utilizzata
            crp_type *temporary_pointer = (crp_type *) realloc(puf->used_crps, puf->total_used_crps * sizeof(crp_type)); // riallocazione della memoria per le CRP utilizzate

            if (temporary_pointer == NULL) { // se la funzione realloc non è andata a buon fine
                ESP_LOGE(TAG, "Errore durante la riallocazione della memoria per le CRP utilizzate\n"); // stampa dell'errore
                free(puf->used_crps); // liberazione della memoria allocata per le CRP utilizzate
                exit(1); // uscita
            }

            puf->used_crps = temporary_pointer; // memorizzazione del nuovo puntatore
        }

        puf->used_crps[puf->total_used_crps - 1].challenge = challenge; // memorizzazione della challenge utilizzata
        puf->used_crps[puf->total_used_crps - 1].response = clean_response; // memorizzazione della risposta senza rumore utilizzata
    }

    return add_noise(puf, clean_response); // restituzione della risposta con rumore
}

// funzione per inizializzare una memoria (DA RIVEDERE)
void initialize_memory(memory_type *memory, int block_size) {
    // Trova la partizione di default
    memory->partition = esp_partition_find_first(ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_ANY, NULL);
    
    if (memory->partition == NULL) {
        ESP_LOGE("MEMORY", "Errore: impossibile trovare la partizione.");
        return;
    }

    memory->size = memory->partition->size;
    memory->block_size = block_size;
    memory->total_blocks = memory->size / memory->block_size;
}

// funzione per ottenere un blocco (DA RIVEDERE)
uint8_t *get_block(memory_type *memory, int block_index) {
    if (block_index < 0 || block_index >= memory->total_blocks) {
        ESP_LOGE("MEMORY", "Errore: indice del blocco non valido.");
        return NULL;
    }

    uint8_t *block = (uint8_t *) malloc(memory->block_size);
    if (block == NULL) {
        ESP_LOGE("MEMORY", "Errore: allocazione della memoria per il blocco fallita.");
        return NULL;
    }

    size_t offset = block_index * memory->block_size;
    esp_err_t err = esp_partition_read(memory->partition, offset, block, memory->block_size);
    if (err != ESP_OK) {
        ESP_LOGE("MEMORY", "Errore durante la lettura del blocco %d: %s", block_index, esp_err_to_name(err));
        free(block);
        return NULL;
    }

    return block;
}

// funzione per stampare un blocco
void print_block(uint8_t *block, int size, char format) {
    printf("\n"); // stampa di una nuova riga

    if (format == 'b') { // se il formato è binario
        for (int i = 0; i < size; i++) { // iterazione sulla dimensione del blocco
            for (int j = 7; j >= 0; j--) { // iterazione sul numero di bit in un byte
                printf("%d", (block[i] >> j) & 1); // stampa di un bit
            }

            printf(" "); // stampa di uno spazio ogni byte

            if ((i + 1) % 8 == 0) { // se sono stati stampati 8 byte
                printf("\n"); // stampa di una nuova riga
            }
        }
    } else if (format == 'h') { // se il formato è esadecimale
        for (int i = 0; i < size; i++) { // iterazione sulla dimensione del blocco
            printf("%02x ", block[i]); // stampa di un byte

            if ((i + 1) % 16 == 0) { // se sono stati stampati 16 byte
                printf("\n"); // stampa di una nuova riga
            }
        }
    } else { // se il formato non è né binario né esadecimale
        ESP_LOGE(TAG, "Errore durante il riconoscimento del formato per la stampa del blocco\n"); // stampa dell'errore
        exit(1); // uscita
    }
}

// funzione per stampare le informazioni di una memoria
void print_memory(memory_type *memory) {
    printf("Informazioni della memoria:\n");
    /* printf("- File descriptor: %d\n", memory->fd); // stampa del file descriptor */
    printf("- Dimensione: %ld byte\n", memory->size); // stampa della dimensione
    printf("- Dimensione di un blocco: %d byte\n", memory->block_size); // stampa della dimensione di un blocco
    printf("- Numero di blocchi: %d\n", memory->total_blocks); // stampa del numero di blocchi
}

// funzione per calcolare l'hash di un blocco
uint8_t *calculate_block_hash(uint8_t *block, int size, uint8_t *previous_hash) {
    mbedtls_sha256_context context; // contesto per il calcolo dell'hash
    uint8_t *block_hash; // hash del blocco

    mbedtls_sha256_init(&context); // inizializzazione del contesto per il calcolo dell'hash

    block_hash = (uint8_t *) malloc(SHA256_DIGEST_LENGTH); // allocazione della memoria per l'hash del blocco

    if (block_hash == NULL) { // se la funzione malloc non è andata a buon fine
        ESP_LOGE(TAG, "Errore durante l'allocazione della memoria per l'hash del blocco\n"); // stampa dell'errore
        exit(1); // uscita
    }

    if (mbedtls_sha256_starts(&context, 0) != 0) { // se l'inizializzazione dell'algoritmo per il calcolo dell'hash non è andata a buon fine
        ESP_LOGE(TAG, "Errore durante l'inizializzazione dell'algoritmo per il calcolo dell'hash\n"); // stampa dell'errore
        free(block_hash); // liberazione della memoria allocata per l'hash del blocco
        mbedtls_sha256_free(&context); // liberazione del contesto
        exit(1); // uscita
    }

    if (mbedtls_sha256_update(&context, /* (const unsigned char *) */ block, size) != 0) { // se l'aggiornamento del contesto per il calcolo dell'hash non è andato a buon fine
        ESP_LOGE(TAG, "Errore durante l'aggiornamento del contesto per il calcolo dell'hash\n"); // stampa dell'errore
        free(block_hash); // liberazione della memoria allocata per l'hash del blocco
        mbedtls_sha256_free(&context); // liberazione del contesto
        exit(1); // uscita
    }

    if (previous_hash != NULL) { // se esiste un hash precedente
        if (mbedtls_sha256_update(&context, /* (const unsigned char *) */ previous_hash, SHA256_DIGEST_LENGTH) != 0) { // se l'aggiornamento del contesto per il calcolo dell'hash non è andato a buon fine
            ESP_LOGE(TAG, "Errore durante l'aggiornamento del contesto per il calcolo dell'hash\n"); // stampa dell'errore
            free(block_hash); // liberazione della memoria allocata per l'hash del blocco
            mbedtls_sha256_free(&context); // liberazione del contesto
            exit(1); // uscita
        }
    }

    if (mbedtls_sha256_finish(&context, block_hash) != 0) { // se la finalizzazione del calcolo dell'hash non è andata a buon fine
        ESP_LOGE(TAG, "Errore durante la finalizzazione del calcolo dell'hash\n"); // stampa dell'errore
        free(block_hash); // liberazione della memoria allocata per l'hash del blocco
        mbedtls_sha256_free(&context); // liberazione del contesto
        exit(1); // uscita
    }

    mbedtls_sha256_free(&context); // liberazione del contesto

    return block_hash; // restituzione dell'hash del blocco
}

// funzione per eseguire l'attestazione
uint8_t *run_attestation(puf_type *puf, memory_type *memory, uint32_t challenge, int nonce, int security_parameter) {
    int total_crps; // numero di CRP da utilizzare
    uint32_t response; // risposta della PUF
    uint32_t *responses; // risposte della PUF
    unsigned int seed; // seme per un generatore di numeri casuali
    int blocks_random_walk_size; // dimensione della random walk dei blocchi
    int *blocks_random_walk; // random walk dei blocchi
    uint8_t *block_hash; // hash di un blocco
    int *crps_random_walk; // random walk delle CRP
    int iterations = 0; // contatore delle iterazioni del ciclo principale
    int response_index; // indice di una risposta della PUF

    total_crps = (int) fmax(MIN_TOTAL_CRPS, memory->total_blocks * security_parameter * 0.01); // calcolo del numero di CRP da utilizzare

    response = query_puf(puf, challenge + nonce); // calcolo della prima risposta della PUF

    printf("Prima risposta della PUF: %lu\n", response); // stampa della prima risposta della PUF

    responses = (uint32_t *) malloc(total_crps * sizeof(uint32_t)); // allocazione della memoria per le risposte della PUF

    if (responses == NULL) { // se la funzione malloc non è andata a buon fine
        ESP_LOGE(TAG, "Errore durante l'allocazione della memoria per le risposte della PUF\n"); // stampa dell'errore
        exit(1); // errore
    }

    memset(responses, 0, total_crps * sizeof(uint32_t)); // inizializzazione delle risposte della PUF
    responses[0] = response; // memorizzazione della prima risposta della PUF

    seed = (response + nonce) % (int) (pow(2, 32) - 1); // inizializzazione del seme per un generatore di numeri casuali

    blocks_random_walk_size = memory->total_blocks * log2(memory->total_blocks); // calcolo della dimensione della random walk dei blocchi
    blocks_random_walk = (int *) malloc(blocks_random_walk_size * sizeof(int)); // allocazione della memoria per la random walk dei blocchi

    if (blocks_random_walk == NULL) { // se la funzione malloc non è andata a buon fine
        ESP_LOGE(TAG, "Errore durante l'allocazione della memoria per la random walk dei blocchi\n"); // stampa dell'errore
        exit(1); // errore
    }

    for (int i = 0; i < blocks_random_walk_size; i++) { // iterazione sulla dimensione della random walk dei blocchi
        blocks_random_walk[i] = rand_r(&seed) % memory->total_blocks; // calcolo di un elemento della random walk dei blocchi
    }

    block_hash = calculate_block_hash(get_block(memory, blocks_random_walk[0]), memory->block_size, NULL); // calcolo dell'hash del primo blocco

    crps_random_walk = (int *) malloc(total_crps * sizeof(int)); // allocazione della memoria per la random walk delle CRP

    if (crps_random_walk == NULL) { // se la funzione malloc non è andata a buon fine
        ESP_LOGE(TAG, "Errore durante l'allocazione della memoria per la random walk delle CRP\n"); // stampa dell'errore
        exit(1); // errore
    }

    for (int i = 0; i < total_crps; i++) { // iterazione sul numero di CRP da utilizzare
        crps_random_walk[i] = rand_r(&seed) % total_crps; // calcolo di un elemento della random walk delle CRP
    }

    for (int i = 1; i < blocks_random_walk_size; i++) { // iterazione sulla dimensione della random walk dei blocchi
        if (iterations >= total_crps) { // se la random walk delle CRP è stata esaurita
            for (int i = 0; i < total_crps; i++) { // iterazione sul numero di CRP da utilizzare
                crps_random_walk[i] = rand_r(&seed) % total_crps; // calcolo di un elemento della random walk delle CRP
            }

            iterations = 0; // azzeramento del contatore delle iterazioni del ciclo principale
        }

        response_index = crps_random_walk[iterations % total_crps]; // calcolo dell'indice di una risposta della PUF

        if (responses[response_index] == 0) { // se una risposta della PUF non è ancora stata calcolata
            response = query_puf(puf, response + nonce); // calcolo di una risposta della PUF
            responses[response_index] = response; // memorizzazione di una risposta della PUF
        } else { // se una risposta della PUF è già stata calcolata
            response = responses[response_index]; // ottenimento di una risposta della PUF
        }

        block_hash = calculate_block_hash(get_block(memory, blocks_random_walk[i]), memory->block_size, block_hash); // calcolo dell'hash di un blocco

        iterations++; // incremento del contatore delle iterazioni del ciclo principale
    }

    free(responses); // liberazione della memoria allocata per le risposte della PUF
    free(blocks_random_walk); // liberazione della memoria allocata per la random walk dei blocchi
    free(crps_random_walk); // liberazione della memoria allocata per la random walk delle CRP

    return block_hash; // restituzione dell'hash finale
}

// funzione per convertire dei bit in un intero senza segno
uint32_t bits_to_unsigned_int(int *bits, int total_bits) {
    uint32_t unsigned_integer = 0; // intero senza segno

    if (total_bits < 0 || total_bits > MAX_TOTAL_BITS) { // se il numero di bit non rispetta i limiti
        ESP_LOGE(TAG, "Errore nel numero di bit\n"); // stampa dell'errore
        exit(1); // uscita
    }

    for (int i = 0; i < total_bits; i++) { // iterazione sul numero di bit
        unsigned_integer |= (bits[total_bits - i - 1] << i); // conversione di un bit
    }

    return unsigned_integer; // restituzione dell'intero senza segno
}

// funzione per stampare un hash SHA-256
void print_sha256_hash(uint8_t hash[SHA256_DIGEST_LENGTH]) {
    for (int i = 0; i < SHA256_DIGEST_LENGTH; i++) { // iterazione sul numero di byte di un hash SHA-256
        printf("%02x", hash[i]); // stampa di un byte
    }

    printf("\n"); // stampa di una nuova riga
}

// funzione per convertire un intero senza segno in bit
int *unsigned_int_to_bits(uint32_t unsigned_integer, int total_bits) {
    int *bits; // bit

    if (total_bits < 0 || total_bits > MAX_TOTAL_BITS) { // se il numero di bit non rispetta i limiti
        ESP_LOGE(TAG, "Errore nel numero di bit\n"); // stampa dell'errore
        exit(1); // uscita
    }

    bits = (int *) malloc(total_bits * sizeof(int)); // allocazione della memoria per i bit

    if (bits == NULL) { // se la funzione malloc non è andata a buon fine
        ESP_LOGE(TAG, "Errore durante l'allocazione della memoria per i bit\n"); // stampa dell'errore
        exit(1); // uscita
    }

    for (int i = 0; i < total_bits; i++) { // iterazione sul numero di bit
        bits[total_bits - i - 1] = (unsigned_integer >> i) & 1; // conversione di un bit
    }

    return bits; // restituzione dei bit
}