import json
import socket
import ssl
import struct
import time
import urllib.request
import uuid
from pathlib import Path

import argparse
parser = argparse.ArgumentParser(description='Independent test-device/backend receipt ACK checks; does not run the Go simulator.')
parser.add_argument('--credentials', required=True, help='Private JSON handoff or ESP32 connection.txt file')
parser.add_argument('--ca', required=True, help='Test broker public CA file')
parser.add_argument('--host', default='109.122.252.129')
parser.add_argument('--plaintext-port', type=int, default=1884)
parser.add_argument('--tls-port', type=int, default=8884)
args = parser.parse_args()
HOST = args.host
private = Path(args.credentials).read_text()
if args.credentials.endswith('.json'):
    connection = json.loads(private)
    device_user, device_password = connection['username'], connection['password']
else:
    connection = dict(line.split('=', 1) for line in private.splitlines() if '=' in line)
    device_user, device_password = connection['mqtt_user'].strip(), connection['mqtt_pass'].strip()

def varint(n):
    out = bytearray()
    while True:
        byte = n % 128
        n //= 128
        out.append(byte | (128 if n else 0))
        if not n:
            return bytes(out)

def string(s):
    raw = s.encode()
    return struct.pack('!H', len(raw)) + raw

def exact(sock, n):
    result = b''
    while len(result) < n:
        chunk = sock.recv(n - len(result))
        if not chunk:
            raise RuntimeError('broker closed connection')
        result += chunk
    return result

def read(sock):
    header = exact(sock, 1)[0]
    length = 0
    factor = 1
    while True:
        digit = exact(sock, 1)[0]
        length += (digit & 127) * factor
        if not digit & 128:
            break
        factor *= 128
    return header, exact(sock, length)

def send(sock, header, body):
    sock.sendall(bytes([header]) + varint(len(body)) + body)

def connect(version, tls, username, password, client_id=None, persistent=False, expect_session=False, anonymous=False):
    sock = socket.create_connection((HOST, args.tls_port if tls else args.plaintext_port), timeout=8)
    if tls:
        context = ssl.create_default_context(cafile=args.ca)
        sock = context.wrap_socket(sock, server_hostname=HOST)
    body = string('MQTT') + bytes([version, 2 if anonymous else (0xc0 if persistent else 0xc2)]) + struct.pack('!H', 30)
    if version == 5:
        body += b'\x05\x11' + struct.pack('!I', 30) if persistent else b'\x00'
    body += string(client_id or 'rm5-independent-' + uuid.uuid4().hex[:12]) + (b'' if anonymous else string(username) + string(password))
    send(sock, 0x10, body)
    header, body = read(sock)
    if header >> 4 != 2 or body[1] != 0:
        raise RuntimeError('authentication/CONNECT rejected: code ' + str(body[1]))
    if expect_session and not body[0] & 1:
        raise RuntimeError('persistent session absent; retransmission test inconclusive')
    return sock

def subscribe(sock, version, topic):
    body = b'\x00\x01' + (b'\x00' if version == 5 else b'') + string(topic) + b'\x01'
    send(sock, 0x82, body)
    header, body = read(sock)
    if header >> 4 != 9 or body[-1] >= 128:
        raise RuntimeError('backend subscription rejected')

def publish(sock, version, topic, payload):
    body = string(topic) + b'\x00\x02' + (b'\x00' if version == 5 else b'') + payload.encode()
    send(sock, 0x32, body)
    header, body = read(sock)
    if header >> 4 != 4 or (version == 5 and len(body) > 2 and body[2] >= 128):
        raise RuntimeError('publication denied by broker')

def received(sock, version):
    header, body = read(sock)
    if header >> 4 != 3:
        raise RuntimeError('expected subscribed publication')
    topic_len = int.from_bytes(body[:2], 'big')
    position = 2 + topic_len
    if (header >> 1) & 3:
        packet_id = body[position:position+2]
        position += 2
        send(sock, 0x40, packet_id)
    if version == 5:
        size = 0
        factor = 1
        while True:
            digit = body[position]
            position += 1
            size += (digit & 127) * factor
            if not digit & 128:
                break
            factor *= 128
        position += size
    return body[position:].decode()

def viewer_has(topic, payload, kind):
    until = time.monotonic() + 5
    while time.monotonic() < until:
        with urllib.request.urlopen('http://127.0.0.1:8080/api/snapshot', timeout=3) as response:
            snapshot = json.load(response)
        if any(event.get('topic') == topic and event.get('kind') == kind and event.get('payload', {}).get('text', '') == payload for event in snapshot['events']):
            return
        time.sleep(0.1)
    raise RuntimeError('broker delivered but viewer did not show ' + kind)


