"""VISA HiSLIP driver for the PLANAR SN9000 (Иридиум series) multi-port VNA.

The SN9000 hardware connects to a host PC over USB 2.0; the PC runs the SNVNA
application, which exposes the SCPI HiSLIP server. This service connects to
that server, configures a 2x4 virtual switch matrix (ports 1 and 2 stimulate,
ports 3..6 receive), and acquires the full matrix per call in a single
synchronized SCPI round trip.
"""

from __future__ import annotations

from dataclasses import dataclass, field
import time
from typing import Any

import numpy as np
import pyvisa

from python_app.models.dataset_model import ComboKey, SweepCollection, TraceData
from python_app.models.run_config_model import RadarSweepModel

_OUTPUT_PORT_BY_INDEX: tuple[int, ...] = (1, 2)
_INPUT_PORT_BY_INDEX: tuple[int, ...] = (3, 4, 5, 6)
_REFLECTION_BY_OUTPUT: tuple[str, ...] = tuple(
    f"S{port}{port}" for port in _OUTPUT_PORT_BY_INDEX
)
_S_PARAMETER_QUERY_ORDER: tuple[str, ...] = tuple(
    name
    for output_port in _OUTPUT_PORT_BY_INDEX
    for name in (
        f"S{output_port}{output_port}",
        *(f"S{receiver_port}{output_port}" for receiver_port in _INPUT_PORT_BY_INDEX),
    )
)


