"""LoRa frame synchronization block.

Detects preamble, extracts sync word (NETWORKID), locates SFD,
and outputs aligned data symbols.

Frame structure:
    [Preamble: N upchirps at bin 0]
    [Sync Word: 2 upchirps encoding NETWORKID nibbles]
    [SFD: 2.25 downchirps]
    [Data: encoded payload symbols]
"""

import numpy as np
from numpy.typing import NDArray
from dataclasses import dataclass
from enum import Enum, auto
from typing import Optional, Callable

from .networkid import networkid_from_symbols, sync_word_to_networkid


class FrameSyncState(Enum):
    """State machine states for frame synchronization."""
    SEARCH = auto()      # Searching for preamble
    PREAMBLE = auto()    # Tracking preamble chirps
    SYNC_WORD = auto()   # Capturing sync word symbols
    SFD = auto()         # Detecting SFD downchirps
    DATA = auto()        # Outputting data symbols


@dataclass
class SyncResult:
    """Result from frame synchronization."""
    found: bool                    # Frame detected
    networkid: int                 # Extracted NETWORKID
    cfo_bin: float                 # Carrier frequency offset (bins)
    data_symbols: list[int]        # Aligned data symbol bins
    preamble_count: int            # Number of preamble symbols detected
    sync_word_raw: tuple[int, int] # Raw sync word symbol values


@dataclass
class FrameSyncConfig:
    """Configuration for frame synchronizer."""
    sf: int = 9                    # Spreading factor
    sample_rate: float = 250e3    # Input sample rate
    bw: float = 125e3             # LoRa bandwidth
    preamble_min: int = 4         # Minimum preamble symbols to detect
    expected_networkid: Optional[int] = None  # Filter by NETWORKID (None = any)
    sfd_threshold: float = 0.5    # SFD detection threshold


