"""Parallel executor using ThreadPoolExecutor."""

from __future__ import annotations

import logging
import threading
from concurrent.futures import Future, ThreadPoolExecutor, as_completed
from typing import TYPE_CHECKING, Any, Callable, TypeVar

from .partitioner import CaseItem

if TYPE_CHECKING:
    from test_runner.common.database_executor import DatabaseExecutor, DatabaseExecutorFactory
    from test_runner.common.models import TestFile

LOGGER = logging.getLogger(__name__)

T = TypeVar("T")


class ParallelExecutor:
    """Run test cases across a thread pool.

    Each worker thread is expected to own its own database connection
    (created via the *case_fn* or *create_connection* callback) so
    there is no shared mutable state.
    """

    def __init__(self, max_workers: int) -> None:
        if max_workers < 1:
            raise ValueError("max_workers must be >= 1")
        self._max_workers = max_workers

    def execute_read_only(
        self,
        cases: list[CaseItem],
        case_fn: Callable[[TestFile, int, list[Any]], T],
    ) -> list[T]:
        """Run read-only cases across the thread pool.

        *case_fn* receives ``(test_file, case_idx, test_case)`` and must
        be thread-safe.  In practice each invocation should create or
        reuse a thread-local database executor so connections are never
        shared across threads.

        Results are returned in the **original case order** (not
        completion order).
        """
        if not cases:
            return []

        results: list[T | None] = [None] * len(cases)

        with ThreadPoolExecutor(max_workers=self._max_workers) as pool:
            future_to_idx: dict[Future[T], int] = {}
            for idx, (test_file, case_idx, test_case) in enumerate(cases):
                fut = pool.submit(case_fn, test_file, case_idx, test_case)
                future_to_idx[fut] = idx

            for fut in as_completed(future_to_idx):
                idx = future_to_idx[fut]
                try:
                    results[idx] = fut.result()
                except Exception:
                    LOGGER.exception(
                        "Worker failed on case %d", idx,
                    )
                    raise

        return results  # type: ignore[return-value]

    def execute_with_factory(
        self,
        cases: list[CaseItem],
        factory: DatabaseExecutorFactory,
        config: Any,
        case_fn: Callable[[DatabaseExecutor, TestFile, int, list[Any]], T],
    ) -> list[T]:
        """Run read-only cases, creating a per-thread executor via *factory*.

        Like :meth:`execute_read_only`, but manages thread-local database
        connections automatically.  *case_fn* receives ``(executor,
        test_file, case_idx, test_case)`` — the executor is created once
        per worker thread and reused for all subsequent cases on that thread.

        All per-thread executors are closed when the pool finishes to
        avoid leaking database connections.
        """
        _tls = threading.local()
        _executors: list[DatabaseExecutor] = []
        _lock = threading.Lock()

        def _wrapped(
            test_file: TestFile, case_idx: int, test_case: list[Any],
        ) -> T:
            if not hasattr(_tls, "executor"):
                _tls.executor = factory.create_executor(config)
                with _lock:
                    _executors.append(_tls.executor)
            return case_fn(_tls.executor, test_file, case_idx, test_case)

        try:
            return self.execute_read_only(cases, _wrapped)
        finally:
            for executor in _executors:
                try:
                    executor.close()
                except Exception:
                    LOGGER.debug("Failed to close worker executor", exc_info=True)
