import threading
from tqdm import tqdm
from queue import Queue
from typing import Iterable, Optional
from src.logger import get_logger

logger = get_logger(__name__)


def batch_with_preload(
    items: list,
    batch_size: Optional[int] = None,
    workers: int = 1,
    process_func=lambda x: x,
    progress: bool = True,
) -> Iterable[list]:
    batch_size = batch_size or len(items)
    if not items:
        return

    queue = Queue(maxsize=workers)
    sentinel = object()

    def producer_task():
        try:
            for i in range(0, len(items), batch_size):
                batch = items[i : i + batch_size]
                processed_batch = process_func(batch)

                queue.put(processed_batch)
        except Exception as e:
            logger.error(f"Error in producer task: {e}")
        finally:
            queue.put(sentinel)

    producer_thread = threading.Thread(target=producer_task, daemon=True)
    producer_thread.start()

    total_batches = (len(items) + batch_size - 1) // batch_size
    pbar = tqdm(total=total_batches, desc="Processing batches", disable=not progress)

    try:
        while True:
            item = queue.get()
            if item is sentinel:
                break

            pbar.update(1)
            yield item
    finally:
        pbar.close()
        producer_thread.join()
