import collections

class Request:
    def __init__(self, request_id: str, context_length: int = 0, text: str = "", answer: str = ""):
        self.id = request_id
        self.text = text
        self.context_length = context_length
        self.answer = answer
    
    def __repr__(self):
        return f"Request(id='{self.id}', len={self.context_length})"

class GroupAwareRRScheduler:
    """
    Implements a hybrid scheduling strategy for pre-grouped requests.
    It prioritizes filling a batch from an active cohort, then greedily
    fills remaining capacity with lead queries from new groups.
    """
    def __init__(self, pre_grouped_requests: list[list[Request]], max_prefill_tokens: int):
        if not pre_grouped_requests or not any(pre_grouped_requests):
            raise ValueError("Input request list cannot be empty.")
        
        self.max_prefill_tokens = max_prefill_tokens
        self.group_pool = [collections.deque(group) for group in pre_grouped_requests if group]
        
        self.active_groups = collections.deque()
        self.next_group_idx_for_cohort = 0
        self.total_requests = sum(len(group) for group in self.group_pool)
        self.processed_count = 0
        print(f"Scheduler initialized with {len(self.group_pool)} pre-defined groups.")

    def get_next_batch(self) -> list[Request] | None:
        if self.processed_count >= self.total_requests:
            return None

        if not self.active_groups:
            if self.next_group_idx_for_cohort < len(self.group_pool):
                print("--- Starting a new active set ---")
                new_group = self.group_pool[self.next_group_idx_for_cohort]
                self.active_groups.append(new_group)
                self.next_group_idx_for_cohort += 1
            else:
                return None

        batch = []
        current_tokens = 0
        
        active_groups_count = len(self.active_groups)
        for _ in range(active_groups_count):
            group = self.active_groups.popleft()
            if group:
                candidate = group[0]
                # FIX 2: Allow at least one item in a batch, even if it's oversized.
                if not batch or (current_tokens + candidate.context_length) <= self.max_prefill_tokens:
                    req = group.popleft()
                    batch.append(req)
                    current_tokens += req.context_length
            if group:
                self.active_groups.append(group)

        while (self.next_group_idx_for_cohort < len(self.group_pool)):
            next_group_candidate = self.group_pool[self.next_group_idx_for_cohort]
            if not next_group_candidate:
                self.next_group_idx_for_cohort += 1
                continue
                
            lead_query = next_group_candidate[0]
            # FIX 2: Apply the same logic here for adding new groups.
            if not batch or (current_tokens + lead_query.context_length) <= self.max_prefill_tokens:
                print(f"--- Greedily adding lead query '{lead_query.id}' to batch ---")
                req = next_group_candidate.popleft()
                batch.append(req)
                current_tokens += req.context_length
                if next_group_candidate:
                    self.active_groups.append(next_group_candidate)
                self.next_group_idx_for_cohort += 1
            else:
                break

        if batch:
            self.processed_count += len(batch)

        return batch if batch else None