from types import TracebackType

from nortl.core import Const
from nortl.core.operations import to_renderable
from nortl.core.protocols import EngineProto, Renderable, StateProto

__all__ = [
    'ForLoop',
    'WhileLoop',
]


class ForLoop:
    """Context manager to realize a for-loop within a noRTL engine."""

    def __init__(
        self, engine: EngineProto, start: Renderable | int, stop: Renderable | int, step: Renderable | int = 1, counter_width: int = 16
    ) -> None:
        """Initializes the ForLoop context manager.

        Its intention is to emulate the behavior of a for loop in a engine.

        Example:
        ```python
        engine = CoreEngine("my_engine")
        out = engine.define_output("test_output", width=8)

        with ForLoop(engine, 0, 100, 2) as i:
            # code that should be run multiple times
            engine.set(out, i)
        ```

        The context manager cares about the realization of the counter variable and returns it to be used in the loop.
        This means, multiple (non-colliding) loops may use the same count registers.

        As in other programming languages: Only modify the counter variable within the loop if you really know, what you do.
        Generally, it is a very bad idea to write to the loop variable.

        Arguments:
            engine: The CoreEngine instance.
            start: A signal or int representing the start value of the loop counter
            stop: A signal or int representing the final value of the loop counter. The value is non-inclusive.
            step: A signal or int representing the step value of the counter (default: 1)
            counter_width: Width of the counter variable (default: 16). The counter width must fit the stop value.
        """
        self.engine = engine
        self.start = to_renderable(start)
        self.stop = to_renderable(stop)
        self.step = to_renderable(step)

        self.start_state = engine.current_state
        self.final_state = engine.create_state()

        # TODO Validate start and stop values against counter_width, if possible
        # FIXME width must be 1 bit wider than the actual counter!
        self.counter = self.engine.define_scratch(counter_width)
        self.counter_nxt = self.counter + self.step

    def __enter__(self) -> Renderable:
        """Executes the provided function if the condition is met."""
        self.engine.set(self.counter, self.start)
        self.engine.sync()
        self.start_state = self.engine.current_state

        other_state = self.engine.create_state()

        self.engine.jump_if(self.counter == self.stop, self.final_state, other_state)

        self.engine.current_state = other_state

        # TODO consider moving this into the engine itself
        self.engine.scratch_manager.enter_context()

        return self.counter

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        self.engine.sync()
        self.engine.set(self.counter, self.counter_nxt)
        self.engine.jump_if(Const(1), self.start_state)

        # Goto final state
        self.engine.current_state = self.final_state

        self.engine.scratch_manager.exit_context()
        self.counter.release()

        self.engine.sync()


class WhileLoop:
    """Context manager to realize a While loop within a noRTL engine."""

    def __init__(self, engine: EngineProto, condition: Renderable) -> None:
        """Initializes a While construct.

        Arguments:
            engine: The CoreEngine instance.
            condition: The condition of the while loop.
        """
        self.engine = engine
        self.condition = to_renderable(condition)

        self.start_state: StateProto
        self.final_state: StateProto

    def __enter__(self) -> None:
        self.start_state = self.engine.current_state
        self.final_state = self.engine.create_state()

        # Add conditional jump to start state
        self.engine.jump_if(self.condition, self.engine.next_state, self.final_state)

        # Switch to loop content and yield back
        self.engine.current_state = self.engine.next_state

        # TODO consider moving this into the engine itself
        self.engine.scratch_manager.enter_context()

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        if len(self.engine.current_state.transitions) > 0:
            raise RuntimeError('Last state in a While loop must not have any outgoing transitions!')

        self.engine.jump_if(Const(1), self.start_state)

        # Goto final state
        self.engine.current_state = self.final_state

        self.engine.scratch_manager.exit_context()
