StarPilot/starpilot/system/obdyssey/diagnostics.py

from __future__ import annotations

import re
from dataclasses import dataclass
from enum import IntEnum
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
  from openpilot.starpilot.system.obdyssey.elm327 import Elm327, ElmContext, DiagnosticResponse


class SERVICE_TYPE(IntEnum):
  DIAGNOSTIC_SESSION_CONTROL = 0x10
  ECU_RESET = 0x11
  CLEAR_DIAGNOSTIC_INFORMATION = 0x14
  READ_DTC_INFORMATION = 0x19
  READ_DATA_BY_IDENTIFIER = 0x22
  READ_MEMORY_BY_ADDRESS = 0x23
  READ_SCALING_DATA_BY_IDENTIFIER = 0x24
  SECURITY_ACCESS = 0x27
  COMMUNICATION_CONTROL = 0x28
  WRITE_DATA_BY_IDENTIFIER = 0x2E
  INPUT_OUTPUT_CONTROL_BY_IDENTIFIER = 0x2F
  ROUTINE_CONTROL = 0x31
  REQUEST_DOWNLOAD = 0x34
  REQUEST_UPLOAD = 0x35
  TRANSFER_DATA = 0x36
  REQUEST_TRANSFER_EXIT = 0x37
  WRITE_MEMORY_BY_ADDRESS = 0x3D
  TESTER_PRESENT = 0x3E


class SESSION_TYPE(IntEnum):
  DEFAULT = 1
  PROGRAMMING = 2
  EXTENDED_DIAGNOSTIC = 3
  SAFETY_SYSTEM_DIAGNOSTIC = 4


class RESET_TYPE(IntEnum):
  HARD = 1
  KEY_OFF_ON = 2
  SOFT = 3
  ENABLE_RAPID_POWER_SHUTDOWN = 4
  DISABLE_RAPID_POWER_SHUTDOWN = 5


class ACCESS_TYPE(IntEnum):
  REQUEST_SEED = 1
  SEND_KEY = 2


class ROUTINE_CONTROL_TYPE(IntEnum):
  START = 1
  STOP = 2
  REQUEST_RESULTS = 3


class DTC_REPORT_TYPE(IntEnum):
  NUMBER_OF_DTC_BY_STATUS_MASK = 0x01
  DTC_BY_STATUS_MASK = 0x02
  DTC_SNAPSHOT_IDENTIFICATION = 0x03
  DTC_SNAPSHOT_RECORD_BY_DTC_NUMBER = 0x04
  DTC_EXTENDED_DATA_RECORD_BY_DTC_NUMBER = 0x06
  SUPPORTED_DTC = 0x0A


UDS_NRC_DESCRIPTIONS: dict[int, str] = {
  0x10: "General Reject",
  0x11: "Service Not Supported",
  0x12: "Sub-function Not Supported",
  0x13: "Incorrect Message Length Or Invalid Format",
  0x14: "Response Too Long",
  0x21: "Busy Repeat Request",
  0x22: "Conditions Not Correct",
  0x24: "Request Sequence Error",
  0x25: "No Response From Subnet Component",
  0x26: "Failure Prevents Execution Of Requested Action",
  0x31: "Request Out Of Range",
  0x33: "Security Access Denied",
  0x35: "Invalid Key",
  0x36: "Exceed Number Of Attempts",
  0x37: "Required Time Delay Not Expired",
  0x70: "Upload Download Not Accepted",
  0x71: "Transfer Data Suspended",
  0x72: "General Programming Failure",
  0x73: "Wrong Block Sequence Counter",
  0x78: "Response Pending",
  0x7E: "Sub-function Not Supported In Active Session",
  0x7F: "Service Not Supported In Active Session",
}


class UdsNegativeResponseError(Exception):
  def __init__(self, service_id: int, nrc: int):
    self.service_id = service_id
    self.nrc = nrc
    desc = UDS_NRC_DESCRIPTIONS.get(nrc, f"NRC 0x{nrc:02X}")
    super().__init__(f"UDS Negative Response: Service 0x{service_id:02X} failed with {desc} (0x{nrc:02X})")


@dataclass(frozen=True)
class DiagnosticTroubleCode:
  code: str
  ecu: str | None = None
  status: int | None = None
  source: str = "OBD"
  description: str | None = None
  raw: bytes = b""

  def to_dict(self) -> dict[str, Any]:
    return {
      "code": self.code,
      "ecu": self.ecu,
      "status": self.status,
      "source": self.source,
      "description": self.description,
      "raw": self.raw.hex().upper() if self.raw else "",
    }


