"""Concurrent multi-project FormsDb search utilities."""

from __future__ import annotations

import asyncio
from typing import Callable, TypeVar, Awaitable

from pycarta.formsdb import FormsDb
from pycarta.auth.agent import CartaAgent

T = TypeVar('T')


class MultiProjectSearch:
    """Concurrent search across multiple FormsDb project instances.

    Encapsulates asyncio.TaskGroup, asyncio.Semaphore, and asyncio.timeout
    patterns for multi-project fan-out with configurable concurrency and
    timeout controls.

    Parameters
    ----------
    credentials : CartaAgent
        Authentication credentials for FormsDb API access.
    project_ids : list[str]
        List of project IDs to search across.
    max_concurrent : int
        Maximum number of concurrent FormsDb instances (Semaphore limit).
        Default: 10. Tune based on Lambda memory and API rate limits.
    per_request_timeout : float | None
        Timeout in seconds for each individual project search.
        Default: 30.0. Set to None to disable per-request timeout.
    total_timeout : float | None
        Timeout in seconds for the entire multi-project search operation.
        Default: 600.0 (10 minutes). Set to None to disable session timeout.
    host : str | None
        Optional FormsDb host override. Default: None (uses credentials.host).

    Example
    -------
    async def get_keys(db: FormsDb) -> dict:
        keys = await db.data.keys()
        return {'project_id': db.project_id, 'count': len(keys)}

    searcher = MultiProjectSearch(
        credentials=credentials,
        project_ids=['proj_1', 'proj_2', 'proj_3'],
        max_concurrent=10,
        per_request_timeout=30.0,
        total_timeout=600.0
    )
    results = await searcher.run(get_keys)
    """

    def __init__(
        self,
        credentials: CartaAgent,
        project_ids: list[str],
        *,
        max_concurrent: int = 10,
        per_request_timeout: float | None = 30.0,
        total_timeout: float | None = 600.0,
        host: str | None = None
    ):
        self.credentials = credentials
        self.project_ids = project_ids
        self.max_concurrent = max_concurrent
        self.per_request_timeout = per_request_timeout
        self.total_timeout = total_timeout
        self.host = host

    async def run(
        self,
        search_fn: Callable[[FormsDb], Awaitable[T]]
    ) -> list[T | Exception]:
        """Execute search_fn across all projects concurrently.

        Uses asyncio.TaskGroup for structured concurrency, asyncio.Semaphore
        for limiting concurrent connections, and asyncio.timeout for both
        per-request and session-level timeout control.

        Per-project errors are captured gracefully and returned as Exception
        objects in the results list, enabling partial success scenarios.

        Session-level timeout (total_timeout) raises TimeoutError when the
        entire operation exceeds the configured limit.

        Parameters
        ----------
        search_fn : Callable[[FormsDb], Awaitable[T]]
            Async function that takes a FormsDb instance and returns a result.
            Called once per project_id with a FormsDb instance connected to
            that project.

        Returns
        -------
        list[T | Exception]
            Results in the same order as project_ids. Each element is either:
            - The successful result from search_fn for that project
            - An Exception object if that project's search failed

        Raises
        ------
        TimeoutError
            If total_timeout is exceeded for the entire search operation.
            Per-request timeouts are captured as exceptions in the results list.

        Examples
        --------
        >>> async def get_form_count(db: FormsDb) -> int:
        ...     keys = await db.data.keys()
        ...     return len(keys)
        >>> searcher = MultiProjectSearch(
        ...     credentials=agent,
        ...     project_ids=['proj_1', 'proj_2'],
        ...     max_concurrent=5
        ... )
        >>> results = await searcher.run(get_form_count)
        >>> # results might be [42, 17] if both succeeded
        >>> # or [42, TimeoutError(...)] if second project timed out
        """
        # Create semaphore for concurrency limiting
        semaphore = asyncio.Semaphore(self.max_concurrent)

        async def _search_project(project_id: str) -> T | Exception:
            """Search a single project with semaphore and timeout control."""
            try:
                async with semaphore:
                    # Wrap in per-request timeout if configured
                    if self.per_request_timeout is not None:
                        async with asyncio.timeout(self.per_request_timeout):
                            async with FormsDb(
                                self.credentials,
                                project_id,
                                host=self.host
                            ) as db:
                                return await search_fn(db)
                    else:
                        # No per-request timeout
                        async with FormsDb(
                            self.credentials,
                            project_id,
                            host=self.host
                        ) as db:
                            return await search_fn(db)
            except Exception as e:
                # Graceful per-project failure - return exception object
                return e

        # Collect tasks in order to preserve project_ids ordering
        tasks = []

        # Wrap TaskGroup in session-level timeout if configured
        if self.total_timeout is not None:
            async with asyncio.timeout(self.total_timeout):
                async with asyncio.TaskGroup() as tg:
                    for project_id in self.project_ids:
                        task = tg.create_task(_search_project(project_id))
                        tasks.append(task)
        else:
            # No session-level timeout
            async with asyncio.TaskGroup() as tg:
                for project_id in self.project_ids:
                    task = tg.create_task(_search_project(project_id))
                    tasks.append(task)

        # Collect results in same order as project_ids
        return [task.result() for task in tasks]