class FrameSync:
    """Frame synchronization for LoRa signals.

    Performs preamble detection, sync word extraction, SFD detection,
    and symbol alignment for the receiver chain.
    """

    def __init__(self, sf: int = 9, sample_rate: float = 250e3,
                 bw: float = 125e3, preamble_min: int = 4,
                 expected_networkid: Optional[int] = None):
        """Initialize frame synchronizer.

        Args:
            sf: Spreading factor (7-12)
            sample_rate: Input sample rate in Hz
            bw: LoRa signal bandwidth in Hz
            preamble_min: Minimum preamble symbols to consider valid
            expected_networkid: Only accept frames with this NETWORKID (None = all)
        """
        self.config = FrameSyncConfig(
            sf=sf, sample_rate=sample_rate, bw=bw,
            preamble_min=preamble_min, expected_networkid=expected_networkid
        )
        self.sf = sf
        self.N = 1 << sf
        self.sps = int(self.N * sample_rate / bw)

        # Generate reference chirps
        n = np.arange(self.sps)
        phase_up = 2 * np.pi * (n * n / (2 * self.sps))
        self._upchirp = np.exp(1j * phase_up).astype(np.complex64)
        self._downchirp = np.conj(self._upchirp)

        # State
        self.reset()

    def reset(self):
        """Reset synchronizer state."""
        self._state = FrameSyncState.SEARCH
        self._preamble_bins = []
        self._preamble_count = 0
        self._sync_bins = []
        self._data_bins = []
        self._cfo_estimate = 0.0
        self._sfd_count = 0  # Count SFD downchirps (need 2 full ones)

    def _dechirp_and_peak(self, samples: NDArray[np.complex64],
                          use_downchirp: bool = False) -> tuple[int, float]:
        """Dechirp samples and find FFT peak.

        Args:
            samples: One symbol of IQ samples
            use_downchirp: If True, detect downchirp (for SFD)

        Returns:
            Tuple of (peak_bin, peak_magnitude)
        """
        if use_downchirp:
            # For downchirp detection, multiply by upchirp
            dechirped = samples[:self.sps] * self._upchirp
        else:
            # For upchirp detection, multiply by downchirp
            dechirped = samples[:self.sps] * self._downchirp

        spectrum = np.abs(np.fft.fft(dechirped, n=self.N))
        peak_bin = int(np.argmax(spectrum))
        peak_mag = spectrum[peak_bin] / np.mean(spectrum)

        return peak_bin, peak_mag

    def _is_preamble_chirp(self, peak_bin: int, peak_mag: float) -> bool:
        """Check if a chirp looks like part of the preamble.

        Preamble chirps should have:
        - Strong FFT peak (high SNR)
        - Bin very close to 0 (or consistent with CFO)

        The tolerance must be tight enough to distinguish preamble (bin ~0)
        from sync word symbols which can be as low as 8 for RYLR998.
        """
        if peak_mag < 3.0:  # Minimum SNR threshold
            return False

        # If we have a CFO estimate, check against it
        if self._preamble_count > 0:
            expected_bin = int(round(self._cfo_estimate)) % self.N
            # Tight tolerance: must be within 3 bins of expected
            # This distinguishes preamble (bin ~0) from sync word (bin >= 8)
            tolerance = 3
            distance = min(abs(peak_bin - expected_bin),
                           self.N - abs(peak_bin - expected_bin))
            return distance <= tolerance
        else:
            # First preamble chirp - accept bins close to 0 or N-1 (CFO)
            return peak_bin < 4 or peak_bin > self.N - 4

    def _is_downchirp(self, samples: NDArray[np.complex64]) -> tuple[bool, float]:
        """Detect if samples contain a downchirp (SFD).

        Returns:
            Tuple of (is_downchirp, peak_magnitude)
        """
        # Dechirp with upchirp (inverse of normal)
        dechirped = samples[:self.sps] * self._upchirp
        spectrum = np.abs(np.fft.fft(dechirped, n=self.N))
        peak_mag = np.max(spectrum) / np.mean(spectrum)

        # Downchirp produces strong peak when dechirped with upchirp
        return peak_mag > 5.0, peak_mag

    def _estimate_cfo(self) -> float:
        """Estimate CFO from preamble bin measurements."""
        if not self._preamble_bins:
            return 0.0

        # Average the preamble bin values
        bins = np.array(self._preamble_bins)

        # Handle wraparound (bins near N-1 and 0 are close)
        # Convert to complex unit vectors and average
        angles = 2 * np.pi * bins / self.N
        avg_angle = np.angle(np.mean(np.exp(1j * angles)))
        return avg_angle * self.N / (2 * np.pi)

    def process_symbol(self, samples: NDArray[np.complex64]) -> Optional[SyncResult]:
        """Process one symbol's worth of samples.

        This is a streaming interface - call repeatedly with each symbol.

        Args:
            samples: Complex IQ samples (length >= sps)

        Returns:
            SyncResult when a complete frame is detected, None otherwise
        """
        if len(samples) < self.sps:
            return None

        peak_bin, peak_mag = self._dechirp_and_peak(samples)

        if self._state == FrameSyncState.SEARCH:
            # Looking for preamble start
            if self._is_preamble_chirp(peak_bin, peak_mag):
                self._preamble_bins.append(peak_bin)
                self._preamble_count = 1
                self._cfo_estimate = peak_bin
                self._state = FrameSyncState.PREAMBLE

        elif self._state == FrameSyncState.PREAMBLE:
            # Tracking preamble
            if self._is_preamble_chirp(peak_bin, peak_mag):
                self._preamble_bins.append(peak_bin)
                self._preamble_count += 1
                self._cfo_estimate = self._estimate_cfo()
            else:
                # Preamble ended - check if we have enough
                if self._preamble_count >= self.config.preamble_min:
                    # This symbol is first sync word
                    self._sync_bins = [peak_bin]
                    self._state = FrameSyncState.SYNC_WORD
                else:
                    # False alarm, reset
                    self.reset()

        elif self._state == FrameSyncState.SYNC_WORD:
            # Capturing sync word (2 symbols)
            self._sync_bins.append(peak_bin)
            if len(self._sync_bins) >= 2:
                self._state = FrameSyncState.SFD

        elif self._state == FrameSyncState.SFD:
            # Detecting SFD downchirps (2 full + 0.25 fractional)
            is_dc, _ = self._is_downchirp(samples)
            if is_dc:
                self._sfd_count += 1
                # After 2 full downchirps, transition to DATA
                # The 0.25 fractional downchirp is handled in sync_from_samples
                if self._sfd_count >= 2:
                    self._state = FrameSyncState.DATA
            else:
                # Not a downchirp after expecting SFD - could be data already
                # (This handles cases where SFD detection fails)
                self._data_bins.append(peak_bin)
                self._state = FrameSyncState.DATA

        elif self._state == FrameSyncState.DATA:
            self._data_bins.append(peak_bin)

        return None

    def sync_from_samples(self, samples: NDArray[np.complex64],
                          max_data_symbols: int = 100) -> SyncResult:
        """Synchronize and extract frame from a block of samples.

        The key challenge is the fractional SFD: the LoRa SFD is 2.25 downchirps,
        meaning data symbols start at a 0.25 symbol offset after the last full
        downchirp. This method uses a two-phase approach:

        Phase 1: State machine to detect preamble, sync word, and SFD
        Phase 2: Extract data from the correct fractional offset

        Args:
            samples: Complex IQ samples containing a LoRa frame
            max_data_symbols: Maximum data symbols to extract

        Returns:
            SyncResult with extracted frame data
        """
        self.reset()

        n_symbols = len(samples) // self.sps
        sfd_end_symbol = None  # Track when we find the SFD

        # Phase 1: Find frame structure using state machine
        for i in range(n_symbols):
            symbol_samples = samples[i * self.sps:(i + 1) * self.sps]

            prev_state = self._state
            self.process_symbol(symbol_samples)

            # Record when we transition to DATA state
            if prev_state == FrameSyncState.SFD and self._state == FrameSyncState.DATA:
                sfd_end_symbol = i
                break

        # Phase 2: Extract data symbols at correct fractional offset
        # SFD is 2.25 downchirps, so add 0.25 symbol offset after SFD detection
        if sfd_end_symbol is not None:
            # Clear any bins captured during state machine (they're misaligned)
            self._data_bins = []

            # Data starts after 2.25 SFD downchirps from the start of SFD
            # When we transition to DATA at symbol index i, that's the 2nd full downchirp
            # We still need to skip: that symbol (1) + fractional part (0.25) = 1.25
            # So data_start = (sfd_end_symbol + 1.25) * sps
            data_start_sample = int((sfd_end_symbol + 1.25) * self.sps)

            # Extract data symbols from the correct offset
            for i in range(max_data_symbols):
                start = data_start_sample + i * self.sps
                end = start + self.sps
                if end > len(samples):
                    break
                symbol_samples = samples[start:end]
                peak_bin, peak_mag = self._dechirp_and_peak(symbol_samples)
                self._data_bins.append(peak_bin)

        # Build result
        found = len(self._data_bins) > 0 and self._preamble_count >= self.config.preamble_min

        # Extract NETWORKID from sync word
        sync_raw = (0, 0)
        networkid = 0
        if len(self._sync_bins) >= 2:
            # Sync word bins are relative to preamble bin
            cfo_int = int(round(self._cfo_estimate))
            sync_raw = (self._sync_bins[0], self._sync_bins[1])
            # Convert to deltas from preamble
            d1 = (self._sync_bins[0] - cfo_int) % self.N
            d2 = (self._sync_bins[1] - cfo_int) % self.N
            networkid = sync_word_to_networkid((d1, d2))

        # Filter by expected NETWORKID if configured
        if found and self.config.expected_networkid is not None:
            if networkid != self.config.expected_networkid:
                found = False

        return SyncResult(
            found=found,
            networkid=networkid,
            cfo_bin=self._cfo_estimate,
            data_symbols=list(self._data_bins),
            preamble_count=self._preamble_count,
            sync_word_raw=sync_raw,
        )


