"""
RBSP Parallel Execution Module.

Provides true parallelism for CPU-bound RBSP operations (tokenization,
hashing, scoring) via multiprocessing, and threading for I/O-bound work
(file reading).  Bypasses Python's GIL for multi-core utilization.
"""

from __future__ import annotations

import heapq
import multiprocessing
import os
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
from typing import Any, Callable, Dict, List, Optional, Tuple


# ---------------------------------------------------------------------------
# Utility
# ---------------------------------------------------------------------------

def chunk_list(lst: list, n: int) -> list[list]:
    """Split *lst* into *n* roughly equal chunks.

    Returns up to *n* non-empty sub-lists.  If *lst* has fewer than *n*
    elements, the number of chunks equals ``len(lst)``.
    """
    if n <= 0:
        raise ValueError("n must be positive")
    if not lst:
        return []
    k, remainder = divmod(len(lst), n)
    chunks: list[list] = []
    start = 0
    for i in range(n):
        size = k + (1 if i < remainder else 0)
        if size == 0:
            break
        chunks.append(lst[start : start + size])
        start += size
    return chunks


# ---------------------------------------------------------------------------
# Worker functions (module-level for pickling by multiprocessing)
# ---------------------------------------------------------------------------

def _index_file_batch(
    batch: List[str],
    tokenize_fn: Callable[[str], List[str]],
    hash_fn: Callable[[str], str],
) -> List[Dict[str, Any]]:
    """Worker: tokenize and hash a batch of files.

    Reads each file, applies *tokenize_fn* and *hash_fn*, and returns a
    list of dicts with keys ``path``, ``tokens``, ``content_hash``.
    Files that cannot be read are silently skipped.
    """
    results: List[Dict[str, Any]] = []
    for path in batch:
        try:
            with open(path, "r", encoding="utf-8", errors="replace") as fh:
                content = fh.read()
        except OSError:
            continue
        tokens = tokenize_fn(content)
        content_hash = hash_fn(content)
        results.append({
            "path": path,
            "tokens": tokens,
            "content_hash": content_hash,
        })
    return results


def _score_batch(
    candidates: List[Tuple],
    score_fn: Callable[..., float],
) -> List[Tuple[str, float]]:
    """Worker: score a batch of candidates.

    Each element in *candidates* is unpacked as positional args to
    *score_fn*.  The first element is treated as the doc_id.
    Returns ``(doc_id, score)`` pairs.
    """
    results: List[Tuple[str, float]] = []
    for candidate in candidates:
        doc_id = candidate[0]
        score = score_fn(*candidate)
        results.append((doc_id, score))
    return results


def _search_shard(
    query_fn: Callable[..., list],
    shard: Any,
) -> list:
    """Worker: run *query_fn* on a single shard."""
    return query_fn(shard)


# ---------------------------------------------------------------------------
# ParallelIndexer
# ---------------------------------------------------------------------------

class ParallelIndexer:
    """Parallel file tokenization and hashing using multiprocessing.

    Distributes file I/O + tokenization across worker processes to bypass
    the GIL for CPU-bound stemming / hashing work.
    """

    def __init__(self, num_workers: Optional[int] = None) -> None:
        self.num_workers = num_workers or os.cpu_count() or 1

    def index_files(
        self,
        file_paths: List[str],
        tokenize_fn: Callable[[str], List[str]],
        hash_fn: Callable[[str], str],
    ) -> List[Dict[str, Any]]:
        """Tokenize and hash *file_paths* in parallel.

        Each worker processes a batch of files and returns token lists
        plus content hashes.  Results are gathered and returned as a
        flat list of dicts with keys ``path``, ``tokens``, ``content_hash``.
        """
        if not file_paths:
            return []

        batches = chunk_list(file_paths, self.num_workers)

        # Use ProcessPoolExecutor for true parallelism on CPU-bound work
        all_results: List[Dict[str, Any]] = []
        with ProcessPoolExecutor(max_workers=self.num_workers) as pool:
            futures = [
                pool.submit(_index_file_batch, batch, tokenize_fn, hash_fn)
                for batch in batches
            ]
            for future in futures:
                all_results.extend(future.result())

        return all_results


# ---------------------------------------------------------------------------
# ParallelScorer
# ---------------------------------------------------------------------------

class ParallelScorer:
    """Score candidate documents in parallel across worker processes.

    Partitions the candidate list across *num_workers* processes so that
    CPU-intensive BM25 / similarity scoring runs without GIL contention.
    """

    def __init__(self, num_workers: Optional[int] = None) -> None:
        self.num_workers = num_workers or os.cpu_count() or 1

    def score_batch(
        self,
        candidates: List[Tuple],
        score_fn: Callable[..., float],
    ) -> List[Tuple[str, float]]:
        """Score *candidates* in parallel.

        Each candidate tuple is unpacked as positional arguments to
        *score_fn*.  The first element of each tuple is used as the
        document identifier.

        Returns a list of ``(doc_id, score)`` pairs.
        """
        if not candidates:
            return []

        batches = chunk_list(candidates, self.num_workers)

        all_results: List[Tuple[str, float]] = []
        with ProcessPoolExecutor(max_workers=self.num_workers) as pool:
            futures = [
                pool.submit(_score_batch, batch, score_fn)
                for batch in batches
            ]
            for future in futures:
                all_results.extend(future.result())

        return all_results


# ---------------------------------------------------------------------------
# ParallelSearcher
# ---------------------------------------------------------------------------

class ParallelSearcher:
    """Search multiple index shards in parallel and merge results.

    Uses a thread pool (suitable when each shard search is I/O-bound or
    releases the GIL internally) to query shards concurrently, then
    merges the per-shard sorted results via ``heapq.merge``.
    """

    def __init__(self, num_shards: int = 4) -> None:
        self.num_shards = num_shards

    def search(
        self,
        query_fn: Callable[..., list],
        shards: List[Any],
    ) -> list:
        """Run *query_fn* on each shard in parallel and return merged results.

        *query_fn* receives a single shard object and must return a list
        of ``(doc_id, score)`` tuples sorted descending by score.
        """
        if not shards:
            return []

        with ThreadPoolExecutor(max_workers=len(shards)) as pool:
            futures = [
                pool.submit(_search_shard, query_fn, shard)
                for shard in shards
            ]
            shard_results = [f.result() for f in futures]

        return self.merge_results(shard_results)

    @staticmethod
    def merge_results(
        shard_results: List[List[Tuple[str, float]]],
        top_k: int = 10,
    ) -> List[Tuple[str, float]]:
        """Merge pre-sorted shard results into a single top-K ranking.

        Each shard result must be sorted descending by score.  Uses
        ``heapq.merge`` for efficient O(N log K) merging.
        """
        if not shard_results:
            return []

        # heapq.merge works on ascending iterables.  We negate scores so
        # that the default min-merge produces descending-score order.
        def _negated(results: List[Tuple[str, float]]):
            for doc_id, score in results:
                yield (-score, doc_id)

        merged = heapq.merge(*[_negated(sr) for sr in shard_results])
        top: List[Tuple[str, float]] = []
        for neg_score, doc_id in merged:
            top.append((doc_id, -neg_score))
            if len(top) >= top_k:
                break
        return top
