import asyncio
import threading
import time  # Import time for time.sleep()
import warnings
from abc import ABC
from collections import deque
from typing import Dict, Iterable, List, Literal, Optional, Union

import can

from pylabrobot.arms.backend import (
  AccessPattern,
  HorizontalAccess,
  SCARABackend,
  VerticalAccess,
)
from pylabrobot.resources import Coordinate, Rotation


# Placeholder for KX2 specific types and error codes
# These will be populated as we translate the C# code.
class KX2CartesianCoords:
  # Placeholder for KX2 specific Cartesian coordinates
  pass


class KX2JointCoords:
  # Placeholder for KX2 specific Joint coordinates
  pass


class KX2Error(Exception):
  def __init__(self, replycode: int, message: str):
    self.replycode = replycode
    self.message = message
    super().__init__(f"KX2Error {replycode}: {message}")


# --- CANopen COB-IDs ---
# Base COB-ID for NMT, SYNC, EMCY, PDOs, SDOs
# NMT (Network Management)
NMT_BASE_COB_ID = 0x000

# SYNC (Synchronization)
SYNC_COB_ID = 0x080

# EMCY (Emergency)
EMCY_BASE_COB_ID = 0x081  # + Node ID

# PDOs (Process Data Objects)
# TPDO (Transmit PDO) - sent by node
TPDO1_BASE_COB_ID = 0x180  # + Node ID
TPDO2_BASE_COB_ID = 0x280  # + Node ID
TPDO3_BASE_COB_ID = 0x380  # + Node ID
TPDO4_BASE_COB_ID = 0x480  # + Node ID

# RPDO (Receive PDO) - received by node
RPDO1_BASE_COB_ID = 0x200  # + Node ID
RPDO2_BASE_COB_ID = 0x300  # + Node ID
RPDO3_BASE_COB_ID = 0x400  # + Node ID
RPDO4_BASE_COB_ID = 0x500  # + Node ID

# SDOs (Service Data Objects)
SDO_TX_BASE_COB_ID = 0x580  # Node -> Master
SDO_RX_BASE_COB_ID = 0x600  # Master -> Node