import hashlib
for tls in (False, True):
    for version in (4,5):
        device='ack-test-'+uuid.uuid4().hex[:8]
        source='robot-market/test/v1/'+device+'/telemetry'
        acktopic='robot-market/test/v1/'+device+'/acks'
        reader=connect(version,tls,device_user,device_password)
        writer=connect(version,tls,device_user,device_password)
        try:
            subscribe(reader,version,acktopic)
            for payload in ['hello from device', json.dumps({'message_id':'receipt-001','device_id':'anything','timestamp':'{{$isoTimestamp}}'}), '']:
                publish(writer,version,source,payload)
                ack_payload=received(reader,version)
                ack=json.loads(ack_payload)
                viewer_has(source,payload,'message_received')
                viewer_has(acktopic,ack_payload,'receipt_ack_sent')
                assert ack['type']=='receipt_ack' and ack['ok'] is True
                assert ack['source_topic']==source and ack['payload_bytes']==len(payload.encode())
                assert ack['payload_sha256']==hashlib.sha256(payload.encode()).hexdigest()
                if 'receipt-001' in payload: assert ack['message_id']=='receipt-001'
            reader.settimeout(.3)
            try:
                extra=read(reader)
            except socket.timeout:
                pass
            else:
                raise RuntimeError('Unexpected reply loop or extra ACK: '+str(extra[0]))
            print(('TLS' if tls else 'plaintext')+' MQTT '+('5' if version==5 else '3.1.1')+': device receives backend ACK for text, arbitrary JSON, empty payload; no loop PASS')
        finally:
            for sock in (writer,reader):
                send(sock,0xe0,b'')
                sock.close()
# Native MQTT 5 custom response topic and binary correlation echo.
source='robot-market/test/v1/native-ack-test/telemetry'
response='robot-market/test/v1/native-ack-test/custom-reply'
corr=bytes([0,255,1,2])
reader=connect(5,False,device_user,device_password)
writer=connect(5,False,device_user,device_password)
try:
    subscribe(reader,5,response)
    props=b'\x08'+string(response)+b'\x09'+struct.pack('!H',len(corr))+corr
    body=string(source)+b'\x00\x02'+varint(len(props))+props+b'native properties test'
    send(writer,0x32,body)
    header,body=read(writer)
    assert header>>4==4 and (len(body)==2 or body[2]<128)
    header,body=read(reader)
    assert header>>4==3
    pos=2+int.from_bytes(body[:2],'big')
    packetid=body[pos:pos+2];pos+=2
    send(reader,0x40,packetid)
    size=0;factor=1
    while True:
        digit=body[pos];pos+=1;size+=(digit&127)*factor
        if not digit&128:break
        factor*=128
    receivedprops=body[pos:pos+size]
    assert b'\x09'+struct.pack('!H',len(corr))+corr in receivedprops
    ack=json.loads(body[pos+size:])
    assert ack['source_topic']==source and ack['ok'] is True
    print('MQTT 5 Response Topic and exact binary Correlation Data echo: PASS')
finally:
    for sock in (writer,reader):
        send(sock,0xe0,b'')
        sock.close()

# Controlled wire-level retransmission: same persistent client/session and packet ID,
# original PUBACK deliberately unread, then DUP=1 after a dropped socket.
source = 'robot-market/test/v1/wire-dup-' + uuid.uuid4().hex[:8] + '/telemetry'
acktopic = source.rsplit('/', 1)[0] + '/acks'
reader = connect(5, False, device_user, device_password)
client_id = 'rm5-dup-' + uuid.uuid4().hex[:12]
writer = connect(5, False, device_user, device_password, client_id, True)
try:
    subscribe(reader, 5, acktopic)
    body = string(source) + b'\x00\x07\x00' + b'controlled wire retransmission'
    send(writer, 0x32, body)
    first = json.loads(received(reader, 5))
    writer.close()  # no PUBACK read, no DISCONNECT
    time.sleep(.2)
    writer = connect(5, False, device_user, device_password, client_id, True, True)
    send(writer, 0x3a, body)  # QoS 1, DUP set, identical packet identifier and payload
    header, response = read(writer)
    assert header >> 4 == 4 and response[:2] == b'\x00\x07'
    assert len(response) == 2 or response[2] < 128
    second = json.loads(received(reader, 5))
    assert first['payload_sha256'] == second['payload_sha256']
    print('Controlled MQTT 5 wire QoS 1 DUP retransmission after unread PUBACK: PASS; two receipt ACKs, no device action claimed')
finally:
    for sock in (writer, reader):
        send(sock, 0xe0, b'')
        sock.close()

for label, username, password, anonymous in [('Anonymous', '', '', True), ('Invalid password', device_user, uuid.uuid4().hex, False)]:
    try:
        sock = connect(5, True, username, password, anonymous=anonymous)
    except RuntimeError as error:
        assert 'CONNECT rejected' in str(error), 'failure was not broker authentication rejection'
        print(label + ' authentication rejection: PASS')
    else:
        sock.close()
        raise RuntimeError(label + ' authentication unexpectedly accepted')
for label, context, hostname in [('Untrusted CA', ssl.create_default_context(), HOST), ('Wrong hostname', ssl.create_default_context(cafile=args.ca), 'wrong.invalid')]:
    raw = socket.create_connection((HOST, args.tls_port), timeout=8)
    try:
        secured = context.wrap_socket(raw, server_hostname=hostname)
    except ssl.SSLCertVerificationError:
        print(label + ' TLS rejection: PASS')
    else:
        secured.close()
        raise RuntimeError(label + ' TLS unexpectedly accepted')
    finally:
        raw.close()