import base64
from typing import Any, Dict, List, Optional

import numpy as np
import scipy.fft as fft


def _find_measurement(data: List[Dict], averaging_num: int, data_num: int) -> Optional[Dict]:
    for m in data:
        if m.get("averaging_num") == averaging_num and m.get("data_num") == data_num:
            return m
    return None


def _get_channel_data(measurement: Dict, channel_num: int) -> Optional[str]:
    for ch in measurement.get("channel_data", []):
        if ch.get("channel_num") == channel_num:
            return ch.get("channel_data")
    return None


def _channel_value(value: Any, channel_num: int):
    if isinstance(value, list):
        return value[channel_num]
    return value


def decode_measurement(
    raw_json: Any,
    averaging_num: int = 1,
    data_num: int = 1,
    channel_num: int = 0,
    zero_fill: int = 0,
    voltage_range: float = 0.02,
) -> Dict[str, Any]:
    """
    Decodes raw measurement JSON from the spectrometer.

    raw_json can be:
      - list of measurement dicts (native format)
      - single measurement dict
      - dict with a "data" or "results" key wrapping a list
    """
    if isinstance(raw_json, list):
        data = raw_json
    elif isinstance(raw_json, dict):
        # unwrap common envelope keys
        inner = raw_json.get("data") or raw_json.get("results") or raw_json.get("measurements")
        if isinstance(inner, list):
            data = inner
        else:
            data = [raw_json]
    else:
        raise ValueError(f"Unexpected raw_json type: {type(raw_json)}")

    meas = _find_measurement(data, averaging_num, data_num)
    if meas is None:
        raise ValueError(
            f"Measurement not found (averaging_num={averaging_num}, data_num={data_num})"
        )

    raw_b64 = _get_channel_data(meas, channel_num)
    if raw_b64 is None:
        raise ValueError(f"Channel {channel_num} not found in measurement")

    points_raw = meas.get("measurement_points")
    rate_raw = meas.get("measurement_rate")

    if rate_raw is None:
        raise ValueError("measurement_rate is missing from measurement")

    rate = float(_channel_value(rate_raw, channel_num))
    points = int(_channel_value(points_raw, channel_num)) if points_raw is not None else None

    decoded = base64.b64decode(raw_b64)
    int_data = np.frombuffer(decoded, dtype="<i2")
    if points is not None:
        int_data = int_data[:points]

    scaled = voltage_range * int_data.astype(float) / 32768.0

    segment = scaled - np.mean(scaled)
    n = len(segment)
    if zero_fill > 0:
        segment = np.append(segment, np.zeros(zero_fill))

    transform = fft.rfft(segment) * 2 / n
    freqs = fft.rfftfreq(len(segment), d=1.0 / rate)

    return {
        "measurement_id": meas.get("measurement_id"),
        "points": n,
        "rate_hz": rate,
        "duration_s": n / rate,
        "signal": scaled.tolist(),
        "spectrum": {
            "freqs": freqs.tolist(),
            "amplitudes": np.abs(transform).tolist(),
            "phases": np.angle(transform).tolist(),
        },
    }