class CANMotor:
  """
  Encapsulates low-level CAN communication for the KX2 robot.
  Mimics the functionality of the C# `clsCANMotor` class.
  """

  def __init__(self, channel: str = "PCAN_USBBUS1", bustype: str = "pcan", bitrate: int = 500000):
    self.channel = channel
    self.bustype = bustype
    self.bitrate = bitrate
    self.bus: Optional[can.Bus] = None
    self._rx_queue: deque[can.Message] = deque()
    self._tx_queue: deque[can.Message] = deque()
    self._event_queue: deque[Dict] = deque()  # General events for higher-level processing

    self._read_thread: Optional[threading.Thread] = None
    self._write_thread: Optional[threading.Thread] = None
    self._event_thread: Optional[threading.Thread] = None

    self._running = False
    self._message_received_event = threading.Event()  # For synchronous message waiting

  def connect(self):
    """Establishes connection to the CAN bus."""
    if self.bus and self.bus.is_connected:
      print("CAN bus already connected.")
      return

    try:
      self.bus = can.interface.Bus(channel=self.channel, bustype=self.bustype, bitrate=self.bitrate)
      print(f"Successfully connected to CAN bus: {self.channel}")
      self._running = True

      # Start threads
      self._read_thread = threading.Thread(target=self._can_read_thread, daemon=True)
      self._write_thread = threading.Thread(target=self._can_write_thread, daemon=True)
      self._event_thread = threading.Thread(target=self._event_processing_thread, daemon=True)

      self._read_thread.start()
      self._write_thread.start()
      self._event_thread.start()

    except Exception as e:
      self.bus = None
      self._running = False
      raise KX2Error(-1, f"Failed to connect to CAN bus: {e}")

  def disconnect(self):
    """Closes the CAN bus connection."""
    if self.bus:
      self._running = False
      # Wait for threads to finish (with a timeout)
      if self._read_thread:
        self._read_thread.join(timeout=1)
      if self._write_thread:
        self._write_thread.join(timeout=1)
      if self._event_thread:
        self._event_thread.join(timeout=1)

      self.bus.shutdown()
      self.bus = None
      print("Disconnected from CAN bus.")
    else:
      print("CAN bus not connected.")

  def send_message(self, msg: can.Message):
    """Sends a CAN message asynchronously by adding it to the TX queue."""
    if not self.bus:
      raise KX2Error(-1, "CAN bus not connected. Cannot send message.")
    self._tx_queue.append(msg)

  def _send_can_message_sync(self, msg: can.Message, timeout: float = 0.1) -> Optional[can.Message]:
    """
    Sends a CAN message synchronously and waits for a response.
    This is primarily for setup commands that require immediate feedback.
    """
    if not self.bus:
      raise KX2Error(-1, "CAN bus not connected. Cannot send message synchronously.")

    self._message_received_event.clear()
    self.bus.send(msg)  # Send directly, bypassing TX queue for immediate send

    # Wait for a response. The _can_read_thread will populate _rx_queue
    # and potentially set _message_received_event if it's a response we're waiting for.
    # For now, we'll just wait for any message and then check the queue.
    # A more robust implementation would match response COB-IDs.
    if self._message_received_event.wait(timeout=timeout):
      # Attempt to find the relevant response in the RX queue
      # This is a simplified approach. A real SDO/PDO implementation would
      # need to match the response COB-ID and potentially SDO index/subindex.
      while self._rx_queue:
        response_msg = self._rx_queue.popleft()
        # TODO: Add logic to match response to the sent message
        return response_msg
    return None

  def _can_read_thread(self):
    """Continuously reads messages from the CAN bus and puts them into the RX queue."""
    # Mimics C# CAN_Read_Thread
    print("CAN read thread started.")
    while self._running:
      try:
        msg = self.bus.recv(timeout=0.1)  # Non-blocking or short timeout
        if msg:
          self._rx_queue.append(msg)
          self._message_received_event.set()  # Signal that a message was received
          self._process_incoming_can_message(msg)
      except Exception as e:
        if self._running:  # Only print errors if still running deliberately
          print(f"Error in CAN read thread: {e}")
        # Allow thread to exit cleanly if _running is False
    print("CAN read thread stopped.")

  def _can_write_thread(self):
    """Continuously takes messages from the TX queue and sends them to the CAN bus."""
    # Mimics C# CAN_Write_Thread
    print("CAN write thread started.")
    while self._running:
      try:
        if self._tx_queue:
          msg = self._tx_queue.popleft()
          self.bus.send(msg)
        else:
          time.sleep(0.001)  # Small delay to prevent busy-waiting
      except can.exceptions.CanOperationError as e:
        print(f"CAN write error: {e}")
        # Optionally push an error event to _event_queue
      except Exception as e:
        if self._running:
          print(f"Error in CAN write thread: {e}")
    print("CAN write thread stopped.")

  def _event_processing_thread(self):
    """Processes events from the event queue and dispatches them."""
    # Mimics C# EventThread
    print("Event processing thread started.")
    while self._running:
      try:
        if self._event_queue:
          event = self._event_queue.popleft()
          # TODO: Process events here. This will involve calling event handlers
          # similar to how the C# code does it.
          # For now, just print.
          print(f"Processing event: {event}")
        else:
          time.sleep(0.001)
      except Exception as e:
        if self._running:
          print(f"Error in event processing thread: {e}")
    print("Event processing thread stopped.")

  def _process_incoming_can_message(self, msg: can.Message):
    """
    Processes an incoming CAN message, similar to the switch-case in C# CAN_Read().
    This will eventually parse SDOs, PDOs, EMCYs, etc.
    For now, it just adds a generic event.
    """
    # This is where the bulk of the C# CAN_Read parsing logic will go.
    # It will identify message type (NMT, SDO, PDO, EMCY) and payload.
    # Then it will create a structured event and push it to _event_queue.
    self._event_queue.append({"type": "CAN_MESSAGE", "msg": msg})
    # print(f"Received CAN message: {msg}") # Commented out to reduce verbosity

  # --- NMT Commands ---
  def nmt_command(self, command_code: int, node_id: int = 0):
    """
    Sends an NMT command to a specific node or all nodes (node_id=0).
    command_code: 1 for start, 2 for stop, 128 for pre-operational, 129 for reset node, 130 for reset communication.
    """
    arbitration_id = NMT_BASE_COB_ID
    data = [command_code, node_id]
    msg = can.Message(arbitration_id=arbitration_id, data=data, is_extended_id=False)
    self.send_message(msg)
    print(f"Sent NMT command: {command_code} to node {node_id}")


