"""CS8000 Python API - Main Class"""

import serial
import serial.tools.list_ports
import time
from typing import Optional, Dict, Any
from .protocol import CMD, EOF


class CS8000:
    """CS8000 RF Switch Controller"""

    def __init__(self, model: str, serial_number: Optional[str] = None):
        """
        Initialize CS8000 instance
        
        Args:
            model: Device model (e.g., "CS8216")
            serial_number: Optional device serial number for multi-device setups
        """
        self.model = model
        self.serial_number = serial_number
        self.port_path: Optional[str] = None
        self.baud_rate = 115200
        self.port: Optional[serial.Serial] = None
        self.buffer = []

        # Internal state for prepared connections
        self._prepared_connections: Dict[int, int] = {}
        self._sequence_steps = []

        # Trigger options (50R termination)
        self.trigger50r = False
        self.trigger_armed = False

        # Channel remapping table (user channel -> hardware channel)
        self._channel_map = {
            1: 13, 2: 14, 3: 15, 4: 16,
            5: 8, 6: 7, 7: 6, 8: 5,
            9: 12, 10: 11, 11: 10, 12: 9,
            13: 1, 14: 2, 15: 3, 16: 4
        }

    @staticmethod
    def _auto_detect(serial_number: Optional[str] = None) -> str:
        """Auto-detect CS8000 device"""
        ports = serial.tools.list_ports.comports()
        found = None

        if serial_number:
            # Look for specific serial number
            for port in ports:
                if port.serial_number and str(serial_number) in str(port.serial_number):
                    found = port
                    break
        else:
            # Look for any CS8000 device (serial numbers starting with 3080216)
            for port in ports:
                if port.serial_number and str(port.serial_number).startswith("3080216"):
                    found = port
                    break

        if not found:
            msg = f"No CS8000 device found with serial number {serial_number}" if serial_number else "No CS8000-compatible device found"
            raise RuntimeError(msg)

        return found.device

    def open(self, serial_number: Optional[str] = None) -> None:
        """
        Open connection to the device
        
        Args:
            serial_number: Optional serial number to connect to specific device
        """
        # Use provided serial number or the one from constructor
        target_serial = serial_number or self.serial_number

        self.port_path = CS8000._auto_detect(target_serial)
        print(f"Connected to {self.model} at {self.port_path}")

        self.port = serial.Serial(
            port=self.port_path,
            baudrate=self.baud_rate,
            timeout=1
        )

        # Clear prepared connections after opening
        self._prepared_connections = {}

        # Initialize the device
        self._sync()

    def close(self) -> None:
        """Close connection to the device"""
        if self.port and self.port.is_open:
            self.port.close()

        # Clear the port reference and buffer
        self.port = None
        self.buffer = []

    # IMMEDIATE MODE METHODS

    def _remap_channel(self, channel: int) -> int:
        """Remap user channel to hardware channel"""
        if channel == 0:
            return 0  # 0 stays 0 (disconnect)
        if channel in self._channel_map:
            return self._channel_map[channel]
        raise ValueError(f"Invalid channel: {channel}")

    def _remap_com_port(self, common_port: int) -> int:
        """No remapping - COM port numbers map directly to hardware"""
        if common_port in [1, 2]:
            return common_port  # Direct mapping
        raise ValueError(f"Invalid common port: {common_port}")

    def prepare(self, common_port: int, channel: int) -> None:
        """
        Prepare a connection for a COMMON port to a channel
        
        Args:
            common_port: COMMON port number (1 or 2)
            channel: Channel number (0 for disconnect, 1-16 for channels)
        """
        if common_port not in [1, 2]:
            raise ValueError("Common port must be 1 or 2")
        if channel < 0 or channel > 16:
            raise ValueError("Channel must be between 0 (disconnect) and 16")

        # Remap both COM port and channel before storing
        hw_com_port = self._remap_com_port(common_port)
        hw_channel = self._remap_channel(channel)

        self._prepared_connections[hw_com_port] = hw_channel

    def commit(self) -> None:
        """Apply all prepared connections immediately"""
        if not self.port or not self.port.is_open:
            raise RuntimeError("Device not connected. Call open() first.")

        # Load the current prepared connections as a single step sequence and trigger it
        self.arm_trigger(False, self.trigger50r)  # disarm trigger to avoid unintended step advance
        self.reset_sequence()
        self.load_step()
        self.start_sequence(1, True)  # single step to apply changes
        self.send_trigger()  # force trigger to advance to the loaded step

    # SEQUENCE MODE METHODS

    def reset_sequence(self) -> None:
        """Clear any previously loaded sequence steps"""
        if not self.port or not self.port.is_open:
            raise RuntimeError("Device not connected. Call open() first.")

        self._sequence_steps = []
        self._send_frame(0, 0, 0, CMD.SEQ_RESET)

    def load_step(self) -> None:
        """Load the currently prepared connections as a step in the sequence"""
        if not self.port or not self.port.is_open:
            raise RuntimeError("Device not connected. Call open() first.")

        # Store the current prepared connections as a sequence step
        step = self._prepared_connections.copy()
        self._sequence_steps.append(step)

        # Convert to hardware command and load into device memory
        step_index = len(self._sequence_steps) - 1
        word = self._convert_connections_to_word(step)
        self._send_frame(step_index & 0xFF, word & 0xFF, (word >> 8) & 0xFF, CMD.LOAD_STEP)

        # Don't clear prepared connections here - they might be reused

    def start_sequence(self, loops: int = 1, arm_trigger: bool = True) -> None:
        """
        Start sequence execution with specified number of loops
        
        Args:
            loops: Number of times to loop the sequence
            arm_trigger: Whether to arm the trigger
        """
        if not self.port or not self.port.is_open:
            raise RuntimeError("Device not connected. Call open() first.")

        if not self._sequence_steps:
            raise RuntimeError("No sequence steps loaded. Call load_step() first.")

        # Set the number of steps and loops
        self._send_frame(len(self._sequence_steps) & 0xFF, 0, 0, CMD.STEPS_COUNT)
        self._send_frame(loops & 0xFF, (loops >> 8) & 0xFF, 0, CMD.LOOPS_COUNT)

        # Arm the trigger
        self.arm_trigger(arm_trigger, self.trigger50r)

    def arm_trigger(self, arm: bool = True, enabled_50r: bool = False) -> None:
        """
        Arm or disarm the trigger
        
        Args:
            arm: Whether to arm the trigger
            enabled_50r: Whether to enable 50R termination
        """
        self.trigger50r = enabled_50r
        self.trigger_armed = arm
        trig_word = (1 << 0 if arm else 0) | (1 << 1 if enabled_50r else 0)
        self._send_frame(trig_word, 0, 0, CMD.ARM_TRIGGER)

    def send_trigger(self) -> None:
        """Send a trigger signal to advance to the next step"""
        if not self.port or not self.port.is_open:
            raise RuntimeError("Device not connected. Call open() first.")

        self._send_frame(0, 0, 0, CMD.FORCE_TRIG)

    def stop_sequence(self) -> None:
        """Stop the sequence by disarming the trigger"""
        if not self.port or not self.port.is_open:
            raise RuntimeError("Device not connected. Call open() first.")

        self.arm_trigger(False, self.trigger50r)

    # STATUS METHODS

    def get_status(self, timeout: float = 0.5) -> Dict[str, Any]:
        """
        Get current device status
        
        Args:
            timeout: Timeout in seconds
            
        Returns:
            Dictionary containing device status information
        """
        if not self.port or not self.port.is_open:
            raise RuntimeError("Device not connected. Call open() first.")

        # Send get status command
        self._request_status()

        # Wait for response
        return self._read_status(timeout)

    # PRIVATE/INTERNAL METHODS

    def _send_frame(self, p0: int = 0, p1: int = 0, p2: int = 0, cmd: int = CMD.NONE) -> None:
        """Send a command frame to the device"""
        if not self.port or not self.port.is_open:
            raise RuntimeError("Port not open")

        frame = bytes([p0, p1, p2, cmd, EOF])
        print(f"Sending frame: {list(frame)}")
        self.port.write(frame)
        time.sleep(0.005)  # Small delay to ensure data propagates

    def _sync(self) -> None:
        """Initialize device communication"""
        for _ in range(5):
            self._send_frame(0, 0, 0, CMD.NONE)

    def _convert_connections_to_word(self, connections: Dict[int, int]) -> int:
        """Convert connection state to a 16-bit word for the hardware"""
        word = 0
        if 1 in connections:
            word |= (connections[1] & 0xFF)
        if 2 in connections:
            word |= ((connections[2] & 0xFF) << 8)
        return word

    def _request_status(self) -> None:
        """Request status from device"""
        self._send_frame(0, 0, 0, CMD.GET_STATUS)

    def _read_status(self, timeout: float = 0.5) -> Dict[str, Any]:
        """Read and parse status response from device"""
        deadline = time.time() + timeout

        while time.time() < deadline:
            if self.port and self.port.in_waiting >= 5:
                # Read 5 bytes for status response
                raw_data = self.port.read(5)
                if len(raw_data) == 5:
                    data = list(raw_data)
                    s = data[:4]  # First 4 bytes contain status info
                    return {
                        "raw": s,
                        "version": s[0] & 0x0F,
                        "pll_locked": bool(s[0] & 0x10),
                        "sequence_done": bool(s[0] & 0x20),
                        "hw_revision": (s[0] >> 6) & 0x03,
                        "last_spi_byte": s[1],
                        "seq_addr": s[2],
                        "loop_counter": (s[3] << 8) | data[4],
                    }
            time.sleep(0.01)

        raise RuntimeError("Timeout waiting for status")
