from typing import Union

from ..bl.delay_execution import DelayExecution
from ..bl.event import Event
from ..interfaces.accelerator_interface import AcceleratorInterface
from ..interfaces.element_interface import ElementInterface


class AcceleratorImpl(AcceleratorInterface):
    def __init__(self, acc, proxy_factory, twiss_calculator, orbit_calculator, *, delay=0.1):
        """
        Initializes the accelerator implementation.
        This manages the accelerator state, interactions, and calculation flow.
        Registers event handlers and calculations.

        Args:
            acc: Accelerator model.
            proxy_factory: Factory for creating element proxies.
            twiss_calculator: Calculator for twiss parameters.
            orbit_calculator: Calculator for orbit calculations.
            delay (float): Delay for calculation executions.
        """
        super().__init__()
        self.acc = acc
        self.proxy_factory = proxy_factory
        self.twiss_calculator = twiss_calculator
        self.orbit_calculator = orbit_calculator

        # Store calculation results
        self.twiss = twiss_calculator.calculate()
        self.orbit = orbit_calculator.calculate()

        # Events for notification
        self.on_new_orbit = Event()
        self.on_new_twiss = Event()

        async def cb_twiss():
            """Internal method to calculate twiss parameters."""
            self.twiss = self.twiss_calculator.calculate()
            await self.on_new_twiss.trigger(self.twiss)

        async def cb_orbit():
            """Internal method to calculate orbit parameters."""
            self.orbit = self.orbit_calculator.calculate()
            await self.on_new_orbit.trigger(self.orbit)

        self.twiss_calculation_delay = DelayExecution(callback=cb_twiss, delay=delay)
        self.orbit_calculation_delay = DelayExecution(callback=cb_orbit, delay=delay)

        self.on_changed_value = Event()

        self.on_orbit_calculation_request = self.orbit_calculation_delay.on_calculation_requested
        self.on_orbit_calculation = self.orbit_calculation_delay.on_calculation
        self.on_twiss_calculation_request = self.twiss_calculation_delay.on_calculation_requested
        self.on_twiss_calculation = self.twiss_calculation_delay.on_calculation

    def set_delay(self, delay: Union[float, None]):
        """How much to delay twiss and orbit calculation after last received command"""
        self.twiss_calculation_delay.set_delay(delay)
        self.orbit_calculation_delay.set_delay(delay)

    async def get_element(self, element_id) -> ElementInterface:
        """
        Retrieves an element and subscribes to its changes.

        Args:
            element_id (str): The element identifier.

        Returns:
            ElementInterface: Proxy object for the element.
        """
        proxy = self.proxy_factory.get(element_id)
        proxy.on_changed_value.subscribe(self.on_changed_value.trigger)

        async def cb(unused):
            await self.orbit_calculation_delay.request_execution()
            await self.twiss_calculation_delay.request_execution()

        proxy.on_update_finished.subscribe(cb)
        return proxy

    def calculate_twiss(self):
        return self.twiss_calculator.calculate()

    def calculate_orbit(self):
        return self.orbit_calculator.calculate()
