# originally written by Delf Byonoy
# https://github.com/byonoy-dne/report-parsing/blob/main/parser.py

import re
import struct

# Map report_id to the corresponding payload structure
PAYLOAD_FORMATS = {
  0x0010: (
    "<BB29H",
    "SUPPORTED_REPORTS_IN",
    ["seq", "seq_len", "ids"],
  ),  # supported_reports_in_t
  0x0050: ("<I", "API_VERSION_IN", ["version_no"]),  # api_version_in_t
  0x0020: (
    "<HH56s",
    "ACK_INOUT",
    ["report_id", "return_code", "error_description"],
  ),  # ack_inout_t
  0x0500: (
    "<BBhhI12fBB",
    "ABS96_MEASUREMENT_IN",
    [
      "seq",
      "seq_len",
      "signal_wavelength_nm",
      "reference_wavelength_nm",
      "duration_ms",
      "data",
      "flags",
      "progress",
    ],
  ),  # abs96_measurement_in_t
  0x0300: (
    "<BBBIBB",
    "STATUS_IN",
    [
      "is_initialized",
      "is_slot_occupied",
      "error_code",
      "time_since_boot_s",
      "is_measurement_in_progress",
      "boot_completed",
    ],
  ),  # status_in_t
  0x0070: (
    "<BII50s",
    "DIAGNOSTIC_DATA_IN",
    ["level", "offset", "data_length", "block"],
  ),  # diagnostic_data_in_t
  0x0030: (
    "<II52s",
    "BULK_DATA_INOUT",
    ["offset", "data_length", "block"],
  ),  # bulk_data_inout_t
  0x0200: (
    "<HB52s",
    "DEVICE_DATA_READ_IN",
    ["field_index", "flags", "data"],
  ),  # device_data_read_in_t
  0x0210: (
    "<HB52s",
    "DEVICE_DATA_WRITE_OUT",
    ["field_index", "flags", "data"],
  ),  # device_data_write_out_t
  0x0220: (
    "<HH55sB",
    "DEVICE_DATA_FIELDS_IN",
    ["field_index", "field_count", "field_name", "flags"],
  ),  # device_data_fields_in_t
  0x0240: (
    "<HBI",
    "DEVICE_FILE_READ_IN",
    ["file_index", "flags", "file_size"],
  ),  # device_file_read_in_t
  0x0250: (
    "<HBI",
    "DEVICE_FILE_WRITE_OUT",
    ["file_index", "flags", "file_size"],
  ),  # device_file_write_out_t
  0x0260: (
    "<HH51sIBx",  # Added 1 byte padding to align with struct
    "DEVICE_FILE_NAMES_IN",
    ["file_index", "file_count", "file_name", "file_size", "flags"],
  ),  # device_file_names_in_t
  0x0280: (
    "<HH55sB",
    "RPC_ENDPOINTS_IN",
    ["rpc_index", "rpc_count", "function_name", "flags"],
  ),  # rpc_endpoints_in_t
  0x0290: (
    "<H",
    "RPC_EXECUTE_OUT",
    ["rpc_index"],
  ),  # rpc_execute_out_t
  0x0320: (
    "<hhBB",
    "ABS_TRIGGER_MEASUREMENT_OUT",
    [
      "signal_wavelength_nm",
      "reference_wavelength_nm",
      "is_reference_measurement",
      "flags",
    ],
  ),  # abs_trigger_measurement_out_t
  0x0330: (
    "<30h",
    "ABS_WAVELENGTHS_IN",
    ["wavelengths_nm"],
  ),  # abs_wavelengths_in_t
  0x0340: (
    "<i12sBB",
    "LUM_TRIGGER_MEASUREMENT_OUT",
    [
      "integration_time_us",
      "channels_selected",
      "is_reference_measurement",
      "flags",
    ],
  ),  # lum_trigger_measurement_out_t
  0x0341: (
    "<i48sBB",
    "LUM384_TRIGGER_MEASUREMENT_OUT",
    [
      "integration_time_us",
      "channels_selected",
      "is_reference_measurement",
      "flags",
    ],
  ),  # lum384_trigger_measurement_out_t
  0x0600: (
    "<BBII12fBB",
    "LUM96_MEASUREMENT_IN",
    [
      "seq",
      "seq_len",
      "integration_time_us",
      "duration_ms",
      "row",
      "flags",
      "progress",
    ],
  ),  # lum96_measurement_in_t
  0x0610: (
    "<BBII12fBB",
    "LUM384_MEASUREMENT_IN",
    [
      "seq",
      "seq_len",
      "integration_time_us",
      "duration_ms",
      "row",
      "flags",
      "progress",
    ],
  ),  # lum384_measurement_in_t
  0x1000: (
    "<BBBIBIBh",
    "ESP_STATUS_IN",
    [
      "connection_state",
      "last_connection_error",
      "is_connected_to_cloud",
      "last_cloud_connection_error_code",
      "self_test_status",
      "self_test_errors",
      "wifi_phy_mode",
      "wifi_rssi",
    ],
  ),  # esp_status_in_t
  0x0040: (
    "<B",
    "HEARTBEAT_IN",
    ["enabled"],
  ),  # heartbeat_in_t
}

# Create a reverse mapping from payload name to report ID for easier encoding
REPORT_IDS = {
  payload_name: report_id for report_id, (_, payload_name, _) in PAYLOAD_FORMATS.items()
}


def parse_hid_report(binary_data):
  # Parse the outer hid_report_t structure
  header_fmt = "<H"  # Little-endian: H=uint16 for report_id
  header_size = struct.calcsize(header_fmt)
  (report_id,) = struct.unpack(header_fmt, binary_data[:header_size])

  if report_id not in PAYLOAD_FORMATS:
    raise ValueError(f"Unknown report_id: 0x{report_id:04x}")

  payload_fmt, payload_name, field_names = PAYLOAD_FORMATS[report_id]
  payload_size = struct.calcsize(payload_fmt)
  payload_data = struct.unpack(payload_fmt, binary_data[header_size : header_size + payload_size])

  result = {
    "report_id": report_id,
    "payload_name": payload_name,
    "payload": {field: value for field, value in zip(field_names, payload_data)},
  }

  return result


def encode_hid_report(report_dict):
  report_id = report_dict.get("report_id")
  if report_id not in PAYLOAD_FORMATS:
    raise ValueError(f"Unknown report_id: 0x{report_id:04x}")

  payload_fmt, payload_name, field_names = PAYLOAD_FORMATS[report_id]

  # Prepare the payload data in the correct order, zeroing any fields that
  # were not explicitly provided.
  payload = report_dict.get("payload", {})

  # Find all format characters (e.g., 'i', 's', 'B') in the format string.
  format_chars = re.findall(r"\d*([a-zA-Z])", payload_fmt)

  payload_data = []
  for i, field in enumerate(field_names):
    default = b"" if format_chars[i] == "s" else 0
    p = payload.get(field, default)
    if isinstance(p, list):
      payload_data.extend(p)
    else:
      payload_data.append(p)

  # Encode the payload
  binary_payload = struct.pack(payload_fmt, *payload_data)

  # Encode the full report (header + payload)
  header_fmt = "<H"
  binary_header = struct.pack(header_fmt, report_id)
  packet = binary_header + binary_payload
  packet += b"\x00" * (64 - len(packet))

  return packet
