from __future__ import annotations

import abc
import contextlib
from collections.abc import Callable, AsyncGenerator
from collections.abc import Mapping
from typing import Protocol

from caqtus.device import DeviceController
from caqtus.device import DeviceName, DeviceConfiguration
from caqtus.device.remote import DeviceProxy
from caqtus.shot_compilation import SequenceContext
from caqtus.types.data import DataLabel, Data
from ._initialize_devices import create_devices
from ._shot_compiler import ShotCompilerProtocol
from ._shot_event_dispatcher import DeviceRunConfig, ShotEventDispatcher
from ._shot_primitives import DeviceParameters
from ..device_manager_extension import DeviceManagerExtensionProtocol


class ShotRunnerProtocol(Protocol):
    """Interface for running a shot."""

    @abc.abstractmethod
    async def run_shot(
        self, shot_parameters: DeviceParameters
    ) -> Mapping[DataLabel, Data]:
        """Run a shot.

        Args:
            shot_parameters: The parameters for the shot.
        """

        ...


type ShotRunnerFactory = Callable[
    [SequenceContext, ShotCompilerProtocol, DeviceManagerExtensionProtocol],
    contextlib.AbstractAsyncContextManager[ShotRunnerProtocol],
]
"""A function called to create a object that can run a shot."""


class ShotRunner(ShotRunnerProtocol):
    def __init__(
        self,
        devices: Mapping[DeviceName, DeviceProxy],
        controller_types: Mapping[DeviceName, type[DeviceController]],
    ):
        if set(devices.keys()) != set(controller_types.keys()):
            raise ValueError("The devices and controller_types must have the same keys")
        self.devices = devices
        self.controller_types = controller_types

    async def run_shot(
        self,
        shot_parameters: DeviceParameters,
    ) -> Mapping[DataLabel, Data]:
        event_dispatcher = ShotEventDispatcher(
            {
                name: DeviceRunConfig(
                    device=self.devices[name],
                    controller_type=self.controller_types[name],
                    parameters=shot_parameters.device_parameters[name],
                )
                for name in self.devices
            }
        )
        return await event_dispatcher.run_shot(shot_parameters.timeout)


@contextlib.asynccontextmanager
async def create_shot_runner(
    sequence_context: SequenceContext,
    shot_compiler: ShotCompilerProtocol,
    device_manager_extension: DeviceManagerExtensionProtocol,
) -> AsyncGenerator[ShotRunner, None]:
    """Creates and acquires resources for running a shot.

    Returns:
        A context manager that yields a shot runner.
    """

    initialization_parameters = shot_compiler.compile_initialization_parameters()

    device_configurations_in_use = {
        name: sequence_context.get_device_configuration(name)
        for name in initialization_parameters
    }

    device_types = {
        name: device_manager_extension.get_device_type(config)
        for name, config in device_configurations_in_use.items()
    }

    async with create_devices(
        initialization_parameters=initialization_parameters,
        device_configs=device_configurations_in_use,
        device_types=device_types,
        device_manager_extension=device_manager_extension,
    ) as devices_in_use:
        shot_runner = _create_shot_runner(
            device_proxies=devices_in_use,
            device_configurations=device_configurations_in_use,
            device_manager_extension=device_manager_extension,
        )

        yield shot_runner


def _create_shot_runner(
    device_proxies: Mapping[DeviceName, DeviceProxy],
    device_configurations: Mapping[DeviceName, DeviceConfiguration],
    device_manager_extension: DeviceManagerExtensionProtocol,
) -> ShotRunner:
    device_controller_types = {
        name: device_manager_extension.get_device_controller_type(
            device_configurations[name]
        )
        for name in device_proxies
    }
    shot_runner = ShotRunner(device_proxies, device_controller_types)
    return shot_runner
