"""Run the actual V1/V2/V3 CAN update and SaveFwInfo C functions on Windows.

Requires Visual Studio 2022 C tools. Flash, CAN, and reset are simulated;
the firmware functions are extracted from the current sources, not reimplemented.
"""
from pathlib import Path
import re
import subprocess
import tempfile

ROOT = Path(__file__).resolve().parents[1]
VCVARS = Path(r"C:\Program Files\Microsoft Visual Studio\2022\Professional\VC\Auxiliary\Build\vcvars32.bat")
VARIANTS = [
    ("V1", "mBMS_APP_V3117/APP", "ExtId", 0x08080000, 0x080E0000, 0x08010000),
    ("V2", "APP_R3221/App/Source", "ExtID", 0x60000, 0x80000, 0x10000),
    ("V3", "APP_R3323/App/Source", "ExtID", 0x60000, 0x80000, 0x10000),
]


def function(source, name):
    match = re.search(r"^[ \t]*(?:void|u8|uint8_t|uint32_t) " + name + r"\([^;]*?\)\s*(?://[^\n]*)?\s*\{", source, re.M)
    assert match, name
    end = source.index("{", match.start()) + 1
    depth = 1
    while depth:
        depth += (source[end] == "{") - (source[end] == "}")
        end += 1
    return source[match.start():end]


STUBS = r"""
#include <windows.h>
#include <stdint.h>
#include <stdbool.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
typedef uint8_t u8; typedef uint16_t u16; typedef int16_t s16;
typedef struct { uint32_t ExtId; uint8_t Data[8]; uint8_t DLC; } CanRxMsg;
typedef struct { uint32_t ExtID; uint8_t Data[8]; struct { struct { uint8_t DLC; } Control_f; } Cst; } stc_can_rxframe_t;
typedef struct {
    uint32_t AppFileSize;
    uint8_t BootVer[4], AppVer[4], FileName[30], BootUpdate, AppUpdate;
} FW_INFO;
FW_INFO fwinfo;
#define ADDR_FLASH_SECTOR_8 IMAGE_BASE
#define ADDR_FLASH_SECTOR_10 (IMAGE_END - 0x20000)
#define ADDR_FLASH_SECTOR_11 IMAGE_END
#define mBMS_FW_ADDRESS IMAGE_BASE
#define mBMS_FW_END_ADDRESS IMAGE_END
#define mBMS_FW_DN_ADDRESS IMAGE_BASE
#define mBMS_FW_DN_END_ADDRESS IMAGE_END
#define mBMS_FW_FLASH_SIZE (IMAGE_END - IMAGE_BASE)
#define mBMS_FW_PACKET_SIZE 2048
#define mBMS_FW_RN_ADDRESS 0x20000
#define SECTOR_SIZE 8192
#define mBMS_FW_INFO_BASE_ADDR INFO_BASE
static int fail_mode, erase_count, write_count, reset_count, response_count;
static uint8_t last_response[8];
uint8_t GetDevAddrIO(void) { return 1; }
void FLASH_If_Init(void) {}
uint32_t FLASH_If_SingleSectorErase(uint32_t addr) {
    erase_count++;
    if (fail_mode == 1) return 1;
    memset((void *)addr, 0xFF, addr == INFO_BASE ? sizeof(FW_INFO) : 8192);
    return 0;
}
uint32_t FLASH_If_Write(volatile uint32_t *addr, uint32_t *data, uint32_t words) {
    write_count++;
    if (fail_mode == 2) return 1;
    if (fail_mode == 4) return 0; /* Driver claims success without writing. */
    memcpy((void *)*addr, data, words * 4);
    if (fail_mode == 3) ((uint8_t *)*addr)[0] ^= 1;
    *addr += words * 4;
    return 0;
}
void SendPacket(uint32_t id, uint8_t *data) {
    memcpy(last_response, data, 8); response_count++;
}
void delay_os_ms(uint32_t ms) {}
void NVIC_SystemReset(void) { reset_count++; }
void Sys_Soft_Reset(void) { reset_count++; }
uint32_t SaveFwInfo(void);
"""

