cell_loader_ctrl/robot_control/src/api/grbl_handler.py

import asyncio
import serial_asyncio
from typing import List
import logging
import datetime
import collections

MOVE_RATE = 50
APPROACHE_RATE = 10 
CELL_CIRCLE_RADIUS = 0.35

HOMING_TIMEOUT_S = 80

logger = logging.getLogger(__name__)

class GRBLHandler:
    def __init__(self, port: str, baudrate: int):
        self.port = port
        self.debug = port == "debug"
        self.baudrate = baudrate
        
        self.controller_active = asyncio.Event()
        
        self.reader = None
        self.writer = None
        self.position_callbacks = []

    async def connect(self):
        """Connect to GRBL controller"""
        if self.debug:
            logger.debug("GRBLHandler is not connected in debug mode")
            return

        logger.info("Connecting to GRBL...")
        try:
            self.reader, self.writer = await serial_asyncio.open_serial_connection(
                url = self.port,
                baudrate = self.baudrate
            )            
        except serial_asyncio.serial.SerialException as e:
            raise serial_asyncio.serial.SerialException(f"Failed to connect to robot: {str(e)}")

        init_response = await self._process_response()
        if not any(isinstance(response, collections.abc.Iterable) and "Grbl" in response for response in init_response):
            raise RuntimeError("Failed to connect to GRBL")
        
        logger.info("GRBL connected.")

        # Set GRBL settings:
        # Enable homing ($22=1)
        # Run homing cycle ($H)
        # Set absolute positioning mode (G90)
        # Set units to millimeters (G21) 
        await self.send_gcode(['$22=1']) 
        await self.send_gcode(['$H'], HOMING_TIMEOUT_S) 
        await self.send_gcode(['G90', 'G21']) 

        # if current pos not 0 set working pos to current pos
        # (should not be done every time to save eeprom writes)
        status_response = await self.get_status()
        current_pos = (0.0, 0.0, 0.0)
        if status_response and 'WPos:' in status_response:
            pos_str = status_response.split('WPos:')[1].split('|')[0]
            x, y, z = map(float, pos_str.split(','))
            current_pos = (x, y, z)
        # If not at (0,0,0), set working position to current position
        if any(abs(coord) > 1e-6 for coord in current_pos):
            await self.send_gcode([f'G10 L20 P1 X0 Y0 Z0'])
            logger.info(f"Set working position to current position: {current_pos}")

    async def reset(self):
        """Reset GRBL controller"""
        if self.debug:
            logger.debug("GRBLHandler is not reset in debug mode")
            return

        if self.writer:
            self.writer.write(b"\x18")  # Soft-reset GRBL
            await self.writer.drain()
            await asyncio.sleep(0.5)  # Wait for GRBL to reset
            await self.close()
            
        await self.connect()  # Reconnect after reset

    async def send_gcode(self, commands: List[str], timeout_s = 6):
        """Send GCODE commands"""
        if self.debug:
            logger.debug(f"G-Code commands [{*commands,}] not sent in debug mode")
            return

        self.controller_active.set()
        try:
            for cmd in commands:
                logger.debug(f"Sending G-Code command: {cmd}")
                self.writer.write(f"{cmd}\n".encode())
                await self.writer.drain()
                if not (await self._process_response(timeout_s)):
                    raise RuntimeError("Did not receive response from GRBL")
        except RuntimeError as e:
            logger.error(f"Failed to send G-Code commands: {str(e)}")
        finally:
            self.controller_active.clear()

    async def get_status(self):
        """Get current GRBL status
        
        Returns:
            str: Status response from GRBL, or None if in debug mode
        """
        if self.debug:
            return None
            
        if not self.writer:
            return None
        
        response_lines = None
        try:
            self.writer.write(b"?") # '?' command returns status report
            await self.writer.drain()
            response_lines = str(await self._process_response())
        except Exception as e:
            logger.error(f"Failed to get status: {str(e)}")
            return None
        
        if not response_lines:
            return None
        
        for line in response_lines:
            if not 'ALARM' in line:
                continue
            if ":1" in line:
                raise RuntimeError("Hard Limit was triggered!")
            raise RuntimeError(f"Grbl threw ALARM: {line}")
        
        return response_lines

    async def wait_until_idle(self, timeout_s, position: list[float] = None):
        """Wait until GRBL reports idle status
        
        Args:
            timeout_s: Timeout in seconds
            position: Optional list of 3 floats that will be updated with current position
        """
        if self.debug:
            await asyncio.sleep(1)
            return
            
        start = datetime.datetime.now()
        # response = await self.get_status() # First response can still be idle
        while True:
            response = await self.get_status()
            if response and 'WPos:' in response:
                # Parse position from status reports (<Idle|WPos:0.000,0.000,0.000|...>)
                pos_str = response.split('WPos:')[1].split('|')[0]
                if position is not None:
                    x, y, z = map(float, pos_str.split(','))
                    position = (x,y,z)

            if response and "Idle" in response:
                logger.debug("Movement complete.\nContinuing...")
                break
                    
            now = datetime.datetime.now()
            if (now - start).total_seconds() > timeout_s:
                logger.error("Waiting on idle took too long!")
                raise TimeoutError("GRBL did not report idle status")
                
            await asyncio.sleep(0.5)  # Async delay to prevent flooding

    async def send_and_wait_gcode(self, commands: List[str], timeout_s=60, position: list[float] = None):
        """Send GCODE commands and wait until machine is idle"""
        await self.send_gcode(commands)
        await asyncio.sleep(0.2)  # Delay to allow GRBL to process commands
        await self.wait_until_idle(timeout_s, position)

    def _check_pos_callbacks(self, line):
        # Parse position from status reports (<Idle|WPos:0.000,0.000,0.000|...>)
        if line.startswith('<') and 'WPos:' in line:
            pos_str = line.split('WPos:')[1].split('|')[0]
            x, y, z = map(float, pos_str.split(','))
            self.current_position = (x,y,z)
            
            # Notify all registered callbacks
            for callback in self.position_callbacks:
                callback(self.current_position)

    async def _process_response(self, timeout_s=6.0):
        """Process GRBL responses"""
        if self.reader:
            response_lines = []
            first = True
            while True:
                try:
                    current_timeout = timeout_s if first else 0.1
                    line = await asyncio.wait_for(self.reader.readuntil(), timeout=current_timeout)
                    decoded = line.strip().decode("utf-8")
                    self._check_pos_callbacks(decoded)
                    logger.debug(f"G-Code response: {decoded}")
                    response_lines.append(decoded)
                    first = False
                except asyncio.TimeoutError:
                    # No more data available right now
                    break
                except Exception as e:
                    raise RuntimeError(f"Failed to process response: {e}")
            
            if not response_lines:
                logger.warning(f"No GRBL response received! ({timeout_s}s)")

            return response_lines

    def register_position_callback(self, callback):
        """Register callback for position updates
        
        Args:
            callback: Function taking tuple with X,Y,Z coordinates
        """
        self.position_callbacks.append(callback)

    async def close(self):
        """Close connection"""
        if self.writer:
            self.writer.close()
            await self.writer.wait_closed()
        self.reader = None
        self.writer = None