"""U10-7: a pool that breaks mid-batch must never score the tail INCORRECT."""
from concurrent.futures.process import BrokenProcessPool

import pytest

from align.eval.verifier import Outcome, OutcomeVerifier, VerifierConfig


def _stub_scorer(gold, preds, timeout):
    return gold == preds[0]


class _Fut:
    def __init__(self, val):
        self.val = val

    def result(self, timeout=None):
        if isinstance(self.val, BaseException):
            raise self.val
        return self.val


class _DiesAfter:
    """Worker 'OOM-killed' after `ok_first` submissions: the rest break."""

    _processes = {}

    def __init__(self, ok_first):
        self.n = 0
        self.ok_first = ok_first

    def submit(self, fn, gold, preds, timeout):
        self.n += 1
        if self.n <= self.ok_first:
            return _Fut(fn(gold, preds, timeout))
        return _Fut(BrokenProcessPool("worker died"))

    def shutdown(self, **kw):
        pass


def _verifier(factory):
    return OutcomeVerifier(
        VerifierConfig(workers=2), _score_fn=_stub_scorer, _executor_factory=factory
    )


def test_persistently_broken_tail_raises_not_incorrect():
    v = _verifier(lambda n: _DiesAfter(ok_first=3))
    golds = [str(i) for i in range(10)]
    with pytest.raises(RuntimeError, match="BrokenProcessPool"):
        v.verify_batch(golds, golds)


def test_broken_tail_is_retried_on_a_fresh_pool():
    pools = []

    def factory(n):
        pools.append(_DiesAfter(ok_first=3 if not pools else 99))
        return pools[-1]

    v = _verifier(factory)
    golds = [str(i) for i in range(10)]
    out = v.verify_batch(golds, golds)
    assert out == [Outcome.CORRECT] * 10 and len(pools) == 2