import os

import pytest

from fast_lowcmd import FastLowCmdBuilder, default_state_path
from go1_pro_sdk import Blowfish, LowCmd, MotorCmd, MotorMode, PowerProtectViolation
from go1_pro_sdk import apply_safety
from go1_pro_sdk.codec.lowcmd_builder import build_low_cmd_encrypted, build_low_cmd_plain


@pytest.fixture
def builder():
    return FastLowCmdBuilder(default_state_path())


@pytest.fixture
def bf():
    return Blowfish.from_state_file(default_state_path())


def test_damping_plain_matches_python(builder):
    assert builder.build_plain_damping() == build_low_cmd_plain(LowCmd())


def test_damping_encrypted_matches_python(builder, bf):
    cmd = LowCmd()
    assert builder.build_encrypted_damping() == build_low_cmd_encrypted(cmd, bf)


def test_servo12_plain_matches_python(builder):
    q = [0.1 * i - 0.5 for i in range(12)]
    dq = [0.2 * i for i in range(12)]
    tau = [0.0 for _ in range(12)]
    kp = 36.0
    kd = 1.0

    cmd = LowCmd()
    for i in range(12):
        cmd.set_motor(i, MotorCmd(
            mode=MotorMode.Servo,
            q=q[i],
            dq=dq[i],
            tau=tau[i],
            Kp=kp,
            Kd=kd,
        ))

    assert builder.build_plain_servo12(q, dq=dq, tau=tau, kp=kp, kd=kd) == build_low_cmd_plain(cmd)


def test_servo12_encrypted_matches_python(builder, bf):
    q = [-0.1, 0.8, -1.5, 0.1, 0.8, -1.5, -0.1, 1.0, -1.5, 0.1, 1.0, -1.5]
    cmd = LowCmd()
    for i in range(12):
        cmd.set_motor(i, MotorCmd(mode=MotorMode.Servo, q=q[i], Kp=28.0, Kd=0.7))

    assert builder.build_encrypted_servo12(q, kp=28.0, kd=0.7) == build_low_cmd_encrypted(cmd, bf)


def test_servo12_checked_matches_python_safety(builder, bf):
    q = [-1.0, 0.8, -1.5, 0.1, 0.8, -1.5, -0.1, 1.0, -1.5, 0.1, 1.0, -1.5]
    actual_q = [0.0] * 12
    actual_tau = [0.0] * 12

    class Motor:
        def __init__(self, q_value, tau_value):
            self.q = q_value
            self.tauEst = tau_value

    class State:
        pass

    state = State()
    state.motorState = [Motor(actual_q[i], actual_tau[i]) for i in range(12)]

    cmd = LowCmd()
    for i in range(12):
        cmd.set_motor(i, MotorCmd(mode=MotorMode.Servo, q=q[i], Kp=36.0, Kd=1.0))
    apply_safety(cmd, state, power_factor=9, position_limit_on=True, position_protect_limit=0.5)

    assert builder.build_encrypted_servo12_checked(
        q,
        kp=36.0,
        kd=1.0,
        actual_q=actual_q,
        actual_tau=actual_tau,
        position_protect_limit=0.5,
    ) == build_low_cmd_encrypted(cmd, bf)


def test_servo12_checked_power_protect(builder):
    q = [-0.1, 0.8, -1.5, 0.1, 0.8, -1.5, -0.1, 1.0, -1.5, 0.1, 1.0, -1.5]
    actual_tau = [0.0] * 12
    actual_tau[4] = -24.02
    with pytest.raises(PowerProtectViolation):
        builder.build_encrypted_servo12_checked(q, kp=36.0, kd=1.0, actual_tau=actual_tau)


def test_fields20_plain_matches_python(builder):
    mode = [int(MotorMode.Servo) if i < 12 else int(MotorMode.Damping) for i in range(20)]
    q = [0.05 * i for i in range(20)]
    dq = [0.01 * i for i in range(20)]
    tau = [0.0 for _ in range(20)]
    kp = [20.0 if i < 12 else 0.0 for i in range(20)]
    kd = [1.0 if i < 12 else 0.0 for i in range(20)]

    cmd = LowCmd()
    for i in range(20):
        cmd.set_motor(i, MotorCmd(mode=mode[i], q=q[i], dq=dq[i], tau=tau[i], Kp=kp[i], Kd=kd[i]))

    assert builder.build_plain_fields20(mode=mode, q=q, dq=dq, tau=tau, kp=kp, kd=kd) == build_low_cmd_plain(cmd)


def test_lowcmd_object_matches_python(builder, bf):
    cmd = LowCmd()
    cmd.frameReserve = 3
    cmd.wirelessRemote = bytes(range(40))
    cmd.reserve = b"\x01\x02\x03\x04"
    for i in range(20):
        cmd.set_motor(
            i,
            MotorCmd(
                mode=MotorMode.Servo if i < 12 else MotorMode.Damping,
                q=0.02 * i,
                dq=0.01 * i,
                tau=0.0,
                Kp=12.0 if i < 12 else 0.0,
                Kd=1.0 if i < 12 else 0.0,
                reserve=[i, i + 1, i + 2],
            ),
        )

    assert builder.build_plain_lowcmd(cmd) == build_low_cmd_plain(cmd)
    assert builder.build_encrypted_lowcmd(cmd) == build_low_cmd_encrypted(cmd, bf)


def test_state_file_is_required():
    with pytest.raises(FileNotFoundError):
        FastLowCmdBuilder(os.path.join(os.path.dirname(__file__), "missing_state.bin"))