def frame_sync(sf: int = 9, sample_rate: float = 250e3,
               expected_networkid: Optional[int] = None) -> FrameSync:
    """Factory function for GNU Radio compatibility."""
    return FrameSync(sf=sf, sample_rate=sample_rate,
                     expected_networkid=expected_networkid)


if __name__ == "__main__":
    print("Frame Sync Test")
    print("=" * 50)

    # Generate a test frame
    sf = 9
    N = 1 << sf
    fs = 125e3

    # Simple chirp generation
    n = np.arange(N)

    def upchirp(f_start: int) -> np.ndarray:
        phase = 2 * np.pi * ((f_start * n / N) + (n * n / (2 * N)))
        return np.exp(1j * phase).astype(np.complex64)

    def downchirp() -> np.ndarray:
        return np.conj(upchirp(0))

    # Build test frame
    # Preamble (8 upchirps at bin 0)
    preamble = np.tile(upchirp(0), 8)

    # Sync word for NETWORKID=18 (0x12): nibbles 1, 2 → bins 32, 64
    sync_nibble_hi = 1
    sync_nibble_lo = 2
    sync_bin_1 = sync_nibble_hi * (N // 16)  # 32
    sync_bin_2 = sync_nibble_lo * (N // 16)  # 64
    sync_word = np.concatenate([upchirp(sync_bin_1), upchirp(sync_bin_2)])

    # SFD (2.25 downchirps)
    dc = downchirp()
    sfd = np.concatenate([dc, dc, dc[:N // 4]])

    # Data symbols
    data_bins = [100, 200, 300, 400, 500]
    data = np.concatenate([upchirp(b) for b in data_bins])

    # Complete frame
    frame = np.concatenate([preamble, sync_word, sfd, data])

    print(f"\nTest frame (SF{sf}):")
    print(f"  Preamble: 8 symbols")
    print(f"  Sync word: bins [{sync_bin_1}, {sync_bin_2}] → NETWORKID=18")
    print(f"  SFD: 2.25 downchirps")
    print(f"  Data: {len(data_bins)} symbols, bins {data_bins}")
    print(f"  Total: {len(frame)} samples")

    # Test synchronizer
    sync = FrameSync(sf=sf, sample_rate=fs)
    result = sync.sync_from_samples(frame)

    print(f"\nSync result:")
    print(f"  found: {result.found}")
    print(f"  networkid: {result.networkid}")
    print(f"  cfo_bin: {result.cfo_bin:.2f}")
    print(f"  preamble_count: {result.preamble_count}")
    print(f"  sync_word_raw: {result.sync_word_raw}")
    print(f"  data_symbols: {result.data_symbols}")

    if result.found and result.networkid == 18:
        print("\n✓ Frame sync OK")
    else:
        print("\n✗ Frame sync FAILED")