"""Parallel search execution helper.

Provides :func:`run_parallel_searches`, a small helper that runs a sequence
of search queries concurrently using :class:`~concurrent.futures.ThreadPoolExecutor`
while preserving the research/Flask contexts expected by worker threads.

The helper is intentionally generic:

* ``search_fn`` is a callable accepting the query string and returning
  whatever per-query payload the caller needs (a list of results, a dict
  with metadata, etc.).  Callers wrap it with context-preserving
  decorators (e.g. :func:`preserve_research_context`) before passing it
  in; this helper stays focused on the concurrency concern.
* ``context_factory`` is an optional zero-arg callable (e.g.
  :func:`~local_deep_research.utilities.threading_utils.thread_context`)
  that returns a **fresh** Flask :class:`~flask.ctx.AppContext` for one
  worker thread.  It is invoked once per query *on the calling thread* and
  the resulting context is pushed into the matching worker via
  :func:`~local_deep_research.utilities.threading_utils.thread_with_app_context`.
  A single :class:`~flask.ctx.AppContext` instance must never be shared
  across concurrent threads — Flask tracks the context-var token stack on
  the instance, so concurrent push/pop from multiple threads raises
  ``ValueError: <Token> was created in a different Context``.  ``None``
  disables context propagation, preserving the behavior of strategies that
  never propagated a context (e.g. focused-iteration, progressive).

Returns a list of ``(query, payload)`` tuples in completion order.
Callers that need a question-keyed dict build it from this list; callers
that only need a flat list of results flatten it directly.
"""

from __future__ import annotations

import concurrent.futures
from typing import Callable, List, Optional, Tuple, TypeVar

from flask.ctx import AppContext
from loguru import logger

from ..utilities.threading_utils import thread_with_app_context

T = TypeVar("T")


def run_parallel_searches(
    queries: List[str],
    search_fn: Callable[[str], T],
    max_workers: Optional[int] = None,
    context_factory: Optional[Callable[[], Optional[AppContext]]] = None,
) -> List[Tuple[str, T]]:
    """Run ``search_fn`` for each query in parallel.

    Args:
        queries: Queries to search.  If empty, returns an empty list
            immediately (and logs a warning).
        search_fn: Callable invoked as ``search_fn(query)`` inside a worker
            thread.  Callers are responsible for wrapping it with any
            context-preserving decorators (e.g.
            :func:`preserve_research_context`) before passing it in, and
            for their own error handling (the callable should never raise
            — return an empty payload on failure instead, matching the
            pre-existing contract of the strategies this was extracted
            from).
        max_workers: Size of the thread pool.  Defaults to ``len(queries)``
            when ``None``, matching the historical behavior of the
            source-based, focused-iteration, and progressive strategies.
        context_factory: Optional zero-arg callable returning a **fresh**
            Flask app context per call (e.g.
            :func:`~local_deep_research.utilities.threading_utils.thread_context`).
            Invoked once per query on the calling thread; each worker gets
            its own context.  Passing a single shared context instead would
            crash under concurrency, so a factory (not a context) is
            required.  ``None`` disables context propagation, preserving the
            behavior of strategies that never propagated a context.

    Returns:
        A list of ``(query, payload)`` tuples in completion order.  Each
        ``payload`` is whatever ``search_fn`` returned for that query.
    """
    if not queries:
        logger.warning("No queries provided for parallel search")
        return []

    if max_workers is None:
        max_workers = len(queries)

    @thread_with_app_context
    def _worker(query: str) -> Tuple[str, T]:
        return (query, search_fn(query))

    with concurrent.futures.ThreadPoolExecutor(
        max_workers=max_workers
    ) as executor:
        # context_factory() is called here on the calling thread, once per
        # query, so every worker receives its own fresh AppContext. Sharing
        # one instance across workers is what raises "Token was created in a
        # different Context".
        futures = [
            executor.submit(
                _worker,
                context_factory() if context_factory is not None else None,
                q,
            )
            for q in queries
        ]
        return [
            future.result()
            for future in concurrent.futures.as_completed(futures)
        ]