CASES = r"""
#define CHECK(x) do { if (!(x)) { fprintf(stderr, "line %d: %s\n", __LINE__, #x); exit(1); } } while (0)
static void be32(uint8_t *p, uint32_t n) {
    p[0]=n>>24; p[1]=n>>16; p[2]=n>>8; p[3]=n;
}
static PACKET packet;
static uint32_t byte_sum(const uint8_t *p, uint32_t n) {
    uint32_t sum=0; for (uint32_t i=0; i<n; i++) sum+=p[i]; return sum;
}
static uint32_t image_sum(void) {
    return byte_sum((const uint8_t *)IMAGE_BASE,fwinfo.AppFileSize);
}
static uint32_t wire_addr(void) {
    return IMAGE_BASE + WIRE_OFFSET;
}
static void reset_test(void) {
    memset(&fwinfo, 0, sizeof(fwinfo)); memset(&packet, 0, sizeof(packet));
    memset(&CanDLInfor, 0, sizeof(CanDLInfor));
    DownloadSessionActive=0; DownloadImageVerified=0;
    ResetDownloadBlock(); SET_DLC(packet,8);
    fail_mode=erase_count=write_count=reset_count=response_count=0;
}
static void start_download(void) {
    memset(&packet, 0, sizeof(packet)); be32(packet.Data+4, 8);
    SET_DLC(packet,8);
    CHECK(FwUpdateStartByCan(&packet)==0);
}
static void receive_image(void) {
    start_download();
    be32(packet.Data, wire_addr());
    CHECK(SectorEraseProcessByCan(&packet)==0);
    be32(packet.Data, wire_addr()); be32(packet.Data+4, 1);
    CHECK(WriteAddressSetByCan(&packet)==0);
    ((uint32_t *)packet.Data)[0]=0x20001000;
    ((uint32_t *)packet.Data)[1]=mBMS_FW_RN_ADDRESS|1;
    CHECK(WriteDataByCan(&packet)==0);
    CHECK(write_count==0); /* Flash is untouched while receiving. */
    uint32_t sum=byte_sum(packet.Data,8);
    be32(packet.Data, wire_addr()); be32(packet.Data+4,sum);
    CHECK(WriteDataChkSumByCan(&packet)==0);
    CHECK(DownloadBlockCommitted && write_count==1);
}
static void verify_image(void) {
    be32(packet.Data, fwinfo.AppFileSize); be32(packet.Data+4, image_sum());
    CHECK(FwImageChkSumByCan(&packet)==0);
    CHECK(fwinfo.AppUpdate==3 && DownloadImageVerified);
}
static void commit_rejected(void) {
    int before=response_count;
    packet.Data[0]=3;
    CHECK(UpdateRestartByCan(&packet)==9);
    CHECK(last_response[7]==9 && response_count==before+1 && !reset_count);
}
static void prepare_block(uint32_t count) {
    reset_test(); start_download(); fwinfo.AppFileSize=count*8;
    be32(packet.Data,wire_addr()); CHECK(SectorEraseProcessByCan(&packet)==0);
    be32(packet.Data,wire_addr()); be32(packet.Data+4,count);
    CHECK(WriteAddressSetByCan(&packet)==0);
}
static void fill_frame(uint32_t number) {
    PCanDLHeader h=(PCanDLHeader)&PACKET_ID(packet);
    h->BV.NO=number;
    SET_DLC(packet,8);
    for (uint32_t i=0;i<8;i++) packet.Data[i]=(uint8_t)((number*8+i)*13);
    if (number==0) {
        ((uint32_t *)packet.Data)[0]=0x20001000;
        ((uint32_t *)packet.Data)[1]=mBMS_FW_RN_ADDRESS|1;
    }
}
static uint32_t block_sum(void) {
    return byte_sum(DownloadBlock.bytes,CanDLInfor.WriteLength*8);
}
static uint32_t check_block(uint32_t sum) {
    be32(packet.Data,wire_addr()); be32(packet.Data+4,sum);
    return WriteDataChkSumByCan(&packet);
}
static void block_tests(void) {
    /* A complete 256-frame block may arrive out of order, with exact duplicates. */
    prepare_block(256);
    for (int i=255;i>=0;i--) { fill_frame(i); CHECK(WriteDataByCan(&packet)==0); }
    fill_frame(137); CHECK(WriteDataByCan(&packet)==0);
    CHECK(DownloadPacketCount==256 && write_count==0);
    uint32_t sum=block_sum();
    CHECK(check_block(sum+1)==5 && write_count==0);
    CHECK(check_block(sum)==0 && write_count==1);
    CHECK(check_block(sum)==0 && write_count==1); /* Lost ACK: no reprogramming. */
    verify_image();

    /* Missing frame is distinguishable from a checksum error and never reaches flash. */
    prepare_block(256);
    for (int i=0;i<256;i++) if (i!=137) { fill_frame(i); CHECK(WriteDataByCan(&packet)==0); }
    CHECK(check_block(0)==12 && write_count==0);
    fill_frame(137); CHECK(WriteDataByCan(&packet)==0);
    CHECK(check_block(block_sum())==0);

    prepare_block(1); fill_frame(0); CHECK(WriteDataByCan(&packet)==0);
    packet.Data[0]^=1; CHECK(WriteDataByCan(&packet)==13);
    CHECK(check_block(0)==13 && write_count==0);
    prepare_block(1); fill_frame(1); CHECK(WriteDataByCan(&packet)==12);
    CHECK(check_block(0)==12 && write_count==0);
    prepare_block(1); fill_frame(0); SET_DLC(packet,7);
    CHECK(WriteDataByCan(&packet)==12 && write_count==0);

    for (int mode=2;mode<=4;mode++) {
        prepare_block(1); fill_frame(0); CHECK(WriteDataByCan(&packet)==0);
        sum=block_sum(); fail_mode=mode;
        CHECK(check_block(sum)==14);
        CHECK(!DownloadSessionActive && fwinfo.AppUpdate==0);
        commit_rejected();
    }

    /* Multi-block image with a 3-byte tail: padding belongs to the block, not image sum. */
    prepare_block(256); fwinfo.AppFileSize=2051;
    for (int i=0;i<256;i++) { fill_frame(i); CHECK(WriteDataByCan(&packet)==0); }
    CHECK(check_block(block_sum())==0);
    be32(packet.Data,2051); be32(packet.Data+4,image_sum());
    CHECK(FwImageChkSumByCan(&packet)==9); /* Tail has not been committed. */
    be32(packet.Data,wire_addr()+2048); be32(packet.Data+4,1);
    CHECK(WriteAddressSetByCan(&packet)==0);
    PACKET_ID(packet)=0; memset(packet.Data,0,8);
    packet.Data[0]=7; packet.Data[1]=8; packet.Data[2]=9;
    CHECK(WriteDataByCan(&packet)==0);
    be32(packet.Data,wire_addr()+2048); be32(packet.Data+4,24);
    CHECK(WriteDataChkSumByCan(&packet)==0);
    verify_image();
}
int main(void) {
    uint32_t allocation_base = IMAGE_BASE < 0x08000000 ? 0x10000 : 0x08000000;
    CHECK(VirtualAlloc((void *)allocation_base, IMAGE_END-allocation_base,
        MEM_RESERVE|MEM_COMMIT, PAGE_READWRITE)==(void *)allocation_base);

    reset_test(); commit_rejected();
    fwinfo.AppUpdate=3; commit_rejected(); /* Persisted state is not a live session. */
    reset_test(); start_download(); commit_rejected();
    reset_test(); receive_image(); commit_rejected();
    be32(packet.Data,8); be32(packet.Data+4,image_sum()+1);
    CHECK(FwImageChkSumByCan(&packet)==2); commit_rejected();
    be32(packet.Data,7); be32(packet.Data+4,8);
    CHECK(FwImageChkSumByCan(&packet)==1); commit_rejected();

    reset_test(); receive_image(); verify_image();
    {
        int erases=erase_count, writes=write_count;
        CHECK(SectorEraseProcessByCan(&packet)==9);
        CHECK(WriteAddressSetByCan(&packet)==9);
        CHECK(WriteDataByCan(&packet)==9);
        CHECK(WriteDataChkSumByCan(&packet)==9);
        CHECK(FwImageChkSumByCan(&packet)==9);
        CHECK(erase_count==erases && write_count==writes);
        CHECK(fwinfo.AppUpdate==3 && DownloadImageVerified);
    }
    start_download(); CHECK(!DownloadImageVerified && CanDLInfor.WriteLength==0);
    commit_rejected();
    CHECK(WriteDataByCan(&packet)==9); /* Previous block address was cleared. */

    reset_test(); receive_image(); verify_image();
    memset(packet.Data,0,8); /* Zero length start invalidates verification. */
    CHECK(FwUpdateStartByCan(&packet)==2); commit_rejected();
    reset_test(); receive_image(); verify_image();
    be32(packet.Data+4,0xFFFFFFFF);
    CHECK(FwUpdateStartByCan(&packet)==2); commit_rejected();

    for (int mode=1; mode<=4; mode++) {
        reset_test(); receive_image(); verify_image(); fail_mode=mode;
        packet.Data[0]=3;
        CHECK(UpdateRestartByCan(&packet)==11);
        CHECK(last_response[7]==11 && !reset_count && fwinfo.AppUpdate==3);
        CHECK(DownloadImageVerified && WriteDataByCan(&packet)==9);
        fail_mode=0; packet.Data[0]=3;
        CHECK(UpdateRestartByCan(&packet)==0);
        CHECK(last_response[7]==0 && reset_count==1);
        CHECK(memcmp((void *)INFO_BASE,&fwinfo,sizeof(fwinfo))==0);
        CHECK(fwinfo.AppUpdate==9 && !DownloadSessionActive);
    }
    block_tests();
    queue_tests();
    puts("PASS: update guards, 256-frame RAM staging, missing/duplicate/short frames, flash failures, partial tail, queue overflow/wrap and bounded draining.");
    return 0;
}
"""

