"""Relay runner."""

import threading
from concurrent.futures import Executor
from typing import Any, Callable, Iterable, Iterator

from abacus_common_logic.utils.logging import task_id_var


class RelayRunner[R]:
    """Process items from an iterator / iterable in parallel using a relay pattern.

    Each task follows this sequence:
    1. Grab the next item from the iterator
    2. Submit a task to process the item after that (i.e. continue the relay chain)
    3. Process the current item

    This approach has the benefits of:
    - Iterator consumption is lazy and thread-safe (each task pulls its own item)
    - The thread pool stays saturated with work (next task is queued before processing starts)

    Items are submitted sequentially but may execute in any order depending on thread
    availability. If any task raises an exception, no new tasks are submitted and the
    first exception is propagated to the caller.
    """

    def __init__(
        self,
        executor: Executor,
        iterable: Iterable[R] | Iterator[R],
        handler: Callable[[R, int], Any],
    ):
        """Initialize the relay runner.

        Args:
            executor: Thread pool where tasks will execute.
            iterable: Source of items to process.
            handler: Function called for each item with (item, zero-based index).
        """
        self._executor = executor
        self._iterator = iter(iterable)
        self._handler = handler
        self._task_id = task_id_var.get()

        self._counter = 0
        self._done = False
        self._event = threading.Event()
        self._first_error: Exception | None = None
        self._has_run = False
        self._length = 0
        self._lock = threading.Lock()

    def run(self) -> None:
        """Start processing and block until all items complete or an error occurs.

        Submits the first task to start the relay chain, then waits for all tasks
        to finish. When any task raises an exception, no new tasks are submitted
        and the first exception is re-raised here.

        Raises:
            RuntimeError: If run() is called more than once.
            Exception: The first exception raised by any task.
        """
        if self._has_run:
            raise RuntimeError('Runner cannot be reused.')
        self._has_run = True

        self._next()

        self._event.wait()

        if self._first_error:
            raise self._first_error

    def _next(self) -> int:
        """Submit the next _step task to the executor if processing should continue.

        New tasks are only submitted if the iterator hasn't been exhausted and no
        error has occurred.

        Returns:
            Total number of tasks submitted so far (including the current one).
        """
        with self._lock:
            if not self._done and not self._first_error:
                self._length += 1
                self._executor.submit(self._step)
            return self._length

    def _step(self) -> None:
        """Execute one step in the relay chain.

        1. Pull the next item from the iterator (thread-safe via _next's lock)
        2. Submit a task for the item after that (continue the relay)
        3. Process the current item with the handler
        4. Track completion and signal when all tasks finish

        Any unexpected exception stops the chain.
        """
        token = task_id_var.set(self._task_id)
        try:
            value = next(self._iterator)
            length = self._next()
            self._handler(value, length - 2)
        except StopIteration:
            with self._lock:
                self._done = True
        except Exception as e:
            with self._lock:
                self._done = True
                if not self._first_error:
                    self._first_error = e
        finally:
            task_id_var.reset(token)
            with self._lock:
                self._counter += 1
                if self._done and self._counter == self._length:
                    self._event.set()