@dataclass(slots=True)
class Sn9000Service:
    """Acquire the full 2x4 corrected sweep matrix from an SN9000 through VISA HiSLIP."""

    host: str = "127.0.0.1"
    port: int = 4880
    timeout_ms: int = 20_000
    preset_on_open: bool = True
    visa_library: str = "@ivi"
    _resource_manager: pyvisa.ResourceManager | None = field(init=False, default=None, repr=False)
    _instrument: Any | None = field(init=False, default=None, repr=False)
    _settings: RadarSweepModel | None = field(init=False, default=None, repr=False)
    _frequency_hz: np.ndarray | None = field(init=False, default=None, repr=False)

    def __post_init__(self) -> None:
        """Normalize and validate constructor values."""
        self.host = str(self.host).strip()
        if not self.host:
            raise ValueError("SN9000 host must not be empty")
        self.port = int(self.port)
        if self.port <= 0 or self.port > 65535:
            raise ValueError("SN9000 port must be in 1..65535")
        self.timeout_ms = int(self.timeout_ms)
        if self.timeout_ms <= 0:
            raise ValueError("SN9000 timeout_ms must be > 0")
        self.visa_library = str(self.visa_library).strip() or "@ivi"
        if self.visa_library == "@py" or self.visa_library.endswith("@py"):
            raise ValueError("SN9000 requires an IVI/Vendor VISA backend, not pyvisa-py")

    @property
    def resource(self) -> str:
        """Return the assembled VISA HiSLIP resource string."""
        return f"TCPIP0::{self.host}::hislip0,{self.port}::INSTR"

    @property
    def is_open(self) -> bool:
        """Return whether the VISA session is open."""
        return self._instrument is not None

    def open(self) -> None:
        """Open the VISA session and apply stored sweep settings when available."""
        if self._instrument is not None:
            return

        try:
            self._resource_manager = pyvisa.ResourceManager(self.visa_library)
            self._instrument = self._resource_manager.open_resource(self.resource)
            self._instrument.timeout = self.timeout_ms
            self._instrument.write_termination = "\n"
            self._instrument.read_termination = None
            self._instrument.chunk_size = max(int(getattr(self._instrument, "chunk_size", 20_480)), 8 * 1024 * 1024)
            self._instrument.write("*CLS")

            if self._settings is not None:
                self._apply_configuration(self._settings)
        except Exception:
            self.close()
            raise

    def close(self) -> None:
        """Close VISA sessions."""
        if self._instrument is not None:
            self._instrument.close()
            self._instrument = None
        if self._resource_manager is not None:
            self._resource_manager.close()
            self._resource_manager = None

    def __enter__(self) -> Sn9000Service:
        """Open the service and return it."""
        self.open()
        return self

    def __exit__(self, exc_type: object, exc: object, traceback: object) -> None:
        """Close VISA resources on context exit."""
        self.close()

    def recover(self) -> None:
        """Reopen the VISA session after a transient acquisition failure."""
        self.close()
        time.sleep(0.25)
        self.open()

    def query_identity(self) -> str:
        """Read the analyzer identity string."""
        instrument = self._require_instrument()
        return str(instrument.query("*IDN?")).strip()

    def query_system_error(self) -> str:
        """Read one entry from the analyzer SCPI error queue."""
        instrument = self._require_instrument()
        return str(instrument.query("SYST:ERR?")).strip()

    def configure(self, sweep: RadarSweepModel) -> None:
        """Store and apply sweep settings."""
        self._validate_sweep(sweep)
        self._settings = sweep
        self._frequency_hz = None
        if self._instrument is None:
            return
        self._apply_configuration(sweep)

    def read_device_limits(self) -> dict[str, float | int]:
        """Read analyzer limits through SCPI capability/service queries."""
        opened_here = self._instrument is None
        try:
            if opened_here:
                self.open()
            instrument = self._require_instrument()
            return {
                "min_frequency_hz": float(instrument.query("SERV:SWE:FREQ:MIN?")),
                "max_frequency_hz": float(instrument.query("SERV:SWE:FREQ:MAX?")),
                "max_points": int(float(instrument.query("SERV:SWE:POIN?"))),
                "min_power_dbm": float(instrument.query("SERV:SWE:POW:MIN?")),
                "max_power_dbm": float(instrument.query("SERV:SWE:POW:MAX?")),
            }
        finally:
            if opened_here:
                self.close()

    def frequency_axis(self) -> np.ndarray:
        """Return the cached frequency axis from the most recent configuration."""
        if self._frequency_hz is None:
            raise RuntimeError("SN9000 frequency axis is not configured")
        return self._frequency_hz

    def acquire_collection(self, collection_id: int = 1) -> SweepCollection:
        """Acquire one full virtual 2x4 matrix in canonical combo order."""
        if self._settings is None:
            raise RuntimeError("SN9000 service is not configured")
        if self._frequency_hz is None:
            raise RuntimeError("SN9000 frequency axis is not configured")
        points = int(self._settings.points)
        capture_start_ns = time.monotonic_ns()

        s_parameters = self._query_sweep_s_parameters(points)
        traces = self._assemble_traces(s_parameters)

        return SweepCollection(
            collection_id=int(collection_id),
            monotonic_ns=time.monotonic_ns(),
            traces=traces,
            capture_start_ns=capture_start_ns,
            capture_end_ns=time.monotonic_ns(),
        )

    def _apply_configuration(self, sweep: RadarSweepModel) -> None:
        instrument = self._require_instrument()
        if self.preset_on_open:
            instrument.write("SYST:PRES")
            self._expect_opc("*OPC?", context="SN9000 preset")

        instrument.write(f"CALC:PAR:COUN {len(_S_PARAMETER_QUERY_ORDER)}")
        for trace_index, parameter_name in enumerate(_S_PARAMETER_QUERY_ORDER, start=1):
            instrument.write(f"CALC:PAR{trace_index}:DEF {parameter_name}")

        instrument.write("SENS:SWE:TYPE LIN")
        instrument.write(f"SENS:FREQ:STAR {float(sweep.start_hz):.9f}")
        instrument.write(f"SENS:FREQ:STOP {float(sweep.stop_hz):.9f}")
        instrument.write(f"SENS:SWE:POIN {int(sweep.points)}")
        instrument.write("SENS:SWE:POIN:TIME 0")
        instrument.write(f"SENS:BAND {float(sweep.if_bandwidth_hz):.9f}")
        instrument.write("SOUR:POW:PORT:COUP ON")
        instrument.write(f"SOUR:POW {float(sweep.power_dbm):.3f}")
        instrument.write("SENS:AVER OFF")
        instrument.write("FORM:DATA REAL32")
        instrument.write("FORM:BORD SWAP")
        instrument.write("INIT:CONT ON")
        instrument.write("TRIG:SOUR BUS")
        self._expect_opc("*OPC?", context="SN9000 setup")
        self._frequency_hz = self._query_float32_array("SENS:FREQ:DATA?", int(sweep.points))

    def _query_sweep_s_parameters(self, points: int) -> dict[str, np.ndarray]:
        instrument = self._require_instrument()
        data_queries = ";".join(f":SENS:DATA:CORR? {name}" for name in _S_PARAMETER_QUERY_ORDER)
        instrument.write(f"TRIG:SING;*OPC?;{data_queries}")
        opc_token = self._read_ascii_token()
        if opc_token != "1":
            raise RuntimeError(f"SN9000 sweep returned unexpected *OPC? response: {opc_token!r}")

        complex_values: dict[str, np.ndarray] = {}
        for parameter_name in _S_PARAMETER_QUERY_ORDER:
            interleaved = self._read_float32_block(f"SENS:DATA:CORR? {parameter_name}", points * 2)
            complex_values[parameter_name] = self._complex_from_interleaved(interleaved)
        return complex_values

    def _assemble_traces(self, s_parameters: dict[str, np.ndarray]) -> list[TraceData]:
        frequency_hz = self._require_frequency_axis()
        traces: list[TraceData] = []
        for output_position, output_port in enumerate(_OUTPUT_PORT_BY_INDEX):
            reflection = s_parameters[_REFLECTION_BY_OUTPUT[output_position]]
            for input_position, receiver_port in enumerate(_INPUT_PORT_BY_INDEX):
                transmission = s_parameters[f"S{receiver_port}{output_port}"]
                traces.append(
                    TraceData(
                        combo=ComboKey(input=input_position, output=output_position),
                        frequency_hz=frequency_hz,
                        s11=reflection,
                        s21=transmission,
                    )
                )
        return traces

    def _expect_opc(self, command: str, *, context: str) -> None:
        instrument = self._require_instrument()
        response = str(instrument.query(command)).strip()
        if response != "1":
            raise RuntimeError(f"{context} returned unexpected *OPC? response: {response!r}")

    def _query_float32_array(self, command: str, expected_values: int) -> np.ndarray:
        instrument = self._require_instrument()
        instrument.write(command)
        return self._read_float32_block(command, expected_values)

    def _read_ascii_token(self) -> str:
        token = bytearray()
        while True:
            byte = self._read_response_bytes(1)
            if byte in (b";", b"\n", b"\r"):
                if token:
                    return token.decode("ascii").strip()
                continue
            token.extend(byte)
            if len(token) > 64 * 1024:
                raise RuntimeError("SN9000 ASCII response token is too long")

    def _read_float32_block(self, context: str, expected_values: int) -> np.ndarray:
        marker = self._read_response_bytes(1)
        while marker in (b";", b"\n", b"\r"):
            marker = self._read_response_bytes(1)
        if marker != b"#":
            raise RuntimeError(f"SN9000 response for {context!r} does not start with IEEE block marker")

        width = self._read_response_bytes(1)
        if width != b"8":
            raise RuntimeError(f"SN9000 response for {context!r} uses unsupported IEEE block header width")

        payload_size = int(self._read_response_bytes(8).decode("ascii"))
        expected_size = expected_values * np.dtype(np.float32).itemsize
        if payload_size != expected_size:
            raise RuntimeError(
                f"SN9000 response for {context!r} returned {payload_size} payload bytes, "
                f"expected {expected_size}"
            )

        payload = self._read_response_bytes(payload_size)
        array = np.frombuffer(payload, dtype="<f4")
        if array.size != expected_values:
            raise RuntimeError(
                f"SN9000 response for {context!r} returned {array.size} float32 values, "
                f"expected {expected_values}"
            )
        return array

    def _read_response_bytes(self, count: int) -> bytes:
        instrument = self._require_instrument()
        data = instrument.read_bytes(count, break_on_termchar=False)
        if len(data) != count:
            raise RuntimeError(f"SN9000 response ended after {len(data)} bytes, expected {count}")
        return data

    def _require_instrument(self) -> Any:
        if self._instrument is None:
            raise RuntimeError("SN9000 VISA instrument is not open")
        return self._instrument

    def _require_frequency_axis(self) -> np.ndarray:
        if self._frequency_hz is None:
            raise RuntimeError("SN9000 frequency axis is not configured")
        return self._frequency_hz

    @staticmethod
    def _validate_sweep(sweep: RadarSweepModel) -> None:
        if int(sweep.points) < 2:
            raise ValueError("SN9000 sweep points must be >= 2")
        if float(sweep.stop_hz) < float(sweep.start_hz):
            raise ValueError("SN9000 sweep stop_hz must be >= start_hz")
        if float(sweep.if_bandwidth_hz) <= 0.0:
            raise ValueError("SN9000 IF bandwidth must be > 0")

    @staticmethod
    def _complex_from_interleaved(values: np.ndarray) -> np.ndarray:
        if values.size % 2 != 0:
            raise RuntimeError("SN9000 complex trace payload has odd scalar count")
        reshaped = np.asarray(values, dtype=np.float32).reshape((-1, 2))
        return (reshaped[:, 0] + 1j * reshaped[:, 1]).astype(np.complex64)