QUEUE_STUBS = r"""
#define __DMB() ((void)0)
#define SUCCESS 1
#define ERROR 0
#define CAN_FIFO0 0
#define CAN_IT_FMP0 0
#define CanRxIrqFlg 1
#define CanTxPrimaryIrqFlg 2
#define CanBusErrorIrqFlg 3
#define CanArbiLostIrqFlg 4
#define Int005_IRQn 5
typedef struct {
    volatile uint16_t RxInPos,RxOutPos;
    volatile uint32_t RxOverflowCount;
    QueuePacket RxBuf[RX_BUFFER_SIZE];
    uint32_t can;
} QueueInfo;
static QueueInfo Can;
static QueuePacket hardware_packet, processed_packet;
static uint32_t processed_count;
int CAN_IrqFlgGet(int flag) { return flag==CanRxIrqFlg; }
void CAN_IrqFlgClr(int flag) {}
void CAN_ClearITPendingBit(uint32_t can,int flag) {}
void NVIC_ClearPendingIRQ(int flag) {}
#if IS_V1
void CAN_Receive(uint32_t can,int fifo,QueuePacket *p) { *p=hardware_packet; }
#else
void CAN_Receive(QueuePacket *p) { *p=hardware_packet; }
#endif
uint8_t CAN_GetPacket(QueuePacket *p);
uint8_t GetCanPacket(void) { return CAN_GetPacket(&processed_packet); }
void CanRxProcess(void) { processed_count++; }
"""

