import _thread
from machine import UART
import time
import ujson

class SafeUART:
    def __init__(self, uart_id, tx_pin, rx_pin, baudrate=115200):
        # Initialize the hardware UART channel
        self.uart = UART(uart_id, baudrate=baudrate, tx=tx_pin, rx=rx_pin, timeout=10, rxbuf=1024)
        
        # Core thread-safety assets
        self.lock = _thread.allocate_lock()
        self.rx_queue = []
        self.buffer = b""
        
        # Start the background data worker thread
        _thread.stack_size(4096) # Cap the stack size for the UART listener
        _thread.start_new_thread(self._listener_worker, ())
        _thread.stack_size(0) 
        
        print(f"[UART] Thread initialized on UART{uart_id} (TX:{tx_pin}, RX:{rx_pin})")

    def _listener_worker(self):
        """Worker loop that handles nested JSON structures by tracking brace depth."""
        while True:
            messages_found = []
            
            with self.lock:
                if self.uart.any():
                    chunk = self.uart.read()
                    if chunk is not None and isinstance(chunk, bytes):
                        self.buffer += chunk
                
                # Extract complete JSON objects while accounting for nested braces
                while True:
                    start_idx = self.buffer.find(b'{')
                    if start_idx == -1:
                        # No starting brace; clear any garbage bytes currently in buffer
                        self.buffer = b""
                        break
                    
                    # Trim any leading noise before the first '{'
                    if start_idx > 0:
                        self.buffer = self.buffer[start_idx:]
                    
                    # Track depth to find the matching OUTER '}'
                    depth = 0
                    in_string = False
                    escape = False
                    end_idx = -1
                    
                    for i in range(len(self.buffer)):
                        b = self.buffer[i]
                        
                        # Ignore braces inside string literals ("...")
                        if b == 34 and not escape:  # 34 is ASCII for '"'
                            in_string = not in_string
                        elif b == 92 and in_string: # 92 is ASCII for '\'
                            escape = not escape
                            continue
                        elif not in_string:
                            if b == 123:   # '{'
                                depth += 1
                            elif b == 125: # '}'
                                depth -= 1
                                if depth == 0:
                                    end_idx = i
                                    break
                        escape = False
                    
                    if end_idx != -1:
                        # Full nested JSON object extracted safely
                        json_bytes = self.buffer[:end_idx + 1]
                        self.buffer = self.buffer[end_idx + 1:]
                        messages_found.append(json_bytes)
                    else:
                        # The complete outer JSON hasn't fully arrived yet; wait for next UART chunk
                        break

            # Process valid complete frames outside the lock
            for json_bytes in messages_found:
                try:
                    decoded_str = json_bytes.decode('utf-8')
                    with self.lock:
                        self.rx_queue.append(decoded_str)
                except Exception as e:
                    print(f"[UART Parse Error]: {e}")
                    
            time.sleep_ms(10)

    def send(self, message):
        """Safely pushes strings across the serial wire from any thread context."""
        if not message.endswith('\n'):
            message += '\n'
        
        data = message.encode('utf-8')
        
        with self.lock:
            self.uart.write(data)
            
    def send_as_command(self, command: 'UARTCommand'):
        """Safely sends a structured command over UART."""
        json_message = command.to_json()
        self.send(json_message)

    def any(self):
        """Checks if any complete messages are waiting to be read."""
        with self.lock:
            return len(self.rx_queue) > 0

    def read(self):
        """Pulls the oldest unread string from the queue. Returns None if empty."""
        with self.lock:
            if self.rx_queue:
                return self.rx_queue.pop(0)
        return None
    
    def read_as_command(self) -> 'UARTCommand | None':
        """Attempts to read the oldest unread string and parse it as a UARTCommand. Returns None if empty or invalid."""
        raw_message = self.read()
        if raw_message is not None:
            cmd = UARTCommand.from_json(raw_message)
            if cmd is None:
                print("[UART] Impossible de traiter le message brut :", raw_message)
            return cmd
        return None


class UARTCommand:
    """A simple wrapper for commands sent over UART, allowing for structured data."""
    def __init__(self, command_type: str, payload):
        self.command_type = command_type
        self.payload = payload

    def to_json(self):
        """Serializes the command to a JSON string."""
        return ujson.dumps({
            "command_type": self.command_type,
            "payload": self.payload
        })
        
    @staticmethod
    def from_json(json_string: str) -> 'UARTCommand | None':
        """Deserializes a JSON string into a UARTCommand object."""
        try:
            # Remplacement préventif si des guillemets simples sont reçus
            clean_str = json_string.replace("'", '"') if "'" in json_string else json_string
            data = ujson.loads(clean_str)
            
            if not isinstance(data, dict):
                return None
                
            return UARTCommand(data.get("command_type"), data.get("payload"))
        except Exception as err:
            # Affiche l'erreur exacte rencontrée par ujson (ex: syntax error)
            print(f"[UARTCommand Parsing Error]: {err} -> Contenu: {json_string}")
            return None


class UARTCommandType:
    """Enumeration of known UART command types."""
    COOKING_PARAMS = "COOKING_PARAMS"
    COOKING_STATE_UPDATE = "COOKING_STATE_UPDATE"