from collections.abc import Callable, Iterable
from datetime import timedelta
from types import TracebackType
from typing import TYPE_CHECKING, Generic, Self, TypeVar
import numpy as np
from pyautd3.autd_error import InvalidDatagramTypeError, KeyAlreadyExistsError
from pyautd3.controller.timer import SpinSleeper, TimerStrategy
from pyautd3.driver.autd3_device import AUTD3
from pyautd3.driver.datagram import Datagram
from pyautd3.driver.firmware.fpga import FPGAState
from pyautd3.driver.firmware_version import FirmwareInfo
from pyautd3.driver.geometry import Device, Geometry
from pyautd3.driver.link import Link, LinkBuilder
from pyautd3.native_methods.autd3capi import ControllerBuilderPtr, ControllerPtr, RuntimePtr
from pyautd3.native_methods.autd3capi import NativeMethods as Base
from pyautd3.native_methods.autd3capi_driver import DatagramPtr, GeometryPtr, HandlePtr, TimerStrategyWrap
from pyautd3.native_methods.structs import FfiFuture, Quaternion, Vector3
from pyautd3.native_methods.utils import _validate_ptr, _validate_status
from pyautd3.emulator import Emulator, Recorder

K = TypeVar("K")
L = TypeVar("L", bound=Link)

class _Builder():
    devices: list[AUTD3]
    async def open_async(self, link: LinkBuilder[L]) -> Controller[L]: ...
    def __init__(self, iterable: Iterable[AUTD3]) -> None: ...
    def _ptr(self, ) -> ControllerBuilderPtr: ...
    def open(self, link: LinkBuilder[L]) -> Controller[L]: ...
    def into_emulator(self, ) -> Emulator: ...
    def with_fallback_parallel_threshold(self, fallback_parallel_threshold: int) -> _Builder: ...
    def with_fallback_timeout(self, fallback_timeout: timedelta) -> _Builder: ...
    def with_send_interval(self, send_interval: timedelta) -> _Builder: ...
    def with_receive_interval(self, receive_interval: timedelta) -> _Builder: ...
    def with_timer_strategy(self, timer_strategy: TimerStrategyWrap) -> _Builder: ...
    @property
    def fallback_parallel_threshold(self) -> int: ...
    @property
    def fallback_timeout(self) -> timedelta: ...
    @property
    def send_interval(self) -> timedelta: ...
    @property
    def receive_interval(self) -> timedelta: ...
    @property
    def timer_strategy(self) -> TimerStrategyWrap: ...

class _GroupGuard(Generic[K]):
    _controller: Controller
    _keys: list[int]
    _datagrams: list[DatagramPtr]
    _keymap: dict[K, int]
    _k: int
    async def send_async(self, ) -> None: ...
    def __init__(self, group_map: Callable[[Device], K | None], controller: Controller) -> None: ...
    def set(self, key: K, d: Datagram | tuple[Datagram, Datagram]) -> Self: ...
    def send(self, ) -> None: ...

class Controller(Geometry, Generic[L]):
    _runtime: RuntimePtr
    _handle: HandlePtr
    _ptr: ControllerPtr
    _link: L
    _disposed: bool
    async def _open_impl_async(self, link_builder: LinkBuilder[L], timeout: timedelta | None = None) -> Controller[L]: ...
    async def firmware_version_async(self, ) -> list[FirmwareInfo]: ...
    async def close_async(self, ) -> None: ...
    async def fpga_state_async(self, ) -> list[FPGAState | None]: ...
    async def send_async(self, d: Datagram | tuple[Datagram, Datagram]) -> None: ...
    def __init__(self, geometry: GeometryPtr, runtime: RuntimePtr, handle: HandlePtr, ptr: ControllerPtr, link: L) -> None: ...
    def __del__(self, ) -> None: ...
    def _dispose(self, ) -> None: ...
    def __enter__(self, ) -> Self: ...
    def __exit__(self, _exc_type: type[BaseException] | None, _exc_value: BaseException | None, _traceback: TracebackType | None) -> None: ...
    def firmware_version(self, ) -> list[FirmwareInfo]: ...
    def close(self, ) -> None: ...
    def fpga_state(self, ) -> list[FPGAState | None]: ...
    def send(self, d: Datagram | tuple[Datagram, Datagram]) -> None: ...
    def group(self, group_map: Callable[[Device], K | None]) -> _GroupGuard: ...
    def tick(self, tick: timedelta) -> None: ...
    @staticmethod
    def builder(iterable: Iterable[AUTD3]) -> _Builder: ...
    @staticmethod
    def _open_impl(builder: ControllerBuilderPtr, link_builder: LinkBuilder[L], timeout: timedelta | None) -> Controller[L]: ...
    @property
    def link(self) -> L: ...
    @property
    def geometry(self) -> Geometry: ...