def parse_standard_dtcs(data: bytes, source: str = "OBD") -> list[DiagnosticTroubleCode]:
  """Parse 2-byte standard SAE J1979 DTCs (e.g. from Mode 03, 07, 0A)."""
  dtcs: list[DiagnosticTroubleCode] = []
  prefix_map = {0: "P", 1: "C", 2: "B", 3: "U"}

  # Payload starts after response service byte (e.g. 0x43, 0x47, 0x4A)
  payload = data[1:] if len(data) > 0 and data[0] in (0x43, 0x47, 0x4A) else data

  for i in range(0, len(payload) - 1, 2):
    b0, b1 = payload[i], payload[i + 1]
    if b0 == 0 and b1 == 0:
      continue  # padding / no DTC
    prefix = prefix_map.get((b0 >> 6) & 0x03, "P")
    d1 = (b0 >> 4) & 0x03
    d2 = b0 & 0x0F
    d3 = (b1 >> 4) & 0x0F
    d4 = b1 & 0x0F
    code = f"{prefix}{d1}{d2:X}{d3:X}{d4:X}"
    dtcs.append(DiagnosticTroubleCode(
      code=code,
      source=source,
      raw=bytes([b0, b1]),
    ))
  return dtcs


def parse_uds_dtcs(data: bytes, ecu: str | None = None) -> list[DiagnosticTroubleCode]:
  """Parse UDS 3-byte DTCs (2 bytes DTC + 1 byte status mask) from Service 0x19 response."""
  dtcs: list[DiagnosticTroubleCode] = []
  prefix_map = {0: "P", 1: "C", 2: "B", 3: "U"}

  # Positive response: 0x59 <ReportType> <StatusAvailabilityMask> <DTC1_B0> <DTC1_B1> <DTC1_Status> ...
  if len(data) < 3 or data[0] != 0x59:
    return dtcs

  records = data[3:]  # Skip 59, report_type, status_mask
  for i in range(0, len(records) - 2, 3):
    b0, b1, status = records[i], records[i + 1], records[i + 2]
    if b0 == 0 and b1 == 0:
      continue
    prefix = prefix_map.get((b0 >> 6) & 0x03, "P")
    d1 = (b0 >> 4) & 0x03
    d2 = b0 & 0x0F
    d3 = (b1 >> 4) & 0x0F
    d4 = b1 & 0x0F
    code = f"{prefix}{d1}{d2:X}{d3:X}{d4:X}"
    dtcs.append(DiagnosticTroubleCode(
      code=code,
      ecu=ecu,
      status=status,
      source="UDS",
      raw=bytes([b0, b1, status]),
    ))
  return dtcs


# Standard OBD-II Functions

def read_current_data(elm: Elm327, pid: int, context: ElmContext | None = None) -> DiagnosticResponse:
  """Mode 01: Read current powertrain diagnostic data."""
  return elm.request(bytes([0x01, pid & 0xFF]), context, retry=True)


def read_freeze_frame(elm: Elm327, pid: int, frame: int = 0, context: ElmContext | None = None) -> DiagnosticResponse:
  """Mode 02: Read freeze frame data."""
  return elm.request(bytes([0x02, pid & 0xFF, frame & 0xFF]), context, retry=True)


def read_stored_dtcs(elm: Elm327, context: ElmContext | None = None) -> list[DiagnosticTroubleCode]:
  """Mode 03: Read confirmed/stored emission-related DTCs."""
  res = elm.request(bytes([0x03]), context, retry=True)
  return parse_standard_dtcs(res.payload, source="OBD_STORED")


def read_pending_dtcs(elm: Elm327, context: ElmContext | None = None) -> list[DiagnosticTroubleCode]:
  """Mode 07: Read pending DTCs detected during current/last drive cycle."""
  res = elm.request(bytes([0x07]), context, retry=True)
  return parse_standard_dtcs(res.payload, source="OBD_PENDING")


def read_permanent_dtcs(elm: Elm327, context: ElmContext | None = None) -> list[DiagnosticTroubleCode]:
  """Mode 0A: Read permanent DTCs."""
  res = elm.request(bytes([0x0A]), context, retry=True)
  return parse_standard_dtcs(res.payload, source="OBD_PERMANENT")


