"""Background image prefetcher: overlap docker pulls with evaluation.

While the runner executes sample N (its image already local), worker
threads pull the images of upcoming samples N+1.. — so per-sample runs
never wait on a cold multi-GB pull unless the queue drains.

    from evalharness.sandbox.prefetch import BackgroundPrefetcher
    images = images_for_dataset(ds)          # ordered like the dataset
    with BackgroundPrefetcher(images, workers=4, lookahead=8) as bp:
        for sample in ds:
            bp.ensure(sample.sandbox.image)  # blocks only if still pulling
            ... run in sandbox ...
"""

import threading
import time
from typing import Dict, List, Optional

from .prefetch import _pull_one, local_images


class BackgroundPrefetcher:
    """Pull upcoming images on worker threads; evaluation thread consumes."""

    def __init__(self, images: List[str], workers: int = 4, lookahead: int = 8):
        self.images = list(images)
        self.workers = max(1, workers)
        self.lookahead = max(1, lookahead)
        self._cursor = 0
        self._lock = threading.Lock()
        self._ready: Dict[str, bool] = {}
        self._failed: Dict[str, str] = {}
        self._stop = threading.Event()
        self._threads: List[threading.Thread] = []
        for img in self.images:
            self._ready[img] = True if img in local_images() else False

    def __enter__(self) -> 'BackgroundPrefetcher':
        for i in range(self.workers):
            t = threading.Thread(target=self._work, name=f'eh-prefetch-{i}', daemon=True)
            t.start()
            self._threads.append(t)
        return self

    def __exit__(self, *exc) -> None:
        self._stop.set()
        for t in self._threads:
            t.join(timeout=5)

    def _work(self) -> None:
        while not self._stop.is_set():
            img = self._next_pending()
            if img is None:
                time.sleep(0.5)
                continue
            try:
                _pull_one(img)
                with self._lock:
                    self._ready[img] = True
            except Exception as e:
                with self._lock:
                    self._failed[img] = str(e)[:200]

    def _next_pending(self) -> Optional[str]:
        """Claim the next not-ready image within the lookahead window."""
        with self._lock:
            hi = min(self._cursor + self.lookahead, len(self.images))
            for i in range(self._cursor, hi):
                img = self.images[i]
                if not self._ready.get(img) and img not in self._failed:
                    return img  # claimed (pull is idempotent; duplicates are cheap)
            return None

    def ensure(self, image: Optional[str]) -> bool:
        """Advance the cursor to `image`; wait (bounded) until pulled."""
        if not image:
            return True
        with self._lock:  # allow random access ordering too
            if image in self.images and self.images.index(image) >= self._cursor:
                self._cursor = self.images.index(image)
        deadline = time.time() + 3600
        while time.time() < deadline and not self._stop.is_set():
            with self._lock:
                if self._ready.get(image) or image in self._failed:
                    return self._ready.get(image, False)
            time.sleep(1.0)
        return self._ready.get(image, False)

    def stats(self) -> Dict[str, int]:
        with self._lock:
            ready = sum(1 for v in self._ready.values() if v)
            return {'total': len(self.images), 'ready': ready, 'failed': len(self._failed)}