import socket
import struct
import tarfile
from pathlib import Path

import pytest

from go1_pro_sdk import Blowfish, parse_low_state


ROOT = Path(__file__).resolve().parents[1]
CAPTURES = ROOT / "data/captures"
STATE_FILE = ROOT / "go1_pro_sdk/_data/blowfish_state.bin"

if not CAPTURES.is_dir():
    pytest.skip("private data/captures fixtures are not available", allow_module_level=True)


@pytest.fixture(scope="module")
def blowfish():
    return Blowfish.from_state_file(str(STATE_FILE))


def _pcap_udp_packets(path):
    packets = []
    with path.open("rb") as stream:
        header = stream.read(24)
        assert len(header) == 24
        magic = header[:4]
        assert magic in (b"\xd4\xc3\xb2\xa1", b"\xa1\xb2\xc3\xd4")
        endian = "<" if magic == b"\xd4\xc3\xb2\xa1" else ">"
        linktype = struct.unpack(endian + "I", header[20:24])[0]
        assert linktype == 113  # Linux cooked capture v1

        while True:
            packet_header = stream.read(16)
            if not packet_header:
                break
            assert len(packet_header) == 16
            _, _, captured_length, _ = struct.unpack(endian + "IIII", packet_header)
            frame = stream.read(captured_length)
            assert len(frame) == captured_length
            assert struct.unpack(">H", frame[14:16])[0] == 0x0800
            ip_offset = 16
            ip_header_length = (frame[ip_offset] & 0x0F) * 4
            assert frame[ip_offset + 9] == 17
            udp_offset = ip_offset + ip_header_length
            source_port, target_port, udp_length, _ = struct.unpack(
                ">HHHH", frame[udp_offset:udp_offset + 8]
            )
            payload = frame[udp_offset + 8:udp_offset + udp_length]
            packets.append((
                socket.inet_ntoa(frame[ip_offset + 12:ip_offset + 16]),
                socket.inet_ntoa(frame[ip_offset + 16:ip_offset + 20]),
                source_port,
                target_port,
                payload,
            ))
    return packets


def test_real_lowcmd_full_cipher_plain_pair(blowfish):
    cipher = (CAPTURES / "real_lowcmd.bin").read_bytes()
    plain = (CAPTURES / "real_lowcmd_decrypted.bin").read_bytes()
    assert len(cipher) == len(plain) == 616
    assert blowfish.decrypt_ecb(cipher) == plain
    assert blowfish.encrypt_ecb(plain) == cipher
    assert plain[:22] == bytes.fromhex(
        "feefff00000000000000000000000000000000003ac0"
    )
    assert plain[22] == 0x0A  # This fixture contains active servo commands.


def test_real_lowcmd_pcap_matches_extracted_first_frame():
    packets = _pcap_udp_packets(CAPTURES / "real_lowcmd.pcap")
    assert len(packets) == 1406
    assert {(src, dst, sport, dport) for src, dst, sport, dport, _ in packets} == {
        ("192.168.123.161", "192.168.123.10", 8008, 8007)
    }
    assert {len(payload) for *_, payload in packets} == {616}
    assert packets[0][-1] == (CAPTURES / "real_lowcmd.bin").read_bytes()


@pytest.mark.parametrize(
    "filename,soc,temperatures",
    [
        ("mcu_response.bin", 76, [52, 42, 40]),
        ("mcu_response_new.bin", 19, [79, 50, 51]),
    ],
)
def test_real_lowstate_samples(blowfish, filename, soc, temperatures):
    cipher = (CAPTURES / filename).read_bytes()
    assert len(cipher) == 858
    plain = blowfish.decrypt_ecb(cipher[:856])
    assert plain[:4] == b"\xfe\xef\xff\x00"
    state = parse_low_state(plain)
    assert state.bms.SOC == soc
    assert [motor.temperature for motor in state.motorState[:3]] == temperatures
    assert state.wirelessRemote[:2] == b"\x55\x51"


def test_memory_dump_reproduces_packaged_blowfish_state():
    with tarfile.open(CAPTURES / "blowfish_dump.tar.gz", "r:gz") as archive:
        member = archive.extractfile("./rw_7f86a1f000.bin")
        assert member is not None
        memory = member.read()
    extracted = memory[0x232190:0x232190 + 4168]
    assert extracted == STATE_FILE.read_bytes()