"""Oscilloscope SCPI policy for SRAS acquisition.

All the Tektronix-specific instrument setup the scan depends on, in one
Qt-free place: per-channel display/coupling config, trigger programming,
the background average, and per-row FastFrame transfer.
"""
from __future__ import annotations

import logging
import time
from dataclasses import dataclass

import numpy as np

logger = logging.getLogger(__name__)

SAMPLE_RATE_HZ = 6.25e9    # 6.25 GS/s → 160 ps/sample
TRIG_LEVEL_V = 0.500
BACKGROUND_AVERAGES = 1024
BACKGROUND_TIMEOUT_S = 60.0


@dataclass(frozen=True)
class ChannelProfile:
    """Display/input configuration for one scope channel."""
    label: str
    scale_v_div: float
    position_div: float
    termination_ohm: int
    coupling: str
    bandwidth_hz: float


# Standard SRAS front-end configuration.
SRAS_CHANNELS = {
    1: ChannelProfile("RF Acoustic Packet", 0.07, 0.0, 50, "DC", 250e6),
    2: ChannelProfile("Trigger Signal", 0.5, -2.72, 1_000_000, "DC", 20e6),
    3: ChannelProfile("Max Vel Gate", 1.0, -2.72, 1_000_000, "DC", 20e6),
    4: ChannelProfile("Bias - B", 0.1, -2.72, 1_000_000, "DC", 20e6),
}


def configure_channels(scope, profiles=None) -> None:
    """Apply the standard SRAS channel configuration."""
    profiles = profiles if profiles is not None else SRAS_CHANNELS
    for ch, p in profiles.items():
        scope.write(f"SELect:CH{ch} ON")
        scope.set_channel_label_name(ch, p.label)
        scope.set_channel_scale(ch, p.scale_v_div)
        scope.set_channel_position(ch, p.position_div)
        scope.set_channel_termination(ch, p.termination_ohm)
        scope.set_channel_coupling(ch, p.coupling)
        scope.set_channel_bandwidth(ch, p.bandwidth_hz)


def configure_acquisition(scope) -> int:
    """Program the edge trigger and timebase; returns samples per frame.

    Edge trigger on the rising edge of CH2 (laser pulse).  FastFrame stays
    off here so the background capture runs as a single record.
    """
    scope.write("TRIGger:A:TYPe EDGE")
    scope.set_trigger_source(2)
    scope.set_trigger_slope("RISE")
    scope.set_trigger_level(2, TRIG_LEVEL_V)
    scope.set_trigger_mode("NORMAL")       # wait for trigger (don't auto-sweep)
    scope.set_acquire_mode("SAMPLE")
    scope.set_fastframe_state(False)
    scope.set_sample_rate(SAMPLE_RATE_HZ)
    scope.write("HORizontal:POSition 30")  # 10 % trigger offset
    time.sleep(0.3)                        # let the timebase settle before reading back
    return scope.get_record_length()


def read_preambles(scope, channels) -> list[str]:
    """Snapshot WFMOutpre per channel (captures YMULT/YOFF/YZERO)."""
    preambles = []
    for ch in channels:
        scope.set_data_source(ch)
        preambles.append(scope.query_wfmoutpre())
    return preambles


def capture_background(scope, should_abort=lambda: False,
                       on_status=lambda msg: None) -> bytes:
    """Capture one CH1 waveform averaged over BACKGROUND_AVERAGES shots.

    The scope auto-stops after the sequence; poll ACQuire:STATE until it
    does rather than assuming a duration.
    """
    on_status(f"Capturing background waveform ({BACKGROUND_AVERAGES}-average) …")
    scope.set_acquire_mode("AVERAGE")
    scope.write(f"ACQuire:NUMAVg {BACKGROUND_AVERAGES}")
    scope.write("ACQuire:STOPAfter SEQuence")
    scope.set_data_source(1)
    scope.write("ACQuire:STATE RUN")

    deadline = time.time() + BACKGROUND_TIMEOUT_S
    while time.time() < deadline:
        if should_abort():
            break
        if scope.query("ACQuire:STATE?").strip() == "0":
            break
        time.sleep(0.25)
    else:
        scope.write("ACQuire:STATE STOP")
        on_status("Warning: background average timed out; stopping early.")
    time.sleep(0.1)
    return scope.transfer_curve()


def configure_scan_trigger(scope) -> None:
    """Switch to the scan-time logic-AND trigger (CH2 HIGH AND CH3 HIGH).

    CH3 is the BBD202 TRIGOUT_MAXV gate, so frames only accumulate while the
    stage is at full scan velocity.
    """
    scope.write("ACQuire:STOPAfter RUNSTop")
    scope.set_acquire_mode("SAMPLE")
    scope.set_fastframe_state(True)

    scope.write("TRIGger:A:TYPe LOGIc")
    scope.write("TRIGger:A:LOGIc:FUNCtion AND")
    scope.set_trigger_level(2, TRIG_LEVEL_V)
    scope.set_trigger_level(3, TRIG_LEVEL_V)
    scope.write("TRIGger:A:LOGICPattern:CH2 HIGH")
    scope.write("TRIGger:A:LOGICPattern:CH3 HIGH")
    time.sleep(0.2)


def arm_row(scope) -> None:
    """Start acquisition for one scan row."""
    scope.write("ACQuire:STATE RUN")
    time.sleep(0.05)


def finish_row(scope) -> None:
    """Wait for trailing frames, then stop acquisition."""
    time.sleep(0.2)
    scope.write("ACQuire:STATE STOP")


def frames_acquired(scope) -> int:
    return int(scope.query("ACQuire:NUMFRAMESACQuired?"))


def transfer_channel(scope, ch: int) -> list[bytes]:
    """Fetch one channel's FastFrame block as raw int8 frames."""
    scope.set_data_source(ch)
    return scope.transfer_fastframe(parse=False)


def frame_means(waveforms: list[bytes]) -> list[float]:
    """Per-frame DC mean of a raw int8 FastFrame block.

    numpy over the joined buffer: the per-frame struct.unpack this replaces
    allocated a tuple of Python ints per frame (~16k frames per row).
    """
    if not waveforms:
        return []
    n = len(waveforms[0])
    if n == 0 or any(len(w) != n for w in waveforms):
        # Ragged block (shouldn't happen) — fall back to per-frame means.
        return [float(np.frombuffer(w, dtype=np.int8).mean()) if len(w) else 0.0
                for w in waveforms]
    block = np.frombuffer(b"".join(waveforms), dtype=np.int8).reshape(len(waveforms), n)
    return block.mean(axis=1, dtype=np.float32).tolist()