from __future__ import annotations
from dataclasses import dataclass
from typing import Iterable, List, Sequence

from zempy.zosapi.core.interop import run_native, ensure_not_none
from zempy.bridge import zemax_exceptions as _exc


@dataclass(frozen=True, slots=True)
class MatrixData:
    """
    Adapter for ZOSAPI.Common.IMatrixData.

    Provides both the raw COM methods (ReadData, WriteData, Get/SetValueAt)
    and higher-level Python helpers to read/write 2D numeric data.
    """
    zosapi: object
    native: object

    # ---------- lifecycle ----------
    @classmethod
    def from_native(cls, zosapi: object, native: object) -> MatrixData:
        if native is None:
            raise ValueError("MatrixData.from_native: native is None")
        return cls(zosapi, native)

    def _ensure_native(self) -> None:
        ensure_not_none(
            self.native,
            what="MatrixData.native",
            exc_type=_exc.ZemaxObjectGone,
        )

    # ---------- properties ----------
    @property
    def IsReadOnly(self) -> bool:
        return bool(run_native(
            "MatrixData.IsReadOnly get",
            lambda: getattr(self.native, "IsReadOnly"),
            ensure=self._ensure_native,
        ))

    @property
    def Rows(self) -> int:
        return int(run_native(
            "MatrixData.Rows get",
            lambda: getattr(self.native, "Rows"),
            ensure=self._ensure_native,
        ))

    @property
    def Cols(self) -> int:
        return int(run_native(
            "MatrixData.Cols get",
            lambda: getattr(self.native, "Cols"),
            ensure=self._ensure_native,
        ))

    @property
    def TotalLength(self) -> int:
        return int(run_native(
            "MatrixData.TotalLength get",
            lambda: getattr(self.native, "TotalLength"),
            ensure=self._ensure_native,
        ))

    @property
    def Data(self) -> List[List[float]]:
        """Return the full 2D data as a nested Python list (row-major)."""
        raw = run_native(
            "MatrixData.Data get",
            lambda: getattr(self.native, "Data"),
            ensure=self._ensure_native,
        )
        # COM returns a 2D array; convert to list of lists
        return [list(map(float, row)) for row in raw]

    @Data.setter
    def Data(self, values: Sequence[Sequence[float]]) -> None:
        if self.IsReadOnly:
            raise PermissionError("MatrixData is read-only; cannot set Data.")
        run_native(
            "MatrixData.Data set",
            lambda: setattr(self.native, "Data",
                            [list(map(float, row)) for row in values]),
            ensure=self._ensure_native,
        )

    # ---------- COM methods ----------
    def ReadData(self, Size: int, Data: List[float]) -> None:
        """Populate `Data` with Size = Rows * Cols elements."""
        run_native(
            "MatrixData.ReadData",
            lambda: self.native.ReadData(int(Size), Data),
            ensure=self._ensure_native,
        )

    def WriteData(self, Size: int, Data: List[float]) -> None:
        """Write `Size` elements from `Data` into the matrix."""
        if self.IsReadOnly:
            raise PermissionError("MatrixData is read-only; cannot write.")
        run_native(
            "MatrixData.WriteData",
            lambda: self.native.WriteData(int(Size), Data),
            ensure=self._ensure_native,
        )

    def GetValueAt(self, Row: int, Col: int) -> float:
        return float(run_native(
            "MatrixData.GetValueAt",
            lambda: self.native.GetValueAt(int(Row), int(Col)),
            ensure=self._ensure_native,
        ))

    def SetValueAt(self, Row: int, Col: int, Value: float) -> None:
        if self.IsReadOnly:
            raise PermissionError("MatrixData is read-only; cannot set element.")
        run_native(
            "MatrixData.SetValueAt",
            lambda: self.native.SetValueAt(int(Row), int(Col), float(Value)),
            ensure=self._ensure_native,
        )

    # ---------- Pythonic helpers ----------
    def to_list(self) -> List[List[float]]:
        """Return a nested list copy of the matrix (row-major)."""
        return self.Data

    def to_flat_list(self) -> List[float]:
        """Return all values as a single flattened list (row-major)."""
        return [x for row in self.Data for x in row]

    def set_from(self, values: Sequence[Sequence[float]]) -> None:
        """Replace entire matrix contents."""
        rows = len(values)
        cols = len(values[0]) if rows else 0
        flat = [float(x) for row in values for x in row]
        self.WriteData(rows * cols, flat)

    def fill(self, value: float) -> None:
        """Fill all elements with the same value."""
        n = self.TotalLength
        self.WriteData(n, [float(value)] * n)

    def __len__(self) -> int:
        return self.Rows

    def __repr__(self) -> str:
        rows, cols = self.Rows, self.Cols
        readonly = "ro" if self.IsReadOnly else "rw"
        data = self.Data
        preview = [row[:3] for row in data[:3]]
        more = "…" if rows > 3 or cols > 3 else ""
        return f"MatrixData<{readonly}>[{rows}x{cols}]: {preview}{more}"
