import json, os, re, struct, subprocess, threading, time
from pathlib import Path
import numpy as np
import pyautogui
from pywinauto import Desktop

ROOT=Path(__file__).resolve().parent
PROJECT=ROOT.parents[1]
DLL=PROJECT/'tmp/can1-test-20261002/pcan-basic/x64'
os.environ['PATH']=str(DLL)+os.pathsep+os.environ['PATH']
handle=os.add_dll_directory(str(DLL))
import can
CLI=r'D:\STM32CubeCLT_1.19.0\STM32CubeProgrammer\bin\STM32_Programmer_CLI.exe'
NM=r'D:\STM32CubeCLT_1.19.0\GNU-tools-for-STM32\bin\arm-none-eabi-nm.exe'
nm=subprocess.check_output([NM,'-n',str(PROJECT/'firmware/build/Release/AMI_Controller.elf')]).decode().replace('\r','')
symbols={n:int(a,16) for a,k,n in re.findall(r'^(\w+)\s+(\w)\s+(\w+)$',nm,re.M)}
(ROOT/'symbols.json').write_text(json.dumps(symbols,indent=2),encoding='utf-8')
names=['sdc_gpio_state','sdc_gpio_changed','sdc_gpio_timestamp_ms','display_state','can_port_last_rx_ms','can_rx_error_count','uart_rx_error_count']
base=min(symbols[n] for n in names)&~3
length=((symbols['display_state']+56-base+3)//4)*4
MAP=[0,5,2,6,7,1,3,4,8,9,10,11,12,13,14]
ROIS={0:[(15,82,165,118)],1:[(374,82,525,118),(555,146,706,185)],2:[(555,82,706,118)],3:[(374,146,525,185)],4:[(195,146,345,185)],5:[(15,241,225,272)],6:[(255,241,465,272)],7:[(495,241,705,272)],8:[(15,311,165,341)],9:[(195,311,345,341)],10:[(375,311,525,341)],11:[(555,311,705,341)],12:[(15,375,165,405)],13:[(195,375,345,405)],14:[(375,375,525,405)],15:[(555,375,705,405)]}
window=Desktop(backend='win32').window(handle=2034542)
window.set_focus()
canvas=next(c.rectangle() for c in window.children() if c.rectangle().width()==720 and c.rectangle().height()==408)
rows=[]
tx=[]
errors=[]
heartbeat=True
stop=threading.Event()

def memory(name):
    result=subprocess.run([CLI,'-c','port=SWD','mode=HOTPLUG','freq=1000','-r32',hex(base),str(length)],capture_output=True,timeout=15)
    (ROOT/f'{name}_ram.log').write_bytes(result.stdout+result.stderr)
    assert result.returncode==0
    mem={}
    for address,words in re.findall(r'0x([0-9A-Fa-f]+)\s*:\s*((?:[0-9A-Fa-f]{8}\s*)+)',result.stdout.decode(errors='replace')):
        addr=int(address,16)
        for word in words.split():
            for i,b in enumerate(struct.pack('<I',int(word,16))): mem[addr+i]=b
            addr+=4
    def get(n,fmt): return struct.unpack(fmt,bytes(mem[symbols[n]+i] for i in range(struct.calcsize(fmt))))
    return {'state':get('sdc_gpio_state','<H')[0],'changed':get('sdc_gpio_changed','<H')[0],'timestamp':get('sdc_gpio_timestamp_ms','<I')[0],
        'debug_state':[mem[symbols['display_state']+33+i] for i in range(16)],'port_rx_ms':list(get('can_port_last_rx_ms','<2I')),
        'can_errors':get('can_rx_error_count','<I')[0],'uart_errors':get('uart_rx_error_count','<I')[0]}

def expected(levels):
    states=[0]*16
    for bit,index in enumerate(MAP): states[index]=2 if levels&(1<<bit) else 1
    return states

with can.Bus(interface='pcan',channel='PCAN_USBBUS1',bitrate=500000) as bus:
    assert bus.reset()
    def sender():
        sequence=0
        try:
            while not stop.is_set():
                if heartbeat:
                    data=struct.pack('<BBHI',0xa1,1,sequence&0xffff,sequence*100)
                    bus.send(can.Message(arbitration_id=0x6a0,data=data,is_extended_id=False),timeout=.2)
                    tx.append({'id':'0x6A0','data':data.hex(' ')})
                    sequence+=1
                stop.wait(.1)
        except Exception as exc: errors.append(repr(exc))
    thread=threading.Thread(target=sender)
    thread.start()
    def check(name,levels,changed,stamp,states=None,extended=False,dlc=8,raw=None,delay=.55):
        if levels is not None:
            data=struct.pack('<HHI',levels,changed,stamp)[:dlc]
            bus.send(can.Message(arbitration_id=0x6a6,data=data,is_extended_id=extended),timeout=.2)
            tx.append({'id':'0x6A6','data':data.hex(' '),'extended':extended})
        time.sleep(delay)
        values=memory(name)
        target=expected(levels) if states is None else states
        assert values['debug_state']==target,(name,values,target)
        if raw is not None: assert (values['state'],values['changed'],values['timestamp'])==raw,(name,values,raw)
        elif levels is not None: assert (values['state'],values['changed'],values['timestamp'])==(levels&0x7fff,changed&0x7fff,stamp),(name,values)
        assert values['can_errors']==values['uart_errors']==0 and not errors,(name,values,errors)
        assert bus.status()==0,(name,hex(bus.status()))
        shot=pyautogui.screenshot(region=(canvas.left,canvas.top,canvas.width(),canvas.height()))
        shot.save(str(ROOT/f'{name}.png'))
        pixels=np.asarray(shot)
        matching=[]
        for index,state in enumerate(target):
            for x1,y1,x2,y2 in ROIS[index]:
                a=pixels[y1:y2,x1:x2,:3].astype(np.int16)
                r,g,b=a[:,:,0],a[:,:,1],a[:,:,2]
                counts=[int(((r>75)&(r<180)&(abs(r-g)<8)&(abs(r-b)<8)).sum()),int(((g>180)&(r<80)&(b<80)).sum()),int(((r>180)&(g<80)&(b<80)).sum())]
                matching.append(counts[state]>20 and (state==1 or counts[1]==0) and (state==2 or counts[2]==0))
        rows.append({'name':name,'expected':target,'snapshot':values,'hmi_color_match':all(matching),'screenshot':f'{name}.png','pcan_status':hex(bus.status())})
        (ROOT/'mapping-results.json').write_text(json.dumps({'map':MAP,'normal_level':0,'checks':rows,'tx':tx},indent=2),encoding='utf-8')
        print(json.dumps({'name':name,'mcu_passed':True,'hmi_color_match':all(matching)}),flush=True)
    try:
        check('heartbeat_before_sample',None,0,0,[0]*16)
        check('all_normal',0,0x7fff,1000)
        for bit in range(15):
            check(f'PD{bit}_high_fault',1<<bit,1<<bit,2000+2*bit)
            check(f'PD{bit}_low_normal',0,1<<bit,2001+2*bit)
        check('all_fault',0x7fff,0x7fff,4000)
        check('unused_bit15',0x8001,0x8002,5000)
        check('truncated',0xffff,0xffff,6000,expected(1),dlc=7,raw=(1,2,5000))
        check('wrong_extended',0xffff,0xffff,7000,expected(1),extended=True,raw=(1,2,5000))
        check('heartbeat_holds_sample',None,0,0,expected(1),raw=(1,2,5000),delay=1.2)
        heartbeat=False
        check('source_timeout',None,0,0,[0]*16,delay=.8)
        heartbeat=True
        check('heartbeat_alone_after_loss',None,0,0,[0]*16,delay=.8)
        check('fresh_sample_recovery',0,0x7fff,8000)
    finally:
        stop.set()
        thread.join()
    check('final_timeout',None,0,0,[0]*16,delay=.8)
    ui=Desktop(backend='uia').window(handle=2034542)
    counters=[{'id':c.element_info.automation_id,'text':c.window_text()} for c in ui.descendants(control_type='Text') if c.element_info.automation_id.startswith('datalabel')]
    (ROOT/'mapping-results.json').write_text(json.dumps({'map':MAP,'normal_level':0,'checks':rows,'tx':tx,'counters':counters,'sender_errors':errors,'bitrate':500000,'port':'CAN1','mcu_passed':True,'hmi_transitions_passed':all(r['hmi_color_match'] for r in rows)},indent=2),encoding='utf-8')
    print(json.dumps({'mcu_passed':True,'checks':len(rows),'hmi_transitions_passed':all(r['hmi_color_match'] for r in rows),'counters':counters}),flush=True)