#include <dirent.h>
#include <errno.h>
#include <inttypes.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <unistd.h>
#include "driver/spi_master.h"
#include "driver/usb_serial_jtag.h"
#include "esp_log.h"
#include "esp_random.h"
#include "esp_timer.h"
#include "esp_vfs_fat.h"
#include "esp_vfs_fat_nand.h"
#include "diskio_nand.h"
#include "diskio_impl.h"
#include "ff.h"
#include "spi_nand_flash.h"

#define ROOT "/extflash"
#define MAX_PAYLOAD 4352
#define TIMEOUT_MS 12000
/* Same SPI3 NAND pins as ../main/FLASH/flash.h on this board. */
#define PIN_MOSI 39
#define PIN_MISO 41
#define PIN_CLK  40
#define PIN_CS   42
#define FLASH_FREQ_HZ 40000000

enum { INFO=1, LIST=2, MKDIR=3, BEGIN=4, BLOCK=5, END=6, VERIFY=7,
       CANCEL=8, FORMAT_PREPARE=9, FORMAT_CONFIRM=10, HELLO=11 };
enum { OK=0, BAD_FRAME=1, BAD_PATH=2, NOT_MOUNTED=3, IO_ERROR=4,
       BAD_STATE=5, MISMATCH=6, DENIED=7 };

static spi_nand_flash_device_t *nand;
static bool mounted;
static FILE *active;
static char active_path[512];
static uint64_t written;
static uint32_t running_crc;
static uint32_t challenge;
static uint32_t challenge_seq;
static int64_t challenge_deadline;
static uint8_t request[16+MAX_PAYLOAD+4], response[16+MAX_PAYLOAD+4];
static uint8_t out_buf[MAX_PAYLOAD];
static uint8_t crc_buf[4096];
static uint8_t last_response[16+MAX_PAYLOAD+4];
static size_t last_response_size;
static uint32_t last_sequence;
static bool last_valid;
static bool session_ready;

