import unittest
from types import SimpleNamespace

from prain import drafts
from tests.helpers import TmpCase  # noqa: F401  (every test module imports the helpers)


def new_request(req_id, xargs):
    return SimpleNamespace(req_id=req_id, sampling_params=SimpleNamespace(extra_args=xargs))


def step(new=(), finished=()):
    return SimpleNamespace(scheduled_new_reqs=list(new), finished_req_ids=set(finished))


class Drafts(unittest.TestCase):
    def setUp(self):
        drafts._DRAFTS.clear()
        drafts.clear()

    def test_draft_rows_follow_the_batch_layout(self):
        # a: 3 prompt tokens (full); b: 2 prompt tokens (draft)
        drafts.note_step(step([new_request("a", None), new_request("b", {"prain": "draft"})]),
                         ["a", "b"], [0, 3, 5])
        self.assertEqual(drafts.spans(), [(3, 5)])
        # next step: both decode one token each, in a new order
        drafts.note_step(step(), ["b", "a"], [0, 1, 2])
        self.assertEqual(drafts.spans(), [(0, 1)])
        drafts.clear()
        self.assertIsNone(drafts.spans())

    def test_finished_drafts_are_forgotten(self):
        drafts.note_step(step([new_request("b", {"prain": "draft"})]), ["b"], [0, 4])
        drafts.note_step(step(finished=["b"]), [], [0])
        self.assertIsNone(drafts.spans())
        self.assertEqual(drafts._DRAFTS, set())

    def test_requests_without_the_flag_are_full(self):
        drafts.note_step(step([new_request("a", {"other": 1}),
                               SimpleNamespace(req_id="c", sampling_params=None)]),
                         ["a", "c"], [0, 1, 2])
        self.assertIsNone(drafts.spans())