"""
MHS-5200A Function Generator Core Module

Protocol determined through experimental observation.
Serial 57600 8N1, commands start with ':', terminated with '\n'
"""

import serial
import time
from typing import Union
from enum import IntEnum

from .exceptions import ResponseError, CommunicationError


class Waveform(IntEnum):
    """
    Waveform types.
    
    IMPORTANT: Write and read values differ!
    
    Write (:s1wN):
    - Basic: 0-4
    - Arbitrary: 100-115
    
    Read (:r1w) returns:
    - Basic: 00-04
    - Arbitrary: 10-25 (decimal, not hex!)
    
    Example:
    - Write :s1w100 (ARB00) -> Read :r1w10
    - Write :s1w115 (ARB15) -> Read :r1w25
    """
    SINE = 0
    SQUARE = 1
    TRIANGLE = 2
    SAWTOOTH_UP = 3
    SAWTOOTH_DOWN = 4
    ARB00 = 100
    ARB01 = 101
    ARB02 = 102
    ARB03 = 103
    ARB04 = 104
    ARB05 = 105
    ARB06 = 106
    ARB07 = 107
    ARB08 = 108
    ARB09 = 109
    ARB10 = 110
    ARB11 = 111
    ARB12 = 112
    ARB13 = 113
    ARB14 = 114
    ARB15 = 115

    @classmethod
    def from_read_value(cls, read_value: int) -> "Waveform":
        """Convert device read value to Waveform."""
        if 0 <= read_value <= 4:
            write_value = read_value
        elif 10 <= read_value <= 25:
            write_value = 100 + (read_value - 10)
        else:
            raise ValueError(f"Invalid read value: {read_value}")
        return cls(write_value)

    def to_read_value(self) -> int:
        """Convert to device read value."""
        if self.value <= 4:
            return self.value
        else:
            return 10 + (self.value - 100)


class Attenuation(IntEnum):
    """
    Output attenuation setting.
    
    IMPORTANT: Manual has this REVERSED!
    
    Note: Read and write values are identical for this enum.
    
    ATT_20DB (-20dB): Lower output, finer resolution (1mV)
    ATT_0DB  (0dB):   Full output, coarser resolution (10mV)
    
    Amplitude ranges:
    - 0dB mode:   0.15V - 15V,   10mV resolution
    - -20dB mode: 0.015V - 1.5V, 1mV resolution
    """
    ATT_20DB = 0
    ATT_0DB = 1


class SweepMode(IntEnum):
    """Sweep mode setting."""
    LINEAR = 0
    LOG = 1