QUEUE_CASES = r"""
static void queue_push(uint32_t number) {
    memcpy(hardware_packet.Data,&number,4);
    RECEIVE_IRQ();
}
static void queue_tests(void) {
    QueuePacket p;
    memset(&Can,0,sizeof(Can)); processed_count=0;
    for (uint32_t i=0;i<RX_BUFFER_SIZE-1;i++) queue_push(i);
    queue_push(0xDEADBEEF);
    CHECK(Can.RxOverflowCount==1);
    CanCommProcess(); CHECK(processed_count==32);
    CHECK(*(uint32_t *)processed_packet.Data==31);
    for (uint32_t i=32;i<RX_BUFFER_SIZE-1;i++) {
        CHECK(CAN_GetPacket(&p)==1 && *(uint32_t *)p.Data==i);
    }
    CHECK(CAN_GetPacket(&p)==0);
    for (uint32_t i=0;i<1000;i++) {
        queue_push(i); CHECK(CAN_GetPacket(&p)==1 && *(uint32_t *)p.Data==i);
    }
    memset(&Can,0,sizeof(Can));
    for (uint32_t i=0;i<257;i++) queue_push(i);
    CHECK(Can.RxOverflowCount==0);
    for (uint32_t i=0;i<257;i++) CHECK(CAN_GetPacket(&p)==1 && *(uint32_t *)p.Data==i);
    CHECK(CAN_GetPacket(&p)==0);
}
"""


