"""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
    if reference:
        source_mask |= _REFERENCE_FLAG

    for idx, flags in enumerate(datapoint.flags):
        if _extract_stage(int(flags)) != stage:
            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 datapoint_to_raw_point(
    datapoint: VNADatapointPacket,
    settings: VNASweepSettings,
    *,
    num_ports: int,
) -> VNARawPoint:
    """Convert low-level VNADatapoint packet into callback-friendly object."""
    measurements: dict[str, complex] = {}
    stage_by_excited_port = {port: stage for stage, port in enumerate(settings.excited_ports)}

    for excited_port, stage in stage_by_excited_port.items():
        ref = _find_vna_value(datapoint, stage=stage, port_index=excited_port - 1, reference=True)
        for receiver_port in range(1, num_ports + 1):
            measured = _find_vna_value(
                datapoint,
                stage=stage,
                port_index=receiver_port - 1,
                reference=False,
            )
            measurements[f"S{receiver_port}{excited_port}"] = measured / ref

    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)

    stage_by_excited_port = {port: stage for stage, port in enumerate(settings.excited_ports)}
    receiver_limit = min(num_ports, 2)
    trace_names = [
        f"s{receiver_port}{excited_port}"
        for excited_port in settings.excited_ports
        for receiver_port in range(1, receiver_limit + 1)
        if receiver_port <= 2 and excited_port <= 2
    ]
    traces: dict[str, np.ndarray] = {
        name: np.empty(expected_points, dtype=np.complex128) for name in trace_names
    }

    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 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)

        for excited_port, stage in stage_by_excited_port.items():
            ref = _find_vna_value(datapoint, stage=stage, port_index=excited_port - 1, reference=True)
            for receiver_port in range(1, num_ports + 1):
                if receiver_port > 2 or excited_port > 2:
                    # Public SweepResult intentionally exposes 2-port canonical traces.
                    continue
                measured = _find_vna_value(
                    datapoint,
                    stage=stage,
                    port_index=receiver_port - 1,
                    reference=False,
                )
                trace_name = f"s{receiver_port}{excited_port}"
                trace = traces.get(trace_name)
                if trace is not None:
                    trace[idx] = measured / ref

        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=traces, 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