from pathlib import Path
from hashlib import sha256

import pytest

from go1_pro_sdk import LowCmd, MCUClient
from go1_pro_sdk.connection import mcu_client as mcu_module


ROOT = Path(__file__).resolve().parents[1]
CAPTURE_FILE = ROOT / "data/captures/mcu_response_new.bin"
if not CAPTURE_FILE.is_file():
    pytest.skip("private LowState capture is not available", allow_module_level=True)
CAPTURE = CAPTURE_FILE.read_bytes()
DAMPING_SHA256 = "ca911b04a09d8e8b069c96ac31e410b683fcf1bff625dd203f9efde215888400"


class FakeSocket:
    def __init__(self):
        self.sent = []
        self.recv_batches = []
        self.closed = False
        self.bound = None
        self.blocking = None

    def setblocking(self, value):
        self.blocking = value

    def setsockopt(self, *args):
        pass

    def bind(self, address):
        self.bound = address

    def getsockname(self):
        return ("0.0.0.0", 45678)

    def sendto(self, data, address):
        self.sent.append((bytes(data), address))
        return len(data)

    def recvfrom(self, size):
        if not self.recv_batches:
            raise BlockingIOError
        batch = self.recv_batches[0]
        if not batch:
            self.recv_batches.pop(0)
            raise BlockingIOError
        return batch.pop(0)[:size], ("192.168.123.10", 8007)

    def close(self):
        self.closed = True


@pytest.fixture
def fake_socket(monkeypatch):
    sock = FakeSocket()
    monkeypatch.setattr(mcu_module.socket, "socket", lambda *args, **kwargs: sock)
    monkeypatch.setattr(mcu_module.time, "sleep", lambda duration: None)
    return sock


def test_native_client_send_receive_and_close(fake_socket):
    with MCUClient() as client:
        assert client.local_port == 45678
        assert client.send(LowCmd().all_damping()) == 616
        packet, address = fake_socket.sent[-1]
        assert address == ("192.168.123.10", 8007)
        assert sha256(packet).hexdigest() == DAMPING_SHA256

        fake_socket.recv_batches = [[CAPTURE]]
        state = client.recv_latest()
        assert state is client.last_state
        assert [motor.temperature for motor in state.motorState[:3]] == [79, 50, 51]
        assert state.wirelessRemote[:2] == b"\x55\x51"

    assert fake_socket.closed
    assert client.sock is None


def test_native_client_ignores_invalid_state(fake_socket):
    client = MCUClient()
    fake_socket.recv_batches = [[b"\x00" * 858]]
    assert client.recv_latest() is None
    assert client.last_state is None


def test_native_wake_and_safe_stop(fake_socket):
    client = MCUClient()
    fake_socket.recv_batches = [[CAPTURE], [CAPTURE], [CAPTURE]]
    assert client.wake_mcu(n_frames=3, dt=0) == 3
    assert len(fake_socket.sent) == 3
    assert all(sha256(packet).hexdigest() == DAMPING_SHA256 for packet, _ in fake_socket.sent)

    client.safe_stop(n_frames=2, dt=0)
    assert len(fake_socket.sent) == 5
    assert all(sha256(packet).hexdigest() == DAMPING_SHA256 for packet, _ in fake_socket.sent)