from __future__ import annotations


class EditSchedulerMixin:
    def apply_edit_writes_map(self, req, sample_output) -> None:
        req_id_str = str(req.req_id)
        edit_writes_map = sample_output.edit_writes_map[req_id_str]
        for block_id, block_writes in edit_writes_map.items():
            if not block_writes:
                continue

            dllm_block = req.dllm_blocks[int(block_id)]
            for rel_idx, token in block_writes.items():
                rel_idx = int(rel_idx)
                token = int(token)
                prev_token = dllm_block.token_ids[rel_idx]
                dllm_block.write_token(token, rel_idx)
                if prev_token == dllm_block.mask_token_id and token != dllm_block.mask_token_id:
                    req.new_tokens += 1