# SDOs (Service Data Objects)
SDO_TX_BASE_COB_ID = 0x580  # Node -> Master
SDO_RX_BASE_COB_ID = 0x600  # Master -> Node

# --- SDO Command Codes (Client -> Server) ---
SDO_CCS_DOWNLOAD_SEGMENT = 0x00  # Download segment
SDO_CCS_INITIATE_DOWNLOAD = 0x20  # Initiate download request
SDO_CCS_INITIATE_UPLOAD = 0x40  # Initiate upload request
SDO_CCS_UPLOAD_SEGMENT = 0x60  # Upload segment
SDO_CCS_ABORT = 0x80  # Abort SDO transfer

# --- SDO Command Codes (Server -> Client) ---
SDO_SCS_UPLOAD_SEGMENT = 0x00  # Upload segment
SDO_SCS_INITIATE_UPLOAD = 0x40  # Initiate upload response
SDO_SCS_INITIATE_DOWNLOAD = 0x60  # Initiate download response
SDO_SCS_DOWNLOAD_SEGMENT = 0x20  # Download segment
SDO_SCS_ABORT = 0x80  # Abort SDO transfer


class CANMotor:
  """
  Encapsulates low-level CAN communication for the KX2 robot.
  Mimics the functionality of the C# `clsCANMotor` class.
  """

  def __init__(self, channel: str = "PCAN_USBBUS1", bustype: str = "pcan", bitrate: int = 500000):
    self.channel = channel
    self.bustype = bustype
    self.bitrate = bitrate
    self.bus: Optional[can.Bus] = None
    self._rx_queue: deque[can.Message] = deque()
    self._tx_queue: deque[can.Message] = deque()
    self._event_queue: deque[Dict] = deque()  # General events for higher-level processing

    self._read_thread: Optional[threading.Thread] = None
    self._write_thread: Optional[threading.Thread] = None
    self._event_thread: Optional[threading.Thread] = None

    self._running = False
    self._message_received_event = threading.Event()  # For synchronous message waiting

  def connect(self):
    """Establishes connection to the CAN bus."""
    if self.bus and self.bus.is_connected:
      print("CAN bus already connected.")
      return

    try:
      self.bus = can.interface.Bus(channel=self.channel, bustype=self.bustype, bitrate=self.bitrate)
      print(f"Successfully connected to CAN bus: {self.channel}")
      self._running = True

      # Start threads
      self._read_thread = threading.Thread(target=self._can_read_thread, daemon=True)
      self._write_thread = threading.Thread(target=self._can_write_thread, daemon=True)
      self._event_thread = threading.Thread(target=self._event_processing_thread, daemon=True)

      self._read_thread.start()
      self._write_thread.start()
      self._event_thread.start()

    except Exception as e:
      self.bus = None
      self._running = False
      raise KX2Error(-1, f"Failed to connect to CAN bus: {e}")

  def disconnect(self):
    """Closes the CAN bus connection."""
    if self.bus:
      self._running = False
      # Wait for threads to finish (with a timeout)
      if self._read_thread:
        self._read_thread.join(timeout=1)
      if self._write_thread:
        self._write_thread.join(timeout=1)
      if self._event_thread:
        self._event_thread.join(timeout=1)

      self.bus.shutdown()
      self.bus = None
      print("Disconnected from CAN bus.")
    else:
      print("CAN bus not connected.")

  def send_message(self, msg: can.Message):
    """Sends a CAN message asynchronously by adding it to the TX queue."""
    if not self.bus:
      raise KX2Error(-1, "CAN bus not connected. Cannot send message.")
    self._tx_queue.append(msg)

  def _send_can_message_sync(self, msg: can.Message, timeout: float = 0.1) -> Optional[can.Message]:
    """
    Sends a CAN message synchronously and waits for a response.
    This is primarily for setup commands that require immediate feedback.
    """
    if not self.bus:
      raise KX2Error(-1, "CAN bus not connected. Cannot send message synchronously.")

    self._message_received_event.clear()
    self.bus.send(msg)  # Send directly, bypassing TX queue for immediate send

    # Wait for a response. The _can_read_thread will populate _rx_queue
    # and potentially set _message_received_event if it's a response we're waiting for.
    # For now, we'll just wait for any message and then check the queue.
    # A more robust implementation would match response COB-IDs.
    if self._message_received_event.wait(timeout=timeout):
      # Attempt to find the relevant response in the RX queue
      # This is a simplified approach. A real SDO/PDO implementation would
      # need to match the response COB-ID and potentially SDO index/subindex.
      while self._rx_queue:
        response_msg = self._rx_queue.popleft()
        # TODO: Add logic to match response to the sent message
        return response_msg
    return None

  def _can_read_thread(self):
    """Continuously reads messages from the CAN bus and puts them into the RX queue."""
    # Mimics C# CAN_Read_Thread
    print("CAN read thread started.")
    while self._running:
      try:
        msg = self.bus.recv(timeout=0.1)  # Non-blocking or short timeout
        if msg:
          self._rx_queue.append(msg)
          self._message_received_event.set()  # Signal that a message was received
          self._process_incoming_can_message(msg)
      except Exception as e:
        if self._running:  # Only print errors if still running deliberately
          print(f"Error in CAN read thread: {e}")
        # Allow thread to exit cleanly if _running is False
    print("CAN read thread stopped.")

  def _can_write_thread(self):
    """Continuously takes messages from the TX queue and sends them to the CAN bus."""
    # Mimics C# CAN_Write_Thread
    print("CAN write thread started.")
    while self._running:
      try:
        if self._tx_queue:
          msg = self._tx_queue.popleft()
          self.bus.send(msg)
        else:
          time.sleep(0.001)  # Small delay to prevent busy-waiting
      except can.exceptions.CanOperationError as e:
        print(f"CAN write error: {e}")
        # Optionally push an error event to _event_queue
      except Exception as e:
        if self._running:
          print(f"Error in CAN write thread: {e}")
    print("CAN write thread stopped.")

  def _event_processing_thread(self):
    """Processes events from the event queue and dispatches them."""
    # Mimics C# EventThread
    print("Event processing thread started.")
    while self._running:
      try:
        if self._event_queue:
          event = self._event_queue.popleft()
          # TODO: Process events here. This will involve calling event handlers
          # similar to how the C# code does it.
          # For now, just print.
          print(f"Processing event: {event}")
        else:
          time.sleep(0.001)
      except Exception as e:
        if self._running:
          print(f"Error in event processing thread: {e}")
    print("Event processing thread stopped.")

  def _process_incoming_can_message(self, msg: can.Message):
    """
    Processes an incoming CAN message, similar to the switch-case in C# CAN_Read().
    This will eventually parse SDOs, PDOs, EMCYs, etc.
    For now, it just adds a generic event.
    """
    # This is where the bulk of the C# CAN_Read parsing logic will go.
    # It will identify message type (NMT, SDO, PDO, EMCY) and payload.
    # Then it will create a structured event and push it to _event_queue.
    self._event_queue.append({"type": "CAN_MESSAGE", "msg": msg})
    # print(f"Received CAN message: {msg}") # Commented out to reduce verbosity

  # --- NMT Commands ---
  def nmt_command(self, command_code: int, node_id: int = 0):
    """
    Sends an NMT command to a specific node or all nodes (node_id=0).
    command_code: 1 for start, 2 for stop, 128 for pre-operational, 129 for reset node, 130 for reset communication.
    """
    arbitration_id = NMT_BASE_COB_ID
    data = [command_code, node_id]
    msg = can.Message(arbitration_id=arbitration_id, data=data, is_extended_id=False)
    self.send_message(msg)
    print(f"Sent NMT command: {command_code} to node {node_id}")

  # --- SDO Commands ---
  def _sdo_request(
    self,
    node_id: int,
    index: int,
    subindex: int,
    data: Optional[List[int]] = None,
    is_write: bool = False,
    timeout: float = 0.5,
  ) -> Optional[List[int]]:
    """
    Sends an SDO request and waits for a response.
    """
    arbitration_id = SDO_RX_BASE_COB_ID + node_id

    if is_write:
      # Initiate download request (write)
      command_byte = SDO_CCS_INITIATE_DOWNLOAD
      # Data bytes: command_byte, index_low, index_high, subindex, data...
      sdo_data = [
        command_byte,
        index & 0xFF,
        (index >> 8) & 0xFF,
        subindex,
      ]
      if data is not None:
        sdo_data.extend(data)
      msg = can.Message(arbitration_id=arbitration_id, data=sdo_data, is_extended_id=False)
      response = self._send_can_message_sync(msg, timeout=timeout)
      # TODO: Parse response for SDO_SCS_INITIATE_DOWNLOAD or SDO_SCS_ABORT
      if response and response.arbitration_id == SDO_TX_BASE_COB_ID + node_id:
        if response.data[0] == SDO_SCS_INITIATE_DOWNLOAD:
          return []  # Success
        elif response.data[0] == SDO_SCS_ABORT:
          abort_code = int.from_bytes(response.data[4:8], byteorder="little")
          raise KX2Error(abort_code, f"SDO write aborted by device: {abort_code:X}")
      raise KX2Error(-1, "SDO write failed or no valid response.")
    else:
      # Initiate upload request (read)
      command_byte = SDO_CCS_INITIATE_UPLOAD
      sdo_data = [
        command_byte,
        index & 0xFF,
        (index >> 8) & 0xFF,
        subindex,
        0x00,
        0x00,
        0x00,
        0x00,  # Padding
      ]
      msg = can.Message(arbitration_id=arbitration_id, data=sdo_data, is_extended_id=False)
      response = self._send_can_message_sync(msg, timeout=timeout)
      # TODO: Parse response for SDO_SCS_INITIATE_UPLOAD or SDO_SCS_ABORT
      if response and response.arbitration_id == SDO_TX_BASE_COB_ID + node_id:
        if response.data[0] == SDO_SCS_INITIATE_UPLOAD:
          # Data is in response.data[4:]
          return list(response.data[4:])
        elif response.data[0] == SDO_SCS_ABORT:
          abort_code = int.from_bytes(response.data[4:8], byteorder="little")
          raise KX2Error(abort_code, f"SDO read aborted by device: {abort_code:X}")
      raise KX2Error(-1, "SDO read failed or no valid response.")

  def sdo_read(self, node_id: int, index: int, subindex: int, timeout: float = 0.5) -> List[int]:
    """Reads a value from the object dictionary of a node using SDO."""
    data = self._sdo_request(node_id, index, subindex, is_write=False, timeout=timeout)
    if data is None:
      raise KX2Error(
        -1,
        f"SDO read from node {node_id}, index {index:X}, subindex {subindex} timed out or failed.",
      )
    return data

  def sdo_write(
    self, node_id: int, index: int, subindex: int, data: List[int], timeout: float = 0.5
  ):
    """Writes a value to the object dictionary of a node using SDO."""
    response = self._sdo_request(
      node_id, index, subindex, data=data, is_write=True, timeout=timeout
    )
    if response is None:
      raise KX2Error(
        -1,
        f"SDO write to node {node_id}, index {index:X}, subindex {subindex} timed out or failed.",
      )