class MHS5200:
    """
    MHS-5200A DDS Function Generator Control Class
    
    Example:
        gen = MHS5200('/dev/ttyUSB0')
        gen.set_frequency(1, 1000)
        gen.set_amplitude(1, 5.0)
        gen.set_waveform(1, Waveform.SINE)
        gen.set_all_outputs(True)
    """
    
    DEFAULT_TIMEOUT = 1.0
    
    # Arbitrary waveform constants
    ARB_SAMPLES = 2048
    ARB_SLICES = 16
    ARB_SAMPLES_PER_SLICE = 128
    ARB_MAX_VALUE = 4095
    ARB_CENTER = 2048
    
    def __init__(self, port: str, timeout: float = DEFAULT_TIMEOUT):
        """
        Initialize connection to MHS-5200A.
        
        Args:
            port: Serial port (e.g., '/dev/ttyUSB0', 'COM3')
            timeout: Read timeout in seconds
        """
        self._timeout = timeout
        self.ser = serial.Serial(
            port=port,
            baudrate=57600,
            bytesize=serial.EIGHTBITS,
            parity=serial.PARITY_NONE,
            stopbits=serial.STOPBITS_ONE,
            timeout=timeout
        )
        time.sleep(0.1)
        self._clear_buffer()
    
    def _clear_buffer(self) -> None:
        """Clear any pending data in serial buffer."""
        self.ser.write(b':\n') # empty command
        time.sleep(0.05)
        self.ser.read_all()
    
    def _send_command(self, cmd: str) -> str:
        """
        Send command and receive response.
        
        Args:
            cmd: Command without leading ':' and trailing newline
            
        Returns:
            Response string (stripped)
            
        Raises:
            CommunicationError: If no response received
        """
        full_cmd = f':{cmd}\n'.encode()
        self.ser.write(full_cmd)
        
        # Wait for data to arrive
        start = time.time()
        while self.ser.in_waiting == 0:
            if time.time() - start > self._timeout:
                break
            time.sleep(0.01)
        
        time.sleep(0.02)  # Small delay for complete message
        response = self.ser.readline().decode().strip()
        
        if not response:
            raise CommunicationError(f"No response to '{cmd}'")
        
        return response
    
    def _send_command_ok(self, cmd: str) -> None:
        """
        Send command and verify 'ok' response.
        
        Raises:
            ResponseError: If response is not 'ok'
        """
        resp = self._send_command(cmd)
        if resp != 'ok':
            raise ResponseError(f"Expected 'ok', got: '{resp}'")
    
    def _parse_response(self, resp: str, expected_prefix: str) -> str:
        """
        Parse response and extract data portion.
        
        Args:
            resp: Raw response string
            expected_prefix: Expected prefix (e.g., ':r1f')
            
        Returns:
            Data portion of response
            
        Raises:
            ResponseError: If response doesn't match expected format
        """
        if not resp.startswith(expected_prefix):
            raise ResponseError(f"Expected '{expected_prefix}...', got: '{resp}'")
        return resp[len(expected_prefix):]
    
    def _validate_channel(self, channel: int) -> None:
        """Validate channel number."""
        if channel not in (1, 2):
            raise ValueError(f"Invalid channel: {channel}. Must be 1 or 2.")
    
    # =========================================================================
    # Device Information
    # =========================================================================
    
    def get_device_info(self) -> dict:
        """
        Get device information.
        
        Returns:
            Dictionary with 'model' and 'serial' keys
            
        Raises:
            ResponseError: If response format is invalid
        """
        resp = self._send_command('r0c')
        data = self._parse_response(resp, ':r0c')
        
        if len(data) < 5:
            raise ResponseError(f"Device info too short: '{data}'")
        
        return {
            'model': f"MHS-{data[:5]}",
            'serial': data[5:]
        }
    
    def get_model(self) -> str:
        """Get device model name (e.g., 'MHS-5206A')."""
        return self.get_device_info()['model']
    
    def get_serial_number(self) -> str:
        """Get device serial number."""
        return self.get_device_info()['serial']
    
    # =========================================================================
    # Frequency Control
    # =========================================================================
    
    def set_frequency(self, channel: int, freq_hz: float) -> None:
        """
        Set output frequency.
        
        Args:
            channel: 1 or 2
            freq_hz: Frequency in Hz (resolution: 0.01 Hz)
        """
        self._validate_channel(channel)
        value = int(freq_hz * 100)
        self._send_command_ok(f's{channel}f{value}')
    
    def get_frequency(self, channel: int) -> float:
        """
        Get current frequency setting.
        
        Args:
            channel: 1 or 2
            
        Returns:
            Frequency in Hz
        """
        self._validate_channel(channel)
        resp = self._send_command(f'r{channel}f')
        data = self._parse_response(resp, f':r{channel}f')
        return int(data) / 100.0
    
    # =========================================================================
    # Amplitude Control
    # =========================================================================
    
    def set_amplitude(self, channel: int, vpp: float) -> None:
        """
        Set output amplitude.
        
        Args:
            channel: 1 or 2
            vpp: Peak-to-peak voltage
            
        Note:
            Unit depends on current attenuation setting.
        """
        self._validate_channel(channel)
        atten = self.get_attenuation(channel)
        
        if atten == Attenuation.ATT_20DB:
            value = int(vpp * 1000)
        else:
            value = int(vpp * 100)
        
        self._send_command_ok(f's{channel}a{value}')
    
    def get_amplitude(self, channel: int) -> float:
        """
        Get current amplitude setting.
        
        Args:
            channel: 1 or 2
            
        Returns:
            Amplitude in Vpp
        """
        self._validate_channel(channel)
        atten = self.get_attenuation(channel)
        
        resp = self._send_command(f'r{channel}a')
        data = self._parse_response(resp, f':r{channel}a')
        value = int(data)
        
        if atten == Attenuation.ATT_20DB:
            return value / 1000.0
        else:
            return value / 100.0
    
    def set_attenuation(self, channel: int, atten: Attenuation) -> None:
        """
        Set output attenuation.
        
        Args:
            channel: 1 or 2
            atten: Attenuation.ATT_0DB or Attenuation.ATT_20DB
        """
        self._validate_channel(channel)
        self._send_command_ok(f's{channel}y{atten.value}')
    
    def get_attenuation(self, channel: int) -> Attenuation:
        """Get current attenuation setting."""
        self._validate_channel(channel)
        resp = self._send_command(f'r{channel}y')
        data = self._parse_response(resp, f':r{channel}y')
        return Attenuation(int(data))
    
    def set_amplitude_auto(self, channel: int, vpp: float) -> None:
        """
        Set amplitude with automatic attenuation selection.
        
        Args:
            channel: 1 or 2
            vpp: Peak-to-peak voltage
        """
        self._validate_channel(channel)
        
        if vpp <= 1.5:
            self.set_attenuation(channel, Attenuation.ATT_20DB)
        else:
            self.set_attenuation(channel, Attenuation.ATT_0DB)
        
        self.set_amplitude(channel, vpp)
    
    # =========================================================================
    # Waveform Control
    # =========================================================================
    
    def set_waveform(self, channel: int, waveform: Union[Waveform, int]) -> None:
        """
        Set output waveform type.
        
        Args:
            channel: 1 or 2
            waveform: Waveform enum or integer (0-4, 100-115)
        """
        self._validate_channel(channel)
        value = waveform.value if isinstance(waveform, Waveform) else waveform
        self._send_command_ok(f's{channel}w{value}')
    
    def get_waveform(self, channel: int) -> Waveform:
        """
        Get current waveform setting.
        
        Args:
            channel: 1 or 2
            
        Returns:
            Waveform enum
        """
        self._validate_channel(channel)
        resp = self._send_command(f'r{channel}w')
        data = self._parse_response(resp, f':r{channel}w')
        return Waveform.from_read_value(int(data))
    
    def get_waveform_name(self, channel: int) -> str:
        """
        Get current waveform name as string.
        
        Args:
            channel: 1 or 2
            
        Returns:
            Waveform name (e.g., 'SINE', 'SQUARE', 'ARB00')
        """
        return self.get_waveform(channel).name
    
    # =========================================================================
    # Duty Cycle Control
    # =========================================================================
    
    def set_duty_cycle(self, channel: int, duty_percent: float) -> None:
        """
        Set duty cycle (for square wave).
        
        Args:
            channel: 1 or 2
            duty_percent: Duty cycle 0.0 - 99.9 %
        """
        self._validate_channel(channel)
        value = int(duty_percent * 10)
        self._send_command_ok(f's{channel}d{value}')
    
    def get_duty_cycle(self, channel: int) -> float:
        """
        Get current duty cycle setting.
        
        Args:
            channel: 1 or 2
            
        Returns:
            Duty cycle in percent
        """
        self._validate_channel(channel)
        resp = self._send_command(f'r{channel}d')
        data = self._parse_response(resp, f':r{channel}d')
        return int(data) / 10.0
    
    # =========================================================================
    # Offset Control
    # =========================================================================
    
    def set_offset(self, channel: int, offset_percent: float) -> None:
        """
        Set DC offset.
        
        Args:
            channel: 1 or 2
            offset_percent: Offset -120% to +120% (0% = no offset)
        """
        self._validate_channel(channel)
        value = int(offset_percent + 120)
        self._send_command_ok(f's{channel}o{value}')
    
    def get_offset(self, channel: int) -> float:
        """
        Get current DC offset setting.
        
        Args:
            channel: 1 or 2
            
        Returns:
            Offset in percent (-120 to +120)
        """
        self._validate_channel(channel)
        resp = self._send_command(f'r{channel}o')
        data = self._parse_response(resp, f':r{channel}o')
        return int(data) - 120
    
    # =========================================================================
    # Phase Control
    # =========================================================================
    
    def set_phase(self, phase_deg: int) -> None:
        """
        Set phase offset.
        
        Args:
            phase_deg: Phase in degrees (-180 to +180)
        """
        if not -180 <= phase_deg <= 180:
            raise ValueError("phase_deg must be between -180 and +180")
        
        raw_phase = int(phase_deg % 360)
        self._send_command_ok(f's1p{raw_phase}')

    def get_phase(self) -> int:
        """
        Get current phase setting.
        
        Args:
            channel: 1 or 2
            
        Returns:
            Phase in degrees (-180 to +180)
        """
        resp = self._send_command(f'r1p')
        raw_phase = self._parse_response(resp, f':r1p')
        phase_deg = int(raw_phase)
        
        if phase_deg > 180:
            return phase_deg - 360
        return phase_deg
    
    # =========================================================================
    # Output Control
    # =========================================================================
    
    def set_all_outputs(self, enabled: bool) -> None:
        """
        Enable or disable ALL channel outputs.
        
        Args:
            enabled: True to enable, False to disable
            
        Note:
            This controls BOTH channels simultaneously.
            Individual channel control is not available on this device.
        """
        value = 1 if enabled else 0
        self._send_command_ok(f's1b{value}')
    
    def get_all_outputs(self) -> bool:
        """
        Get output enabled state.
        
        Returns:
            True if outputs are enabled
        """
        resp = self._send_command('r1b')
        data = self._parse_response(resp, ':r1b')
        return data == '1'
    
    def outputs_on(self) -> None:
        """Enable all channel outputs."""
        self.set_all_outputs(True)
    
    def outputs_off(self) -> None:
        """Disable all channel outputs."""
        self.set_all_outputs(False)
    
    # =========================================================================
    # Invert Control
    # =========================================================================
    
    def set_invert(self, channel: int, inverted: bool) -> None:
        """
        Set waveform inversion.
        
        Args:
            channel: 1 or 2
            inverted: True to invert waveform
        """
        self._validate_channel(channel)
        ch_code = 'a' if channel == 1 else 'b'
        value = 1 if inverted else 0
        self._send_command_ok(f's{ch_code}b{value}')
    
    def get_invert(self, channel: int) -> bool:
        """
        Get waveform inversion state.
        
        Args:
            channel: 1 or 2
            
        Returns:
            True if waveform is inverted
        """
        self._validate_channel(channel)
        ch_code = 'a' if channel == 1 else 'b'
        resp = self._send_command(f'r{ch_code}b')
        data = self._parse_response(resp, f':r{ch_code}b')
        return data == '1'
    
    # =========================================================================
    # Tracking Mode
    # =========================================================================
    
    def set_tracking(self, enabled: bool) -> None:
        """
        Enable or disable channel tracking mode.
        
        Args:
            enabled: True to enable tracking (CH2 tracks CH1 frequency)
        """
        value = 1 if enabled else 0
        self._send_command_ok(f's3b{value}')
    
    def get_tracking(self) -> bool:
        """
        Get tracking mode state.
        
        Returns:
            True if tracking is enabled
        """
        resp = self._send_command('r3b')
        data = self._parse_response(resp, ':r3b')
        return data == '1'
    
    # =========================================================================
    # Power Amplifier
    # =========================================================================
    
    def set_power_amp(self, enabled: bool) -> None:
        """
        Enable or disable power amplifier (if equipped).
        
        Args:
            enabled: True to enable
        """
        value = 1 if enabled else 0
        self._send_command_ok(f's9b{value}')
    
    def get_power_amp(self) -> bool:
        """
        Get power amplifier state.
        
        Returns:
            True if power amp is enabled
        """
        resp = self._send_command('r9b')
        data = self._parse_response(resp, ':r9b')
        return data == '1'
    
    # =========================================================================
    # Sweep Control
    # =========================================================================
    
    def set_sweep_start_freq(self, freq_hz: float) -> None:
        """Set sweep start frequency."""
        value = int(freq_hz * 100)
        self._send_command_ok(f's3f{value}')
    
    def set_sweep_stop_freq(self, freq_hz: float) -> None:
        """Set sweep stop frequency."""
        value = int(freq_hz * 100)
        self._send_command_ok(f's4f{value}')
    
    def set_sweep_range(self, start_hz: float, stop_hz: float) -> None:
        """
        Set sweep frequency range.
        
        Args:
            start_hz: Start frequency in Hz
            stop_hz: Stop frequency in Hz
        """
        self.set_sweep_start_freq(start_hz)
        self.set_sweep_stop_freq(stop_hz)
    
    def get_sweep_start_freq(self) -> float:
        """Get sweep start frequency in Hz."""
        resp = self._send_command('r3f')
        data = self._parse_response(resp, ':r3f')
        return int(data) / 100.0
    
    def get_sweep_stop_freq(self) -> float:
        """Get sweep stop frequency in Hz."""
        resp = self._send_command('r4f')
        data = self._parse_response(resp, ':r4f')
        return int(data) / 100.0
    
    def set_sweep_time(self, seconds: int) -> None:
        """
        Set sweep time.
        
        Args:
            seconds: Sweep duration in seconds (1-600)
        """
        self._send_command_ok(f's1t{seconds}')
    
    def set_sweep_mode(self, mode: SweepMode) -> None:
        """
        Set sweep mode.
        
        Args:
            mode: SweepMode.LINEAR or SweepMode.LOG
        """
        self._send_command_ok(f's7b{mode.value}')
    
    def get_sweep_mode(self) -> SweepMode:
        """Get sweep mode."""
        resp = self._send_command('r7b')
        data = self._parse_response(resp, ':r7b')
        return SweepMode(int(data))
    
    def set_sweep_channel(self, channel: int) -> None:
        """
        Set sweep target channel.
        
        Args:
            channel: 1 or 2
        """
        self._validate_channel(channel)
        self._send_command_ok(f's2b{channel}')
    
    def set_sweep_enabled(self, enabled: bool) -> None:
        """
        Start or stop sweep.
        
        Args:
            enabled: True to start, False to stop
        """
        value = 1 if enabled else 0
        self._send_command_ok(f's8b{value}')
    
    def get_sweep_enabled(self) -> bool:
        """
        Get sweep running state.
        
        Returns:
            True if sweep is running
        """
        resp = self._send_command('r8b')
        data = self._parse_response(resp, ':r8b')
        return data == '1'
    
    def sweep_on(self) -> None:
        """Start sweep."""
        self.set_sweep_enabled(True)
    
    def sweep_off(self) -> None:
        """Stop sweep."""
        self.set_sweep_enabled(False)
    
    def configure_sweep(
        self,
        start_hz: float,
        stop_hz: float,
        time_sec: int,
        mode: SweepMode = SweepMode.LINEAR,
        channel: int = 1
    ) -> None:
        """
        Configure and start sweep.
        
        Args:
            start_hz: Start frequency in Hz
            stop_hz: Stop frequency in Hz
            time_sec: Sweep duration in seconds
            mode: SweepMode.LINEAR or SweepMode.LOG
            channel: Target channel (1 or 2)
        """
        self.set_sweep_range(start_hz, stop_hz)
        self.set_sweep_time(time_sec)
        self.set_sweep_mode(mode)
        self.set_sweep_channel(channel)
    
    # =========================================================================
    # Arbitrary Waveform
    # =========================================================================
    
    def upload_arbitrary_waveform(self, arb_index: int, samples: list) -> None:
        """
        Upload arbitrary waveform data.
        
        Args:
            arb_index: Arbitrary waveform slot (0-15)
            samples: List of 2048 samples, each 0-4095 (12-bit)
        """
        if not 0 <= arb_index <= 15:
            raise ValueError(f"arb_index must be 0-15, got {arb_index}")
        
        if len(samples) != self.ARB_SAMPLES:
            raise ValueError(f"samples must have {self.ARB_SAMPLES} elements, got {len(samples)}")
        
        arb_hex = format(arb_index, 'X')
        
        for slice_idx in range(self.ARB_SLICES):
            slice_hex = format(slice_idx, 'X')
            start = slice_idx * self.ARB_SAMPLES_PER_SLICE
            end = start + self.ARB_SAMPLES_PER_SLICE
            slice_data = samples[start:end]
            
            slice_data = [max(0, min(self.ARB_MAX_VALUE, int(v))) for v in slice_data]
            data_str = ','.join(str(v) for v in slice_data)
            
            self._send_command_ok(f'a{arb_hex}{slice_hex}{data_str}')
    
    def download_arbitrary_waveform(self, arb_index: int) -> list:
        """
        Download arbitrary waveform data from device.
        
        Args:
            arb_index: Arbitrary waveform slot (0-15)
            
        Returns:
            List of 2048 samples (12-bit values 0-4095)
        """
        if not 0 <= arb_index <= 15:
            raise ValueError(f"arb_index must be 0-15, got {arb_index}")
        
        arb_hex = format(arb_index, 'X')
        samples = []
        
        for slice_idx in range(self.ARB_SLICES):
            slice_hex = format(slice_idx, 'X')
            
            self.ser.write(f':b{arb_hex}{slice_hex}\n'.encode())
            time.sleep(0.1)
            
            response = b''
            while True:
                chunk = self.ser.read(1024)
                if not chunk:
                    break
                response += chunk
                if len(chunk) < 1024:
                    break
            
            response = response.decode().strip()
            prefix = f':b{arb_hex}{slice_hex}'
            data = self._parse_response(response, prefix)
            values = [int(v) for v in data.split(',') if v.strip()]
            samples.extend(values)
        
        return samples
    
    # =========================================================================
    # Bulk Read
    # =========================================================================
    
    def get_channel_status(self, channel: int) -> dict:
        """
        Get all parameters for a channel.
        
        Args:
            channel: 1 or 2
            
        Returns:
            Dictionary with all channel parameters
        """
        self._validate_channel(channel)
        return {
            'channel': channel,
            'frequency_hz': self.get_frequency(channel),
            'amplitude_vpp': self.get_amplitude(channel),
            'attenuation': self.get_attenuation(channel).name,
            'waveform': self.get_waveform_name(channel),
            'duty_cycle_percent': self.get_duty_cycle(channel),
            'offset_percent': self.get_offset(channel),
            'inverted': self.get_invert(channel),
        }
    
    def get_all_status(self) -> dict:
        """Get status of all channels and global settings."""
        return {
            'device': self.get_device_info(),
            'channel1': self.get_channel_status(1),
            'channel2': self.get_channel_status(2),
            'output_enabled': self.get_all_outputs(),
            'phase_deg': self.get_phase(),
            'tracking': self.get_tracking(),
            'power_amp': self.get_power_amp(),
        }
    
    # =========================================================================
    # Utility
    # =========================================================================
    
    def raw_command(self, cmd: str) -> str:
        """
        Send raw command for debugging.
        
        Args:
            cmd: Command string (with or without leading ':')
            
        Returns:
            Raw response string
        """
        if cmd.startswith(':'):
            cmd = cmd[1:]
        return self._send_command(cmd)
    
    def close(self) -> None:
        """Close serial connection."""
        if self.ser and self.ser.is_open:
            self.ser.close()
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
        return False