"""InfraLens Python example: bounded multistage inference scheduler."""

from dataclasses import dataclass, field


@dataclass
class Stage:
    name: str
    capacity: int
    queue: list[str] = field(default_factory=list)
    completed: list[str] = field(default_factory=list)

    def admit(self, item: str) -> bool:
        if len(self.queue) >= self.capacity:
            return False
        self.queue.append(item)
        return True

    def run_one(self) -> str | None:
        if not self.queue:
            return None
        item = self.queue.pop(0)
        output = f"{item}->{self.name}"
        self.completed.append(output)
        return output


def run_pipeline(requests: list[str]) -> tuple[list[str], list[str]]:
    thinker = Stage("autoregressive", capacity=2)
    generator = Stage("diffusion", capacity=1)
    decoder = Stage("vae_decode", capacity=1)
    rejected: list[str] = []
    outputs: list[str] = []
    for request in requests:
        if not thinker.admit(request):
            rejected.append(request)
    while thinker.queue or generator.queue or decoder.queue:
        output = decoder.run_one()
        if output:
            outputs.append(output)
        output = generator.run_one()
        if output and not decoder.admit(output):
            rejected.append(output)
        output = thinker.run_one()
        if output and not generator.admit(output):
            rejected.append(output)
    return outputs, rejected


def smoke_test() -> None:
    outputs, rejected = run_pipeline(["r0", "r1", "r2"])
    assert len(outputs) == 2
    assert rejected == ["r2"]
    assert outputs[0].endswith("->vae_decode")
    print("multimodal_stage_scheduler_toy.py ok", {"outputs": outputs, "rejected": rejected})


if __name__ == "__main__":
    smoke_test()