import logging
from typing import Any, List, Sequence, Union

from p4p.client.thread import Context, Value
from p4p.nt.enum import ntenum
from p4p.nt.ndarray import ntndarray
from p4p.nt.scalar import ntbool, ntfloat, ntint, ntnumericarray, ntstr, ntstringarray

from .pv import PV, PVField

logger = logging.getLogger(__name__)


def unwrap_p4p_structure(res: Any) -> Any:
    """
    Unwraps P4P structures to get clean Python types.
    """
    raw_data = res.raw
    raw_dict = raw_data.todict()
    if isinstance(res, (ntbool, ntint, ntfloat, ntnumericarray, ntstr, ntstringarray)):
        val = raw_dict["value"]
    elif isinstance(res, ntenum):
        val = raw_dict["value"]["index"]
    elif isinstance(res, ntndarray):
        val = raw_dict["value"]
    else:
        logger.warning(
            "Unknown NT type %s; val set to None.",
            type(res).__name__,
        )
        val = None

    return val


def _p4p_value_as_stored_raw(res: Any) -> Any:
    """Normalize p4p get results to stored form; dict when ``todict`` exists."""
    raw_value = getattr(res, "raw", res)
    if hasattr(raw_value, "todict") and callable(getattr(raw_value, "todict", None)):
        return raw_value.todict()
    return raw_value


class BridgeIO:
    """
    Handles synchronous Input/Output operations for Process Variables (PVs) via p4p.

    PVs are read and written by passing PV objects. Each PV's fields (PVField)
    are expanded to full names and batch get/put; pvget mutates each field's
    .val and .raw in-place.
    """

    def __init__(self) -> None:
        """Create a PVA context for synchronous get/put. One instance per control loop."""
        self._ctx = Context("pva")

    @staticmethod
    def _as_sequence(items: Any) -> list[Any]:
        """Normalize a single item or iterable into a list."""
        if isinstance(items, (PV, PVField)):
            return [items]
        return list(items)

    @staticmethod
    def _fields_from(item: Any, expand: bool) -> list[PVField]:
        """Resolve one PV/PVField input to concrete fields."""
        if isinstance(item, PVField):
            return [item]
        if isinstance(item, PV):
            if expand:
                return item.all_fields()
            if not isinstance(item.val_field, PVField):
                raise RuntimeError(
                    f"PV {item.name!r} has invalid VAL descriptor; this violates PV invariants."
                )
            return [item.val_field]
        raise TypeError("Expected PV or PVField input; got " f"{type(item).__name__}.")

    @staticmethod
    def _resolve_fields(items: list[Any], expand: bool, for_put: bool) -> list[PVField]:
        """
        Resolve a list of PV/PVField objects into concrete PVField items.

        When ``expand`` is True, PV entries are expanded to their fields.
        When ``expand`` is False, inputs must already be PVField objects.
        """
        fields = []
        for item in items:
            resolved = BridgeIO._fields_from(item, expand)
            for field in resolved:
                fields.append(field)
        if for_put:
            fields = [field for field in fields if field.val is not None]
        return fields

    def pvget(
        self,
        pvs: Union[PV, PVField, Sequence[PV], Sequence[PVField]],
        expand: bool = True,
        timeout: float = 5.0,
    ) -> None:
        """
        Reads PVs synchronously by expanding each PV's fields and mutates each
        field's .val and .raw in-place.

        Args:
            pvs: A single PV or a list of PVs. Each PV's fields are batch-read by full name.
            timeout: Maximum seconds to wait for the network response.

        On success: each field's .val is set to the unwrapped Python value, .raw to a
        normalized form (``todict()`` when the p4p value provides it, else the object).
        On failure (per-field or catastrophic): that field's .val and .raw are set to None.
        """
        inputs = self._as_sequence(pvs)
        if not inputs:
            return

        expanded = self._resolve_fields(inputs, expand=expand, for_put=False)
        names = [field.pv_name for field in expanded]
        raw_results: List[Union[Value, Exception]] = self._ctx.get(
            names, throw=False, timeout=timeout
        )

        for field, raw in zip(expanded, raw_results):
            if isinstance(raw, Exception):
                field.val = None
                field.raw = None
                logger.error("Error getting PV %s: %s", field.pv_name, raw)
            else:
                field.raw = _p4p_value_as_stored_raw(raw)
                field.val = unwrap_p4p_structure(raw)

        logger.debug("pvget result: %s", inputs)

    def pvput(
        self,
        pvs: Union[PV, PVField, Sequence[PV], Sequence[PVField]],
        expand: bool = True,
        timeout: float = 5.0,
    ) -> None:
        """
        Writes to a batch of PVs synchronously by expanding each PV's fields.
        Uses each field's .name and .val; fields with .val None are skipped.

        Blocks until the server confirms the write.

        Args:
            pvs: A single PV or a list of PVs; each field's .val is written to the field's full .name.
            timeout: Maximum seconds to wait for write confirmation.

        Raises:
            TimeoutError, or other exceptions from the underlying PVA context on failure.

        Note:
            An empty PV list is a no-op (no network call).
        """
        logger.debug("pvput: %s", pvs)

        inputs = self._as_sequence(pvs)
        if not inputs:
            return

        expanded = self._resolve_fields(inputs, expand=expand, for_put=True)
        if not expanded:
            return
        names = [field.pv_name for field in expanded]
        vals = [field.val for field in expanded]

        self._ctx.put(names, vals, wait=True, timeout=timeout)