def read_all_dtcs(elm: Elm327, context: ElmContext | None = None) -> list[DiagnosticTroubleCode]:
  """Read confirmed (03), pending (07), and permanent (0A) DTCs."""
  results: list[DiagnosticTroubleCode] = []
  for func in (read_stored_dtcs, read_pending_dtcs, read_permanent_dtcs):
    try:
      results.extend(func(elm, context))
    except Exception:
      pass
  return results


def clear_dtcs(elm: Elm327, context: ElmContext | None = None) -> DiagnosticResponse:
  """Mode 04: Clear diagnostic trouble codes and reset MIL (Check Engine Light). Mutating!"""
  return elm.request(bytes([0x04]), context, retry=False)


def read_vin(elm: Elm327, context: ElmContext | None = None) -> str:
  """Mode 09 PID 02: Read Vehicle Identification Number (VIN)."""
  res = elm.request(bytes([0x09, 0x02]), context, retry=True)
  payload = res.payload

  # Response format: 49 02 <data_items_or_line_nums> followed by ASCII VIN characters
  # Strip service 49 02 header if present
  if len(payload) >= 2 and payload[0] == 0x49 and payload[1] == 0x02:
    payload = payload[2:]

  # Filter ascii alphanumeric printable characters
  ascii_chars = "".join(chr(b) for b in payload if 32 <= b <= 126)
  vin_match = re.search(r"([A-HJ-NPR-Z0-9]{17})", ascii_chars)
  if vin_match:
    return vin_match.group(1)
  return ascii_chars.strip()


# UDS (ISO 14229) Service Functions

def uds_read_data_by_identifier(elm: Elm327, did: int, context: ElmContext | None = None) -> bytes:
  """UDS Service 0x22: ReadDataByIdentifier."""
  payload = bytes([SERVICE_TYPE.READ_DATA_BY_IDENTIFIER, (did >> 8) & 0xFF, did & 0xFF])
  res = elm.request(payload, context, retry=True)
  data = res.payload

  # Positive response: 0x62 <DID_MSB> <DID_LSB> <DataBytes...>
  if len(data) >= 3 and data[0] == 0x62:
    resp_did = (data[1] << 8) | data[2]
    if resp_did == did:
      return data[3:]
    return data[1:]
  return data


def uds_write_data_by_identifier(elm: Elm327, did: int, data: bytes, context: ElmContext | None = None) -> bytes:
  """UDS Service 0x2E: WriteDataByIdentifier. Mutating!"""
  payload = bytes([SERVICE_TYPE.WRITE_DATA_BY_IDENTIFIER, (did >> 8) & 0xFF, did & 0xFF]) + data
  res = elm.request(payload, context, retry=False)
  return res.payload


def uds_diagnostic_session_control(elm: Elm327, session_type: int, context: ElmContext | None = None) -> bytes:
  """UDS Service 0x10: DiagnosticSessionControl."""
  payload = bytes([SERVICE_TYPE.DIAGNOSTIC_SESSION_CONTROL, session_type & 0xFF])
  res = elm.request(payload, context, retry=False)
  return res.payload


def uds_ecu_reset(elm: Elm327, reset_type: int, context: ElmContext | None = None) -> bytes:
  """UDS Service 0x11: ECUReset. Mutating!"""
  payload = bytes([SERVICE_TYPE.ECU_RESET, reset_type & 0xFF])
  res = elm.request(payload, context, retry=False)
  return res.payload


def uds_read_dtc_information(elm: Elm327, report_type: int = DTC_REPORT_TYPE.DTC_BY_STATUS_MASK,
                             status_mask: int = 0xFF, context: ElmContext | None = None,
                             ecu: str | None = None) -> list[DiagnosticTroubleCode]:
  """UDS Service 0x19: ReadDTCInformation."""
  payload = bytes([SERVICE_TYPE.READ_DTC_INFORMATION, report_type & 0xFF, status_mask & 0xFF])
  res = elm.request(payload, context, retry=True)
  return parse_uds_dtcs(res.payload, ecu=ecu)