static uint32_t u32(const uint8_t *p) {
    return (uint32_t)p[0] | (uint32_t)p[1]<<8 | (uint32_t)p[2]<<16 | (uint32_t)p[3]<<24;
}
static void put32(uint8_t *p, uint32_t v) {
    for (unsigned i=0;i<4;i++) p[i]=(uint8_t)(v>>(8*i));
}
static void put64(uint8_t *p, uint64_t v) {
    put32(p,(uint32_t)v); put32(p+4,(uint32_t)(v>>32));
}
static uint32_t crc_update(uint32_t crc, const uint8_t *p, size_t len) {
    crc=~crc;
    for(size_t i=0;i<len;i++) {
        crc^=p[i];
        for(int bit=0;bit<8;bit++) crc=(crc>>1) ^ (0xedb88320u & -(crc&1u));
    }
    return ~crc;
}
static bool send_bytes(const uint8_t *p, size_t len) {
    while(len) {
        int n=usb_serial_jtag_write_bytes(p,len,pdMS_TO_TICKS(5000));
        if(n<=0) return false;
        p+=n; len-=n;
    }
    return usb_serial_jtag_wait_tx_done(pdMS_TO_TICKS(5000))==ESP_OK;
}
static bool read_exact(uint8_t *p, size_t len) {
    int64_t deadline=esp_timer_get_time()+TIMEOUT_MS*1000LL;
    for(size_t i=0;i<len;i++) {
        int remain=(int)((deadline-esp_timer_get_time())/1000);
        if(remain<=0 || usb_serial_jtag_read_bytes(p+i,1,pdMS_TO_TICKS(remain))!=1) return false;
    }
    return true;
}
static bool read_frame(uint8_t *cmd, uint32_t *seq, uint16_t *len) {
    const uint8_t magic[]={'N','X','F','R'};
    size_t match=0;
    for(;;) {
        uint8_t b;
        if(usb_serial_jtag_read_bytes(&b,1,pdMS_TO_TICKS(1000))!=1) return false;
        if(b==magic[match]) request[match++]=b;
        else match=(b==magic[0]) ? (request[0]=b,1) : 0;
        if(match==4) break;
    }
    if(!read_exact(request+4,12)) return false;
    *cmd=request[5]; *seq=u32(request+8);
    uint32_t length=u32(request+12);
    if(request[4]!=1 || request[6] || request[7] || request[5]==0 || request[5]>HELLO || length>MAX_PAYLOAD) return false;
    *len=(uint16_t)length;
    if(!read_exact(request+16,*len+4)) return false;
    return crc_update(0,request,16+*len)==u32(request+16+*len);
}
static void reply(uint8_t cmd, uint32_t seq, uint8_t status, const uint8_t *data, uint16_t length) {
    memcpy(response,"NXFR",4); response[4]=1; response[5]=cmd|0x80;
    response[6]=response[7]=0; put32(response+8,seq);
    put32(response+12,length+1); response[16]=status;
    if(length) memcpy(response+17,data,length);
    put32(response+17+length,crc_update(0,response,17+length));
    last_response_size=21+length;
    memcpy(last_response,response,last_response_size);
    last_sequence=seq; last_valid=true;
    send_bytes(response,last_response_size);
}
static bool valid_path(const uint8_t *p, size_t length, char *out) {
    if(length<strlen(ROOT) || length>=512 || memchr(p,0,length)) return false;
    memcpy(out,p,length); out[length]=0;
    if(strncmp(out,ROOT,strlen(ROOT)) || (out[strlen(ROOT)] && out[strlen(ROOT)]!='/')) return false;
    if(out[length-1]=='/' && length>strlen(ROOT)) return false;
    if(!strcmp(out,ROOT)) return true;
    const char *s=out+strlen(ROOT);
    while(*s) {
        if(*s++!='/') return false;
        const char *part=s;
        while(*s && *s!='/') {
            unsigned char c=(unsigned char)*s++;
            if(c<32 || c==92 || c==58 || c==42 || c==63 || c==34 || c==60 || c==62 || c==124) return false;
        }
        if(s==part || (s-part==1 && part[0]=='.') || (s-part==2 && part[0]=='.' && part[1]=='.')) return false;
        if(s>part && (s[-1]=='.' || s[-1]==' ')) return false;
    }
    return true;
}
static bool make_dirs(char *path, bool parents_only) {
    size_t start=strlen(ROOT)+1;
    for(size_t i=start;;i++) {
        if(path[i]=='/' || (!parents_only && !path[i])) {
            char c=path[i]; path[i]=0;
            struct stat st;
            if(mkdir(path,0777)!=0 && (stat(path,&st)!=0 || !S_ISDIR(st.st_mode))) {
                path[i]=c; return false;
            }
            path[i]=c;
        }
        if(!path[i]) return true;
    }
}
static bool read_checksum(const char *path, uint64_t *size, uint32_t *crc) {
    FILE *f=fopen(path,"rb");
    if(!f) return false;
    *size=0; *crc=0;
    size_t n;
    while((n=fread(crc_buf,1,sizeof(crc_buf),f))) { *size+=n; *crc=crc_update(*crc,crc_buf,n); }
    bool good=!ferror(f);
    if(fclose(f)!=0) good=false;
    return good;
}
static esp_err_t mount_volume(void) {
    esp_vfs_fat_mount_config_t cfg={.max_files=6,.format_if_mount_failed=false,.allocation_unit_size=16384};
    esp_err_t err=esp_vfs_fat_nand_mount(ROOT,nand,&cfg);
    mounted=(err==ESP_OK);
    return err;
}
static bool format_volume(void) {
    if(mounted) {
        if(esp_vfs_fat_nand_unmount(ROOT,nand)!=ESP_OK) return false;
        mounted=false;
    }
    BYTE drive;
    if(ff_diskio_get_drive(&drive)!=ESP_OK) return false;
    if(ff_diskio_register_nand(drive,nand)!=ESP_OK) return false;
    char name[3]={(char)('0'+drive),':',0};
    void *work=ff_memalloc(4096);
    FRESULT result=FR_NOT_ENOUGH_CORE;
    if(work) {
        MKFS_PARM opt={(BYTE)FM_ANY,0,0,0,16384};
        result=f_mkfs(name,&opt,work,4096);
        ff_memfree(work);
    }
    bool ok=result==FR_OK;
    if(ok) ok=disk_ioctl(drive,CTRL_SYNC,NULL)==RES_OK;
    ff_diskio_unregister(drive);
    ff_diskio_clear_pdrv_nand(nand);
    if(!ok) return false;
    return mount_volume()==ESP_OK;
}
static void close_active(void) {
    if(active) { fclose(active); active=NULL; }
    active_path[0]=0; written=0; running_crc=0;
}
static void execute(uint8_t cmd, uint32_t seq, const uint8_t *payload, uint16_t len) {
    uint8_t *out=out_buf; uint16_t count=0;
    uint8_t status=OK; char path[512];
    if(cmd==HELLO) {
        if(len!=4) status=BAD_FRAME;
        else {
            challenge=0;
            challenge_seq=0;
            challenge_deadline=0;
            close_active();
            session_ready=true;
            memcpy(out,payload,4);
            count=4;
        }
    } else if(!session_ready) status=BAD_STATE;
    else if(cmd==INFO) {
        if(len) status=BAD_FRAME;
        else {
        uint32_t size=0, pages=0; uint64_t total=0,free=0;
        spi_nand_flash_get_page_size(nand,&size);
        spi_nand_flash_get_page_count(nand,&pages);
        if(mounted) esp_vfs_fat_info(ROOT,&total,&free);
        put32(out,size); put32(out+4,pages); out[8]=mounted;
        put64(out+9,total); put64(out+17,free); count=25;
        }
    } else if(cmd==FORMAT_PREPARE) {
        if(active || len) status=BAD_STATE;
        else { challenge=esp_random(); if(!challenge) challenge=1;
            challenge_seq=seq;
            challenge_deadline=esp_timer_get_time()+30000000LL;
            put32(out,challenge); count=4; }
    } else if(cmd==FORMAT_CONFIRM) {
        if(active || len!=8 || !challenge || u32(payload)!=challenge ||
           u32(payload+4)!=challenge_seq || esp_timer_get_time()>challenge_deadline) status=DENIED;
        else { challenge=0; challenge_seq=0; challenge_deadline=0;
            if(!format_volume()) status=IO_ERROR; }
    } else if(cmd==CANCEL) {
        if(len) status=BAD_FRAME;
        else close_active();
    } else if(!mounted) status=NOT_MOUNTED;
    else if(cmd==BLOCK) {
        if(!active || len<=4 || len>4100 || written>UINT32_MAX || u32(payload)!=(uint32_t)written) status=BAD_STATE;
        else if(fwrite(payload+4,1,len-4,active)!=len-4) { close_active(); status=IO_ERROR; }
        else { written+=len-4; running_crc=crc_update(running_crc,payload+4,len-4);
            put64(out,written); count=8; }
    } else if(cmd==END) {
        if(!active) status=BAD_STATE;
        else if(len!=12) { close_active(); status=BAD_STATE; }
        else if(u32(payload)!=(uint32_t)written || u32(payload+4)!=(uint32_t)(written>>32) ||
                u32(payload+8)!=running_crc) { close_active(); status=MISMATCH; }
        else {
            bool ok=fflush(active)==0;
            if(ok) ok=fsync(fileno(active))==0;
            if(fclose(active)!=0) ok=false;
            active=NULL;
            if(ok) ok=spi_nand_flash_sync(nand)==ESP_OK;
            uint64_t size=0; uint32_t crc=0;
            if(ok) ok=read_checksum(active_path,&size,&crc);
            if(!ok) status=IO_ERROR;
            else if(size!=written || crc!=running_crc) status=MISMATCH;
            else { put64(out,size); put32(out+8,crc); count=12; }
            active_path[0]=0; written=0; running_crc=0;
        }
    } else if(cmd==LIST) {
        if(len<2 || !valid_path(payload+2,len-2,path)) status=BAD_PATH;
        else {
            DIR *dir=opendir(path);
            if(!dir) status=IO_ERROR;
            else {
                uint16_t skip=(uint16_t)(payload[0]|payload[1]<<8), index=0;
                struct dirent *ent=NULL;
                bool done=true;
                while((ent=readdir(dir))) {
                    if(!strcmp(ent->d_name,".") || !strcmp(ent->d_name,"..")) continue;
                    if(index++<skip) continue;
                    size_t name_len=strlen(ent->d_name);
                    if(name_len>255) { status=BAD_PATH; break; }
                    if(count+name_len+13>MAX_PAYLOAD) {
                        index--;
                        done=false;
                        break;
                    }
                    char child[768]; struct stat st;
                    snprintf(child,sizeof(child),"%s/%s",path,ent->d_name);
                    if(stat(child,&st)!=0) { status=IO_ERROR; break; }
                    out[count++]=S_ISDIR(st.st_mode)?1:0;
                    out[count++]=(uint8_t)name_len;
                    put64(out+count,(uint64_t)st.st_size); count+=8;
                    memcpy(out+count,ent->d_name,name_len); count+=name_len;
                }
                closedir(dir);
                if(status==OK) {
                    memmove(out+3,out,count);
                    out[0]=(uint8_t)index; out[1]=(uint8_t)(index>>8);
                    out[2]=done; count+=3;
                } else count=0;
            }
        }
    } else if(cmd==MKDIR || cmd==BEGIN || cmd==VERIFY) {
        if(!valid_path(payload,len,path) || !strcmp(path,ROOT)) status=BAD_PATH;
        else if(cmd==MKDIR) { if(!make_dirs(path,false)) status=IO_ERROR; }
        else if(cmd==BEGIN) {
            if(active) status=BAD_STATE;
            else if(!make_dirs(path,true)) status=IO_ERROR;
            else {
                active=fopen(path,"wb");
                if(!active) status=IO_ERROR;
                else { strcpy(active_path,path); written=0; running_crc=0; }
            }
        } else {
            uint64_t size; uint32_t crc;
            if(!read_checksum(path,&size,&crc)) status=IO_ERROR;
            else { put64(out,size); put32(out+8,crc); count=12; }
        }
    } else status=BAD_FRAME;
    reply(cmd,seq,status,out,count);
}
void app_main(void) {
    usb_serial_jtag_driver_config_t usb={.tx_buffer_size=8192,.rx_buffer_size=8192};
    ESP_ERROR_CHECK(usb_serial_jtag_driver_install(&usb));
    spi_bus_config_t bus={.mosi_io_num=PIN_MOSI,.miso_io_num=PIN_MISO,.sclk_io_num=PIN_CLK,
        .quadwp_io_num=-1,.quadhd_io_num=-1,.max_transfer_sz=8192};
    ESP_ERROR_CHECK(spi_bus_initialize(SPI3_HOST,&bus,SPI_DMA_CH_AUTO));
    spi_device_interface_config_t device={.clock_speed_hz=FLASH_FREQ_HZ,.mode=0,.spics_io_num=PIN_CS,
        .queue_size=10,.flags=0};
    spi_device_handle_t spi;
    ESP_ERROR_CHECK(spi_bus_add_device(SPI3_HOST,&device,&spi));
    spi_nand_flash_config_t config={.device_handle=spi,.gc_factor=0,
        .io_mode=SPI_NAND_IO_MODE_SIO,.flags=0};
    ESP_ERROR_CHECK(spi_nand_flash_init_device(&config,&nand));
    uint32_t pages=0, page_size=0;
    ESP_ERROR_CHECK(spi_nand_flash_get_page_count(nand,&pages));
    ESP_ERROR_CHECK(spi_nand_flash_get_page_size(nand,&page_size));
    esp_err_t mount_err=mount_volume();
    char startup[128];
    int startup_len=snprintf(startup,sizeof(startup),
        "NAND capacity: %llu bytes; mount: %s\nNAND-XFER READY\n",
        (unsigned long long)pages*page_size,mount_err==ESP_OK?"ok":"unmounted");
    ESP_ERROR_CHECK(startup_len<0 || startup_len>=sizeof(startup) ? ESP_FAIL : ESP_OK);
    ESP_ERROR_CHECK(send_bytes((const uint8_t *)startup,startup_len) ? ESP_OK : ESP_FAIL);
    for(;;) {
        uint8_t cmd; uint32_t seq; uint16_t len;
        if(!read_frame(&cmd,&seq,&len)) continue;
        if(cmd!=HELLO && last_valid && seq==last_sequence) { send_bytes(last_response,last_response_size); continue; }
        if(cmd==HELLO) last_valid=false;
        execute(cmd,seq,request+16,len);
    }
}