def main():
    if not VCVARS.exists():
        raise SystemExit("Visual Studio 2022 x86 C compiler environment not found")
    with tempfile.TemporaryDirectory(prefix="mbms_update_test_") as tmp:
        tmp = Path(tmp)
        for name, folder, ext, image_base, image_end, info_base in VARIANTS:
            folder = ROOT / folder
            module = (folder / "can_update_process.c").read_bytes().decode("latin1")
            module = re.sub(r'^#include[^\n]*', '', module, flags=re.M)
            save = function((folder / "device_param.c").read_bytes().decode("latin1"), "SaveFwInfo")
            config = f"#define IMAGE_BASE {image_base}u\n#define IMAGE_END {image_end}u\n#define INFO_BASE {info_base}u\n"
            packet_type = "CanRxMsg" if name == "V1" else "stc_can_rxframe_t"
            dlc = "(p).DLC" if name == "V1" else "(p).Cst.Control_f.DLC"
            cases = f"#define PACKET {packet_type}\n#define WIRE_OFFSET {0 if name == 'V1' else 0x08020000}u\n#define SET_DLC(p,n) ({dlc}=(n))\n#define PACKET_ID(p) ((p).{ext})\nstatic void queue_tests(void);\n" + CASES + QUEUE_CASES
            hardware = folder.parent / "HARDWARE" / "can.c" if name == "V1" else folder.parent / "Hardware" / "can.c"
            hardware_source = hardware.read_bytes().decode("latin1")
            capacity = re.search(r"#define RX_BUFFER_SIZE\s+(\d+)", hardware_source).group(1)
            irq = "CAN1_RX0_IRQHandler" if name == "V1" else "CAN_RxIrqCallBack"
            queue = f"#define RX_BUFFER_SIZE {capacity}\n#define IS_V1 {int(name=='V1')}\n#define RECEIVE_IRQ {irq}\ntypedef {packet_type} QueuePacket;\n"
            queue += QUEUE_STUBS + function(hardware_source, irq) + "\n" + function(hardware_source, "CAN_GetPacket")
            queue += "\n" + function((folder / "can_comm.c").read_bytes().decode("latin1"), "CanCommProcess")
            test_c = tmp / (name + ".c")
            test_c.write_bytes((config + STUBS + module + "\n" + save + "\n" + queue + "\n" + cases).encode("latin1"))
            batch = tmp / "compile.bat"
            batch.write_text(f'@call "{VCVARS}" >nul\n@cl /nologo /TC /std:c11 /Od /W3 "{test_c}" /Fe:"{tmp / (name + ".exe")}" /Fo:"{tmp / (name + ".obj")}"\n', encoding="ascii")
            build = subprocess.run(["cmd", "/d", "/c", str(batch)], cwd=tmp, capture_output=True)
            if build.returncode:
                raise RuntimeError(build.stdout.decode(errors="replace") + build.stderr.decode(errors="replace"))
            result = subprocess.run([str(tmp / (name + ".exe"))], cwd=tmp, capture_output=True)
            if result.returncode:
                raise RuntimeError(name + ": " + result.stdout.decode(errors="replace") + result.stderr.decode(errors="replace"))
            print(name + ": " + result.stdout.decode(errors="replace").splitlines()[-1])


if __name__ == "__main__":
    main()