"""SWD protocol decoder.

This module provides ARM Serial Wire Debug (SWD) protocol decoding
with DP/AP access detection and ACK/WAIT/FAULT response handling.


Example:
    >>> from oscura.analyzers.protocols.swd import SWDDecoder
    >>> decoder = SWDDecoder()
    >>> for packet in decoder.decode(swclk=swclk, swdio=swdio):
    ...     print(f"Request: {packet.annotations['request_type']}")

References:
    ARM Debug Interface Architecture Specification ADIv5.0 to ADIv5.2
"""

from __future__ import annotations

from enum import Enum
from typing import TYPE_CHECKING

import numpy as np

from oscura.analyzers.protocols.base import (
    AnnotationLevel,
    ChannelDef,
    SyncDecoder,
)
from oscura.core.types import DigitalTrace, ProtocolPacket

if TYPE_CHECKING:
    from collections.abc import Iterator

    from numpy.typing import NDArray


class SWDResponse(Enum):
    """SWD ACK responses."""

    OK = 0b001
    WAIT = 0b010
    FAULT = 0b100


class SWDDecoder(SyncDecoder):
    """SWD protocol decoder.

    Decodes ARM Serial Wire Debug transactions including read/write
    operations to Debug Port (DP) and Access Port (AP) registers.

    Attributes:
        id: "swd"
        name: "SWD"
        channels: [swclk, swdio] (required)

    Example:
        >>> decoder = SWDDecoder()
        >>> for packet in decoder.decode(swclk=swclk, swdio=swdio, sample_rate=10e6):
        ...     print(f"ACK: {packet.annotations['ack']}")
    """

    id = "swd"
    name = "SWD"
    longname = "Serial Wire Debug"
    desc = "ARM Serial Wire Debug protocol decoder"

    channels = [  # noqa: RUF012
        ChannelDef("swclk", "SWCLK", "Serial Wire Clock", required=True),
        ChannelDef("swdio", "SWDIO", "Serial Wire Data I/O", required=True),
    ]

    optional_channels = []  # noqa: RUF012

    options = []  # noqa: RUF012

    annotations = [  # noqa: RUF012
        ("request", "Request packet"),
        ("ack", "ACK response"),
        ("data", "Data phase"),
        ("parity", "Parity bit"),
        ("error", "Error"),
    ]

    def __init__(self) -> None:
        """Initialize SWD decoder."""
        super().__init__()

    def decode(  # type: ignore[override]
        self,
        trace: DigitalTrace | None = None,
        *,
        swclk: NDArray[np.bool_] | None = None,
        swdio: NDArray[np.bool_] | None = None,
        sample_rate: float = 1.0,
    ) -> Iterator[ProtocolPacket]:
        """Decode SWD transactions.

        Args:
            trace: Optional primary trace.
            swclk: Serial Wire Clock signal.
            swdio: Serial Wire Data I/O signal.
            sample_rate: Sample rate in Hz.

        Yields:
            Decoded SWD transactions as ProtocolPacket objects.

        Example:
            >>> decoder = SWDDecoder()
            >>> for pkt in decoder.decode(swclk=swclk, swdio=swdio, sample_rate=1e6):
            ...     print(f"R/W: {'Read' if pkt.annotations['read'] else 'Write'}")
        """
        if swclk is None or swdio is None:
            return

        n_samples = min(len(swclk), len(swdio))
        swclk = swclk[:n_samples]
        swdio = swdio[:n_samples]

        # Find rising edges of SWCLK (data sampled on rising edge)
        rising_edges = np.where(~swclk[:-1] & swclk[1:])[0] + 1

        if len(rising_edges) == 0:
            return

        trans_num = 0
        edge_idx = 0

        while edge_idx < len(rising_edges):
            # Look for start bit (should be 1)
            start_idx = rising_edges[edge_idx]
            if not swdio[start_idx]:
                edge_idx += 1
                continue

            # Check if we have a low period before this start bit
            # (to avoid decoding line reset sequences as transactions)
            if start_idx > 0 and edge_idx > 0:
                prev_edge_idx = rising_edges[edge_idx - 1]
                # Check if there was a low period between previous and current edge
                # If SWDIO stayed high between edges, this might be part of line reset
                swdio_between = swdio[prev_edge_idx:start_idx]
                if len(swdio_between) > 0 and np.all(swdio_between):
                    # SWDIO was high the entire time - likely line reset, skip
                    edge_idx += 1
                    continue

            # Parse request packet (8 bits total)
            # Bit 0: Start (1)
            # Bit 1: APnDP (0=DP, 1=AP)
            # Bit 2: RnW (0=Write, 1=Read)
            # Bits 3-4: A[2:3] (register address bits)
            # Bit 5: Parity (odd parity of bits 1-4)
            # Bit 6: Stop (0)
            # Bit 7: Park (1)

            if edge_idx + 8 > len(rising_edges):
                break

            request_bits = []
            for i in range(8):
                bit_idx = rising_edges[edge_idx + i]
                request_bits.append(1 if swdio[bit_idx] else 0)

            # Extract fields
            start_bit = request_bits[0]
            apndp = request_bits[1]
            rnw = request_bits[2]
            addr_2 = request_bits[3]
            addr_3 = request_bits[4]
            parity = request_bits[5]
            stop_bit = request_bits[6]
            park_bit = request_bits[7]

            # Validate request format
            errors = []
            if start_bit != 1:
                errors.append("Invalid start bit")
            if stop_bit != 0:
                errors.append("Invalid stop bit")
            if park_bit != 1:
                errors.append("Invalid park bit")

            # Check parity (odd parity of APnDP, RnW, A[2:3])
            expected_parity = (apndp + rnw + addr_2 + addr_3) % 2
            if parity != expected_parity:
                errors.append("Request parity error")

            # Construct register address
            register_addr = (addr_3 << 3) | (addr_2 << 2)

            edge_idx += 8

            # Turnaround period (1 clock, host releases SWDIO)
            edge_idx += 1

            # ACK response (3 bits from target)
            if edge_idx + 3 > len(rising_edges):
                break

            ack_bits = []
            for i in range(3):
                bit_idx = rising_edges[edge_idx + i]
                ack_bits.append(1 if swdio[bit_idx] else 0)

            ack_value = (ack_bits[2] << 2) | (ack_bits[1] << 1) | ack_bits[0]

            # Decode ACK
            if ack_value == SWDResponse.OK.value:
                ack_str = "OK"
            elif ack_value == SWDResponse.WAIT.value:
                ack_str = "WAIT"
                errors.append("Target responded with WAIT")
            elif ack_value == SWDResponse.FAULT.value:
                ack_str = "FAULT"
                errors.append("Target responded with FAULT")
            else:
                ack_str = "INVALID"
                errors.append(f"Invalid ACK: 0b{ack_value:03b}")

            edge_idx += 3

            # If ACK is OK, there's a data phase
            data_value = 0
            if ack_value == SWDResponse.OK.value:
                # Turnaround (1 clock)
                edge_idx += 1

                # Data phase (32 bits + 1 parity)
                if edge_idx + 33 > len(rising_edges):
                    break

                data_bits = []
                for i in range(32):
                    bit_idx = rising_edges[edge_idx + i]
                    data_bits.append(1 if swdio[bit_idx] else 0)

                # Convert to value (LSB first)
                for i, bit in enumerate(data_bits):
                    data_value |= bit << i

                # Parity bit
                parity_idx = rising_edges[edge_idx + 32]
                data_parity = 1 if swdio[parity_idx] else 0

                # Check data parity (odd parity)
                expected_data_parity = sum(data_bits) % 2
                if data_parity != expected_data_parity:
                    errors.append("Data parity error")

                edge_idx += 33

            # Calculate timing
            start_time = start_idx / sample_rate
            end_time = rising_edges[min(edge_idx - 1, len(rising_edges) - 1)] / sample_rate

            # Add annotations
            port_type = "AP" if apndp else "DP"
            access_type = "Read" if rnw else "Write"

            self.put_annotation(
                start_time,
                end_time,
                AnnotationLevel.PACKETS,
                f"{port_type} {access_type} @ 0x{register_addr:02X}: {ack_str}",
            )

            # Create packet
            annotations = {
                "transaction_num": trans_num,
                "apndp": "AP" if apndp else "DP",
                "read": bool(rnw),
                "register_addr": register_addr,
                "ack": ack_str,
                "ack_value": ack_value,
            }

            if ack_value == SWDResponse.OK.value:
                annotations["data"] = data_value

            # Encode data as bytes (little-endian)
            data_bytes = (
                data_value.to_bytes(4, "little") if ack_value == SWDResponse.OK.value else b""
            )

            packet = ProtocolPacket(
                timestamp=start_time,
                protocol="swd",
                data=data_bytes,
                annotations=annotations,
                errors=errors,
            )

            yield packet

            trans_num += 1

            # Turnaround and idle
            edge_idx += 1


def decode_swd(
    swclk: NDArray[np.bool_],
    swdio: NDArray[np.bool_],
    sample_rate: float = 1.0,
) -> list[ProtocolPacket]:
    """Convenience function to decode SWD transactions.

    Args:
        swclk: Serial Wire Clock signal.
        swdio: Serial Wire Data I/O signal.
        sample_rate: Sample rate in Hz.

    Returns:
        List of decoded SWD transactions.

    Example:
        >>> packets = decode_swd(swclk, swdio, sample_rate=10e6)
        >>> for pkt in packets:
        ...     print(f"ACK: {pkt.annotations['ack']}")
    """
    decoder = SWDDecoder()
    return list(decoder.decode(swclk=swclk, swdio=swdio, sample_rate=sample_rate))


__all__ = ["SWDDecoder", "SWDResponse", "decode_swd"]
