#!/usr/bin/env python3
"""Real PTY/socket integration tests for the native RTU and TCP server bridge."""
import ctypes as C,os,pty,select,socket,struct,subprocess,time,tempfile
from pathlib import Path
root=Path(__file__).resolve().parents[1]
with tempfile.TemporaryDirectory() as tmp:
 so=Path(tmp)/'test-native.so'
 subprocess.run(['cc','-shared','-fPIC','-std=c11','-Wall','-Wextra','-Werror','-O2','-I'+str(root/'src'),*[str(root/'src'/f) for f in ['atom_mb.c','atom_rtu.c','atom_server.c','atom_bridge.c']],'-o',str(so)],check=True)
 lib=C.CDLL(str(so));fn=lib.ab_call
 U=C.c_uint32;W=C.c_uint16;B=C.c_uint8
 fn.argtypes=[U,U,U,W,B,B,W,W,C.POINTER(W),U,C.POINTER(U),U,C.c_uint64,U,U,U];fn.restype=U
 data=(W*2000)();status=(U*24)()
 def now():return int(time.monotonic()*1000)
 def call(cmd,h=0,ip=0,port=0,unit=1,f=3,addr=0,q=1,ct=100,rt=100,retry=10,when=None):
  return fn(cmd,h,ip,port,unit,f,addr,q,data,2000,status,24,now() if when is None else when,ct,rt,retry)
 def crc(frame):
  v=0xffff
  for byte in frame:
   v^=byte
   for _ in range(8):v=(v>>1)^(0xa001 if v&1 else 0)
  return v
 def adu(payload):return payload+struct.pack('<H',crc(payload))
 def ready(h):
  end=time.monotonic()+1
  while time.monotonic()<end:
   assert call(5,h)==0
   if status[1]==3:return
   time.sleep(.001)
  raise AssertionError(('not ready',list(status)))
 def rtu_scenario(mode,f=3):
  master,slave=pty.openpty();path=os.ttyname(slave);os.close(slave);os.set_blocking(master,False)
  for i,b in enumerate(path.encode()):data[i]=b
  assert call(20,ip=9600,unit=0,f=8,addr=1,q=len(path),ct=0)==0
  h=status[0];assert call(9,h)==0;assert call(3,h)==0;ready(h)
  for attempt in range(2 if mode in ['crc','timeout','echo'] else 1):
   if f==16:data[0]=123;data[1]=456
   if f==15:data[0]=1;data[1]=0
   assert call(6,h,unit=1,f=f,addr=10,q=2)==0
   received=b'';sent=False;fragments=[];deadline=time.monotonic()+1
   while time.monotonic()<deadline:
    assert call(5,h)==0
    if select.select([master],[],[],0)[0]:
     try:received+=os.read(master,256)
     except OSError:pass
    need=(9+received[6]) if len(received)>=7 and f>=15 else 8
    if len(received)>=need and not sent:
     assert crc(received)==0;sent=True
     payload=(bytes([1,f])+received[2:6]) if f>=15 else bytes([1,f,4,0,42,0,43])
     if mode=='exception':payload=bytes([1,f|0x80,2])
     if mode=='echo' and attempt==0:payload=payload[:3]+bytes([payload[3]^1])+payload[4:]
     response=adu(payload)
     if mode=='crc' and attempt==0:response=response[:-1]+bytes([response[-1]^1])
     if mode!='timeout' or attempt>0:
      fragments=[response[:2],response[2:4],response[4:]] if mode=='fragment' else [response]
    if fragments:os.write(master,fragments.pop(0))
    if status[3] or status[4]:break
    time.sleep(.002)
   assert sent and (status[3] or status[4]),list(status)
   if mode in ['crc','echo'] and attempt==0:assert status[5]==8
   elif mode=='timeout' and attempt==0:assert status[5]==6
   elif mode=='exception':assert status[5]==9 and status[8]==2 and status[1]==3
   else:
    assert status[3]==1,list(status)
    if f<=4:assert call(7,h)==0 and list(data[:2])==[42,43]
   assert call(10,h)==0
   if attempt==0 and mode in ['crc','timeout','echo']:ready(h)
  assert call(6,h,unit=0,f=3,addr=0,q=1)==1 # Broadcast unsupported, no fake completion.
  assert call(2,h)==0;assert call(5,h)==256;os.close(master)
 for mode in ['normal','fragment','crc','timeout','exception']:rtu_scenario(mode);print('PASS RTU',mode)
 rtu_scenario('normal',16);print('PASS RTU register write')
 rtu_scenario('echo',16);print('PASS RTU invalid echo and reopen recovery')
 rtu_scenario('normal',15);print('PASS RTU coil write')
 # Native TCP server, four independent tables mapped at nonzero Modbus address.
 probe=socket.socket();probe.bind(('127.0.0.1',0));port=probe.getsockname()[1];probe.close()
 assert call(30,ip=0x7f000001,port=port)==0;h=status[0]
 for table in range(4):
  for i in range(16):data[i]=i%2 if table<2 else 100+i
  assert call(35,h,f=table,addr=10,q=16)==0
 assert call(32,h)==0;assert call(34,h)==0 and status[1]==1
 peer=socket.create_connection(('127.0.0.1',port));peer.setblocking(False)
 tid=0
 def exchange(pdu):
  global tid
  tid+=1;packet=struct.pack('>HHHB',tid,0,len(pdu)+1,1)+pdu
  peer.send(packet[:3]);assert call(34,h)==0;peer.send(packet[3:])
  reply=b'';end=time.monotonic()+1
  while time.monotonic()<end:
   assert call(34,h)==0
   try:reply+=peer.recv(260)
   except BlockingIOError:pass
   if len(reply)>=7 and len(reply)==6+struct.unpack('>H',reply[4:6])[0]:return reply[7:]
   time.sleep(.001)
  raise AssertionError(('server timeout',pdu,reply))
 for f in [1,2,3,4]:
  rep=exchange(struct.pack('>BHH',f,10,2));assert rep==(bytes([f,1,2]) if f<=2 else bytes([f,4,0,100,0,101]));print('PASS slave read FC',f)
 for f in [5,6]:
  value=0xff00 if f==5 else 456;pdu=struct.pack('>BHH',f,10,value);assert exchange(pdu)==pdu
  assert call(36,h,f=0 if f==5 else 2,addr=10,q=16)==0 and data[0]==(1 if f==5 else 456);print('PASS slave write FC',f)
 for f in [15,16]:
  payload=b'\x01' if f==15 else struct.pack('>HH',321,654)
  pdu=struct.pack('>BHHB',f,10,2,len(payload))+payload
  assert exchange(pdu)==struct.pack('>BHH',f,10,2)
  assert call(36,h,f=0 if f==15 else 2,addr=10,q=16)==0 and list(data[:2])==([1,0] if f==15 else [321,654]);print('PASS slave write FC',f)
 assert exchange(struct.pack('>BHH',3,0,1))==bytes([0x83,2])
 assert exchange(struct.pack('>BHH',16,10,0))==bytes([0x90,3])
 assert exchange(struct.pack('>BHH',99,10,1))==bytes([0xe3,1])
 print('PASS slave exceptions and nonzero table bases')
 assert call(31,h)==0;assert call(34,h)==256;peer.close()
 # Listener bind error recovers after the conflicting listener disappears.
 blocker=socket.socket();blocker.setsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR,1);blocker.bind(('127.0.0.1',port));blocker.listen()
 assert call(30,ip=0x7f000001,port=port)==0;h=status[0];call(32,h);call(34,h);assert status[2]==1
 blocker.close();call(34,h,when=now()+100);assert status[1]==1 and status[2]==0
 call(31,h);print('PASS slave listener recovery and stale handles')
 lib.ab_shutdown()