"""Prodigit 3366G High Power DC Electronic Load RS-232 driver.

3366G specs: 600V, 420A, 6-9kW (turbo 1.5x)
Modes: CC (Constant Current), CR (Constant Resistance),
       CV (Constant Voltage), CP (Constant Power)

RS-232 protocol:
  Baud: 9600-115200 (configurable on front panel)
  Data: 8 bits, No parity, 1 stop bit
  Flow: Hardware RTS/CTS
  Terminator: LF (0x0A) or CR+LF

Commands use "Simple Type" format by default.
Query commands append {?} to the setting command.
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field

import serial


@dataclass
class MeasurementResult:
    """Container for a single measurement snapshot."""

    voltage: float
    current: float
    power: float
    timestamp: float = field(default_factory=time.time)

    def __repr__(self) -> str:
        return (
            f"MeasurementResult(V={self.voltage:.4f}, "
            f"I={self.current:.4f}, P={self.power:.4f})"
        )


class Prodigit3366G:
    """Driver for Prodigit 3366G DC Electronic Load over RS-232.

    Args:
        port: Serial port name, e.g. "COM1" or "/dev/ttyUSB0".
        baudrate: Baud rate (must match front panel setting). Default 115200.
        timeout: Read timeout in seconds.
    """

    # 3366G specifications
    MAX_VOLTAGE = 600.0  # V
    MAX_CURRENT_R1 = 420.0  # A (Range 1)
    MAX_CURRENT_R2 = 630.0  # A (turbo 1.5x)
    MAX_POWER_R1 = 6000.0  # W
    MAX_POWER_R2 = 9000.0  # W (turbo)

    # Mode codes returned by MODE?
    MODE_MAP = {0: "CC", 1: "CR", 2: "CV", 3: "CP"}
    MODE_REVERSE = {"CC": 0, "CR": 1, "CV": 2, "CP": 3}

    def __init__(
        self,
        port: str = "COM1",
        baudrate: int = 115200,
        timeout: float = 2.0,
    ) -> None:
        self._port_name = port
        self._ser = serial.Serial(
            port=port,
            baudrate=baudrate,
            bytesize=serial.EIGHTBITS,
            parity=serial.PARITY_NONE,
            stopbits=serial.STOPBITS_ONE,
            rtscts=True,
            timeout=timeout,
        )

    # -- Low-level communication --

    def write(self, command: str) -> None:
        """Send a command to the instrument (appends LF terminator)."""
        self._ser.write((command + "\n").encode("ascii"))

    def query(self, command: str) -> str:
        """Send a query and return the response string."""
        self._ser.reset_input_buffer()
        self.write(command)
        response = self._ser.readline().decode("ascii").strip()
        return response

    def close(self) -> None:
        """Close the serial connection."""
        if self._ser.is_open:
            try:
                self.local()
            except Exception:
                pass
            self._ser.close()

    def __enter__(self) -> Prodigit3366G:
        return self

    def __exit__(self, *exc) -> None:
        self.close()

    # -- System Commands --

    def remote(self) -> None:
        """Enter remote control mode (RS-232)."""
        self.write("REMOTE")

    def local(self) -> None:
        """Exit remote mode, return to local front panel control."""
        self.write("LOCAL")

    def name(self) -> str:
        """Query the model number of the load."""
        return self.query("NAME?")

    def reset(self) -> None:
        """Reset instrument to factory defaults."""
        self.write("*RST")
        time.sleep(1)

    def recall(self, slot: int) -> None:
        """Recall a saved state from memory (1-150)."""
        self.write(f"RECALL {slot}")

    def store(self, slot: int) -> None:
        """Save current state to memory (1-150)."""
        self.write(f"STORE {slot}")

    # -- Load ON/OFF --

    def load_on(self) -> None:
        """Turn on the load (start sinking current)."""
        self.write("LOAD ON")

    def load_off(self) -> None:
        """Turn off the load (stop sinking current)."""
        self.write("LOAD OFF")

    def get_load_state(self) -> bool:
        """Query load on/off state. Returns True if ON."""
        resp = self.query("LOAD?")
        return resp.strip() == "1"

    # -- Mode Selection --

    def set_mode(self, mode: str) -> None:
        """Set load mode: 'CC', 'CR', 'CV', or 'CP'."""
        mode = mode.upper()
        if mode not in self.MODE_REVERSE:
            raise ValueError(f"Invalid mode '{mode}'. Must be CC, CR, CV, or CP.")
        self.write(f"MODE {mode}")

    def get_mode(self) -> str:
        """Query current load mode. Returns 'CC', 'CR', 'CV', or 'CP'."""
        resp = self.query("MODE?")
        code = int(resp.strip())
        return self.MODE_MAP.get(code, f"UNKNOWN({code})")

    # -- Constant Current (CC) --

    def set_cc_current(self, amps: float, level: str = "HIGH") -> None:
        """Set CC mode current level.

        Args:
            amps: Current in Amperes (0-420A R1, 0-630A turbo).
            level: 'HIGH' or 'LOW' level.
        """
        level = level.upper()
        self.write(f"CURR:{level} {amps:.5f}")

    def get_cc_current(self, level: str = "HIGH") -> float:
        """Query CC mode current setting."""
        level = level.upper()
        resp = self.query(f"CURR:{level}?")
        return float(resp)

    # -- Constant Resistance (CR) --

    def set_cr_resistance(self, ohms: float, level: str = "HIGH") -> None:
        """Set CR mode resistance level.

        Args:
            ohms: Resistance in Ohms.
            level: 'HIGH' or 'LOW' level.
        """
        level = level.upper()
        self.write(f"RES:{level} {ohms:.3f}")

    def get_cr_resistance(self, level: str = "HIGH") -> float:
        """Query CR mode resistance setting."""
        level = level.upper()
        resp = self.query(f"RES:{level}?")
        return float(resp)

    # -- Constant Voltage (CV) --

    def set_cv_voltage(self, volts: float, level: str = "HIGH") -> None:
        """Set CV mode voltage level.

        Args:
            volts: Voltage in Volts (0-600V).
            level: 'HIGH' or 'LOW' level.
        """
        level = level.upper()
        self.write(f"VOLT:{level} {volts:.5f}")

    def get_cv_voltage(self, level: str = "HIGH") -> float:
        """Query CV mode voltage setting."""
        level = level.upper()
        resp = self.query(f"VOLT:{level}?")
        return float(resp)

    # -- Constant Power (CP) --

    def set_cp_power(self, watts: float, level: str = "HIGH") -> None:
        """Set CP mode power level.

        Args:
            watts: Power in Watts (0-6000W R1, 0-9000W turbo).
            level: 'HIGH' or 'LOW' level.
        """
        level = level.upper()
        self.write(f"CP:{level} {watts:.5f}")

    def get_cp_power(self, level: str = "HIGH") -> float:
        """Query CP mode power setting."""
        level = level.upper()
        resp = self.query(f"CP:{level}?")
        return float(resp)

    # -- Level (HIGH/LOW toggle for dynamic mode) --

    def set_level(self, level: str) -> None:
        """Set active level: 'HIGH' or 'LOW'."""
        level = level.upper()
        self.write(f"LEV {level}")

    def get_level(self) -> str:
        """Query active level. Returns 'LOW' or 'HIGH'."""
        resp = self.query("LEV?")
        return "HIGH" if resp.strip() == "1" else "LOW"

    # -- Dynamic Mode --

    def set_dynamic(self, on: bool) -> None:
        """Enable/disable dynamic (switching between HIGH/LOW) mode."""
        self.write(f"DYN {'ON' if on else 'OFF'}")

    def get_dynamic(self) -> bool:
        """Query dynamic mode state."""
        resp = self.query("DYN?")
        return resp.strip() == "1"

    # -- Slew Rate --

    def set_rise_slew(self, slew: float) -> None:
        """Set rise slew rate in A/us."""
        self.write(f"RISE {slew}")

    def get_rise_slew(self) -> float:
        """Query rise slew rate."""
        return float(self.query("RISE?"))

    def set_fall_slew(self, slew: float) -> None:
        """Set fall slew rate in A/us."""
        self.write(f"FALL {slew}")

    def get_fall_slew(self) -> float:
        """Query fall slew rate."""
        return float(self.query("FALL?"))

    # -- Dynamic Period --

    def set_period(self, high_ms: float, low_ms: float) -> None:
        """Set dynamic mode period (THIGH and TLOW) in mS."""
        self.write(f"PERD:HIGH {high_ms:.5f}")
        self.write(f"PERD:LOW {low_ms:.5f}")

    def get_period(self, level: str = "HIGH") -> float:
        """Query dynamic period for HIGH or LOW."""
        return float(self.query(f"PERD:{level.upper()}?"))

    # -- Load ON/OFF Voltage --

    def set_load_on_voltage(self, volts: float) -> None:
        """Set the voltage at which load turns ON (0.25-62.5V for 150V range)."""
        self.write(f"LDONV {volts:.5f}")

    def get_load_on_voltage(self) -> float:
        """Query load ON voltage."""
        return float(self.query("LDONV?"))

    def set_load_off_voltage(self, volts: float) -> None:
        """Set the voltage at which load turns OFF."""
        self.write(f"LDOFFV {volts:.5f}")

    def get_load_off_voltage(self) -> float:
        """Query load OFF voltage."""
        return float(self.query("LDOFFV?"))

    # -- Short Circuit --

    def set_short(self, on: bool) -> None:
        """Enable/disable short-circuit test mode."""
        self.write(f"SHOR {'ON' if on else 'OFF'}")

    def get_short(self) -> bool:
        """Query short-circuit mode state."""
        resp = self.query("SHOR?")
        return resp.strip() == "1"

    # -- Voltage Sense --

    def set_sense(self, mode: str) -> None:
        """Set voltage sense mode: 'ON' (V-sense), 'OFF' (input), or 'AUTO'."""
        self.write(f"SENSe {mode.upper()}")

    def get_sense(self) -> str:
        """Query voltage sense mode."""
        resp = self.query("SENSe?")
        m = {"0": "OFF", "1": "ON", "2": "AUTO"}
        return m.get(resp.strip(), resp)

    # -- CC Range --

    def set_cc_range(self, mode: str) -> None:
        """Set CC mode range: 'AUTO' or 'R2' (force Range II)."""
        self.write(f"CC {mode.upper()}")

    # -- Preset Display --

    def set_preset_display(self, on: bool) -> None:
        """Toggle LCD between current setting (ON) and DWM (OFF)."""
        self.write(f"PRES {'ON' if on else 'OFF'}")

    # -- Protection Status --

    def get_protection_status(self) -> dict[str, bool]:
        """Query protection flags.

        Returns dict with keys: OPP, OTP, OVP, OCP.
        """
        resp = self.query("PROT?")
        bits = int(resp.strip())
        return {
            "OPP": bool(bits & 0x01),
            "OTP": bool(bits & 0x02),
            "OVP": bool(bits & 0x04),
            "OCP": bool(bits & 0x08),
        }

    def clear_protection(self) -> None:
        """Clear error/protection flags."""
        self.write("CLR")

    # -- NG Status --

    def get_ng_status(self) -> bool:
        """Query NG (No Good) flag. True = NG condition."""
        resp = self.query("NG?")
        return resp.strip() == "1"

    def set_ng_enable(self, on: bool) -> None:
        """Enable/disable the GO/NG check function."""
        self.write(f"NGENABLE {'ON' if on else 'OFF'}")

    # -- Polarity --

    def set_polarity(self, positive: bool) -> None:
        """Set voltage display polarity: True=POS, False=NEG."""
        self.write(f"POLAR {'POS' if positive else 'NEG'}")

    # -- Test Configuration --

    def set_tconfig(self, mode: str) -> None:
        """Set test configuration: 'NORMAL', 'OCP', 'OPP', or 'SHORT'."""
        self.write(f"TCONFIG {mode.upper()}")

    def get_tconfig(self) -> str:
        """Query test configuration."""
        resp = self.query("TCONFIG?")
        m = {"1": "NORMAL", "2": "OCP", "3": "OPP", "4": "SHORT"}
        return m.get(resp.strip(), resp)

    # -- OCP Test Settings --

    def set_ocp(self, start: float, step: float, stop: float) -> None:
        """Configure OCP test parameters (all in Amps)."""
        self.write(f"OCP:START {start}")
        self.write(f"OCP:STEP {step}")
        self.write(f"OCP:STOP {stop}")

    # -- OPP Test Settings --

    def set_opp(self, start: float, step: float, stop: float) -> None:
        """Configure OPP test parameters (all in Watts)."""
        self.write(f"OPP:START {start}")
        self.write(f"OPP:STEP {step}")
        self.write(f"OPP:STOP {stop}")

    # -- VTH (Threshold Voltage) --

    def set_vth(self, volts: float) -> None:
        """Set threshold voltage for OCP/OPP test."""
        self.write(f"VTH {volts}")

    def get_vth(self) -> float:
        """Query threshold voltage."""
        return float(self.query("VTH?"))

    # -- Short-circuit Test Time --

    def set_stime(self, ms: float) -> None:
        """Set short-circuit test time in ms (0 = continuous)."""
        self.write(f"STIME {ms}")

    def get_stime(self) -> float:
        """Query short-circuit test time."""
        return float(self.query("STIME?"))

    # -- Test Start/Stop --

    def start_test(self) -> None:
        """Start the configured test (OCP/OPP/SHORT)."""
        self.write("START")

    def stop_test(self) -> None:
        """Stop the current test."""
        self.write("STOP")

    def is_testing(self) -> bool:
        """Query whether a test is currently running."""
        resp = self.query("TESTING?")
        return resp.strip() == "1"

    # -- Measurement Averaging --

    def set_averaging(self, count: int) -> None:
        """Set measurement averaging count (1-64)."""
        self.write(f"AVG {count}")

    def get_averaging(self) -> int:
        """Query measurement averaging count."""
        return int(self.query("AVG?"))

    # -- Turbo Mode --

    def set_turbo(self, on: bool) -> None:
        """Enable/disable turbo mode (1.5x current & power rating)."""
        self.write(f"TURBO {'ON' if on else 'OFF'}")

    # -- Limit Settings (GO/NG Judgment) --

    def set_current_limits(self, high: float, low: float) -> None:
        """Set current high/low limits for GO/NG judgment (Amps)."""
        self.write(f"IH {high}")
        self.write(f"IL {low}")

    def set_voltage_limits(self, high: float, low: float) -> None:
        """Set voltage high/low limits for GO/NG judgment (Volts)."""
        self.write(f"VH {high}")
        self.write(f"VL {low}")

    def set_power_limits(self, high: float, low: float) -> None:
        """Set power high/low limits for GO/NG judgment (Watts)."""
        self.write(f"WH {high}")
        self.write(f"WL {low}")

    # -- Battery Discharge Test --

    def set_battery_test(
        self,
        uvp: float | None = None,
        time_s: int | None = None,
        ah: float | None = None,
        wh: float | None = None,
    ) -> None:
        """Configure battery discharge test parameters.

        Args:
            uvp: Under-voltage protection cutoff (Volts).
            time_s: Discharge time limit in seconds (0 = unlimited).
            ah: AH capacity limit (0 = OFF, 0.1-19999.9).
            wh: WH capacity limit (0 = OFF, 0.1-19999.9).
        """
        if uvp is not None:
            self.write(f"BATT:UVP {uvp}")
        if time_s is not None:
            self.write(f"BATT:TIME {time_s}")
        if ah is not None:
            self.write(f"BATT:AH {ah}")
        if wh is not None:
            self.write(f"BATT:WH {wh}")

    def battery_test_start(self) -> None:
        """Start battery discharge test."""
        self.write("BATT:TEST ON")

    def battery_test_stop(self) -> None:
        """Stop battery discharge test."""
        self.write("BATT:TEST OFF")

    def get_battery_results(self) -> dict[str, str]:
        """Read battery test results (AH, WH, time, voltage)."""
        return {
            "AH": self.query("BATT:RAH?"),
            "WH": self.query("BATT:RWH?"),
            "TIME": self.query("BATT:RTIME?"),
            "VOLT": self.query("BATT:RVOLT?"),
        }

    # -- Measurement Queries (THE KEY METHODS) --

    def measure_current(self) -> float:
        """Read the actual load current in Amperes."""
        resp = self.query("MEAS:CURR?")
        return float(resp)

    def measure_voltage(self) -> float:
        """Read the actual load voltage in Volts."""
        resp = self.query("MEAS:VOLT?")
        return float(resp)

    def measure_power(self) -> float:
        """Read the actual load power in Watts."""
        resp = self.query("MEAS:POW?")
        return float(resp)

    def measure_all(self) -> MeasurementResult:
        """Read voltage, current, and power in one call.

        Uses MEAS:VC? which returns voltage,current in one response.
        Power is queried separately.
        """
        vc_resp = self.query("MEAS:VC?")
        parts = vc_resp.split(",")
        voltage = float(parts[0].strip())
        current = float(parts[1].strip())
        power = self.measure_power()
        return MeasurementResult(
            voltage=voltage,
            current=current,
            power=power,
        )

    # -- Send Raw Command --

    def send(self, command: str) -> str | None:
        """Send a raw command. If it contains '?', returns the response."""
        if "?" in command:
            return self.query(command)
        self.write(command)
        return None