"""Helpers to assemble high-level sweep results from packet streams."""

from __future__ import annotations

import logging
import math
from typing import Iterable

import numpy as np

from ..exceptions import IncompleteSweepError, ParseError
from ..enums import SweepKind
from ..models import VNADatapointPacket, VNARawPoint, VNASweepSettings, SweepResult

logger = logging.getLogger(__name__)


_REFERENCE_FLAG = 0x10


def _extract_stage(flags: int) -> int:
    """Return stage index encoded inside datapoint flag byte."""
    return flags >> 5


def _find_vna_value(
    datapoint: VNADatapointPacket,
    *,
    stage: int,
    port_index: int,
    reference: bool,
) -> complex:
    """Find one complex receiver value for stage/port/ref tuple."""
    source_mask = 1 << port_index

    for idx, flags in enumerate(datapoint.flags):
        if _extract_stage(int(flags)) != stage:
            continue
        if bool(int(flags) & _REFERENCE_FLAG) != reference:
            continue
        if (int(flags) & source_mask) != source_mask:
            continue
        return complex(float(datapoint.real[idx]), float(datapoint.imag[idx]))

    kind = "reference" if reference else "receiver"
    raise ParseError(
        f"Missing {kind} value for stage={stage}, port_index={port_index}, point={datapoint.point_number}"
    )


def _s21(datapoint: VNADatapointPacket) -> complex:
    reference = _find_vna_value(datapoint, stage=0, port_index=0, reference=True)
    receiver = _find_vna_value(datapoint, stage=0, port_index=1, reference=False)
    if not np.isfinite(reference) or abs(reference) <= 1e-12 or not np.isfinite(receiver):
        raise ParseError(f"Invalid receiver/reference at point {datapoint.point_number}")
    return receiver / reference


def datapoint_to_raw_point(
    datapoint: VNADatapointPacket,
    settings: VNASweepSettings,
    *,
    num_ports: int,
) -> VNARawPoint:
    """Convert low-level VNADatapoint packet into callback-friendly object."""
    if num_ports < 2:
        raise ParseError("S21 acquisition requires two receiver ports")
    measurements = {"S21": _s21(datapoint)}

    zero_span = settings.f_start_hz == settings.f_stop_hz and math.isclose(
        settings.power_start_dbm,
        settings.power_stop_dbm,
        rel_tol=0.0,
        abs_tol=0.0,
    )

    if zero_span:
        return VNARawPoint(
            point_number=datapoint.point_number,
            frequency_hz=None,
            time_s=float(datapoint.frequency_or_time) * 1e-6,
            power_dbm=None,
            measurements=measurements,
        )

    return VNARawPoint(
        point_number=datapoint.point_number,
        frequency_hz=float(datapoint.frequency_or_time),
        time_s=None,
        power_dbm=float(datapoint.cdbm) / 100.0,
        measurements=measurements,
    )


def assemble_vna_sweep(
    datapoints: Iterable[VNADatapointPacket],
    settings: VNASweepSettings,
    *,
    num_ports: int,
    expected_points: int,
) -> SweepResult:
    """Build a complete VNA sweep from raw datapoints."""
    if expected_points <= 0:
        raise ValueError("expected_points must be > 0")

    x = np.empty(expected_points, dtype=np.float64)
    seen = np.zeros(expected_points, dtype=bool)

    if num_ports < 2:
        raise ParseError("S21 acquisition requires two receiver ports")
    s21 = np.empty(expected_points, dtype=np.complex128)

    zero_span = settings.f_start_hz == settings.f_stop_hz and math.isclose(
        settings.power_start_dbm,
        settings.power_stop_dbm,
        rel_tol=0.0,
        abs_tol=0.0,
    )
    power_sweep = settings.kind == SweepKind.POWER

    for datapoint in datapoints:
        idx = datapoint.point_number
        if not 0 <= idx < expected_points:
            raise ParseError(f"Received out-of-range point index {idx}, expected < {expected_points}")

        if zero_span:
            x[idx] = float(datapoint.frequency_or_time) * 1e-6
        elif power_sweep:
            x[idx] = float(datapoint.cdbm) / 100.0
        else:
            x[idx] = float(datapoint.frequency_or_time)

        s21[idx] = _s21(datapoint)

        seen[idx] = True

    missing = np.flatnonzero(~seen)
    if missing.size > 0:
        raise IncompleteSweepError(
            f"Sweep is incomplete: received {int(seen.sum())}/{expected_points} points"
        )

    x_label = "frequency_hz"
    if zero_span:
        x_label = "time_s"
    elif power_sweep:
        x_label = "power_dbm"

    result = SweepResult(x=x, traces={"s21": s21}, x_label=x_label)
    logger.debug(
        "Assembled VNA sweep: points=%d x_label=%s traces=%s",
        len(result.x),
        result.x_label,
        sorted(result.traces.keys()),
    )
    return result