def uds_clear_diagnostic_information(elm: Elm327, group: int = 0xFFFFFF, context: ElmContext | None = None) -> bytes:
  """UDS Service 0x14: ClearDiagnosticInformation. Mutating!"""
  payload = bytes([
    SERVICE_TYPE.CLEAR_DIAGNOSTIC_INFORMATION,
    (group >> 16) & 0xFF,
    (group >> 8) & 0xFF,
    group & 0xFF,
  ])
  res = elm.request(payload, context, retry=False)
  return res.payload


def uds_security_access(elm: Elm327, access_type: int, key_data: bytes = b"", context: ElmContext | None = None) -> bytes:
  """UDS Service 0x27: SecurityAccess (Request Seed or Send Key)."""
  payload = bytes([SERVICE_TYPE.SECURITY_ACCESS, access_type & 0xFF]) + key_data
  is_mutating = (access_type == ACCESS_TYPE.SEND_KEY)
  res = elm.request(payload, context, retry=not is_mutating)
  return res.payload


def uds_routine_control(elm: Elm327, routine_type: int, routine_id: int, option_record: bytes = b"", context: ElmContext | None = None) -> bytes:
  """UDS Service 0x31: RoutineControl. Mutating!"""
  payload = bytes([
    SERVICE_TYPE.ROUTINE_CONTROL,
    routine_type & 0xFF,
    (routine_id >> 8) & 0xFF,
    routine_id & 0xFF,
  ]) + option_record
  res = elm.request(payload, context, retry=False)
  return res.payload


def uds_input_output_control(elm: Elm327, did: int, control_option: int, control_state: bytes = b"", context: ElmContext | None = None) -> bytes:
  """UDS Service 0x2F: InputOutputControlByIdentifier. Mutating!"""
  payload = bytes([
    SERVICE_TYPE.INPUT_OUTPUT_CONTROL_BY_IDENTIFIER,
    (did >> 8) & 0xFF,
    did & 0xFF,
    control_option & 0xFF,
  ]) + control_state
  res = elm.request(payload, context, retry=False)
  return res.payload


def uds_tester_present(elm: Elm327, subfunction: int = 0x00, context: ElmContext | None = None) -> bytes:
  """UDS Service 0x3E: TesterPresent."""
  payload = bytes([SERVICE_TYPE.TESTER_PRESENT, subfunction & 0xFF])
  res = elm.request(payload, context, retry=True)
  return res.payload


# Safety Classification Helpers

READ_ONLY_SERVICES = {
  0x01, 0x02, 0x03, 0x07, 0x09, 0x0A,
  SERVICE_TYPE.READ_DTC_INFORMATION,
  SERVICE_TYPE.READ_DATA_BY_IDENTIFIER,
  SERVICE_TYPE.READ_MEMORY_BY_ADDRESS,
  SERVICE_TYPE.READ_SCALING_DATA_BY_IDENTIFIER,
  SERVICE_TYPE.TESTER_PRESENT,
}

MUTATING_SERVICES = {
  0x04,
  SERVICE_TYPE.DIAGNOSTIC_SESSION_CONTROL,
  SERVICE_TYPE.ECU_RESET,
  SERVICE_TYPE.CLEAR_DIAGNOSTIC_INFORMATION,
  SERVICE_TYPE.SECURITY_ACCESS,
  SERVICE_TYPE.WRITE_DATA_BY_IDENTIFIER,
  SERVICE_TYPE.INPUT_OUTPUT_CONTROL_BY_IDENTIFIER,
  SERVICE_TYPE.ROUTINE_CONTROL,
  SERVICE_TYPE.REQUEST_DOWNLOAD,
  SERVICE_TYPE.REQUEST_UPLOAD,
  SERVICE_TYPE.TRANSFER_DATA,
  SERVICE_TYPE.REQUEST_TRANSFER_EXIT,
  SERVICE_TYPE.WRITE_MEMORY_BY_ADDRESS,
}


def is_read_only_service(service: int) -> bool:
  return service in READ_ONLY_SERVICES


def is_mutating_service(service: int) -> bool:
  return service in MUTATING_SERVICES


def is_read_only_payload(payload: bytes) -> bool:
  if not payload:
    return True
  sid = payload[0]
  if sid == SERVICE_TYPE.SECURITY_ACCESS:
    # Subfunction 0x01 (Request Seed) is read-only; Subfunction 0x02 (Send Key) is mutating
    subfn = payload[1] if len(payload) > 1 else 0
    return subfn % 2 == 1  # Odd = request seed
  return is_read_only_service(sid)