from collections import deque
from collections.abc import Iterable
from typing import override

from procset import ProcSet

from pybatsim.batsim import (
    Batsim,
    EDCHelloEvent,
    ExecuteJobEvent,
    Job,
    JobCompletedEvent,
    JobSubmittedEvent,
    RejectJobEvent,
    RxEvent,
    Scheduler,
    SimulationBeginsEvent,
    SimulationEndsEvent,
    SimulationFeatures,
)


class FCFSScheduler(Scheduler):
    waiting_queue: deque[Job]
    idle_resources: ProcSet | None = None
    cluster_size: int | None = None

    @override
    def __init__(self, batsim: Batsim, options: str | None = None):
        super().__init__(batsim, options)

        self.batsim.simulation_metadata.requested_features |= (
            SimulationFeatures.FORWARD_PROFILES_ON_JOB_SUBMISSION
        )
        edc_hello_event = EDCHelloEvent(
            timestamp=self.batsim.time,
            simulation_metadata=self.batsim.simulation_metadata,
            edc_name=type(self).__name__,
            edc_version='0.0.0',
            edc_commit='',
        )
        self.batsim.append_event(edc_hello_event)
        self.waiting_queue = deque()

    @override
    def handle_msg(self, msg: Iterable[RxEvent]) -> None:
        super().handle_msg(msg)
        self.schedule_jobs()

    @override
    def handle_simulation_begins(self, event: SimulationBeginsEvent) -> None:
        self.cluster_size = event.computation_host_number
        self.idle_resources = ProcSet((0, self.cluster_size - 1))

    @override
    def handle_simulation_ends(self, event: SimulationEndsEvent) -> None:
        pass

    @override
    def handle_submitted_job(self, event: JobSubmittedEvent) -> None:
        assert self.cluster_size is not None, (
            'expected initialized cluster_size upon JobSubmittedEvent'
        )

        job = event.job

        if job.resource_request > self.cluster_size:
            # Job requests more resources than cluster_size: trivial reject.
            reject_event = RejectJobEvent(timestamp=self.batsim.time, job=job)
            self.batsim.append_event(reject_event)
        else:
            self.waiting_queue.append(job)

    @override
    def handle_completed_job(self, event: JobCompletedEvent) -> None:
        assert self.idle_resources is not None, (
            'expected initialized idle_resources upon JobCompletedEvent'
        )
        self.idle_resources |= event.job.allocation

    def schedule_jobs(self) -> None:
        scheduled_jobs: list[Job] = []

        while self.waiting_queue:
            # Stop scheduling jobs on the first job lacking resources.
            if self.waiting_queue[0].resource_request > len(self.idle_resources):
                break

            # Extract job from waiting queue.
            job = self.waiting_queue.popleft()

            # Allocate resources to job.
            resources = self.idle_resources[: job.resource_request]
            job.allocation = ProcSet(*resources)
            self.idle_resources -= job.allocation

            # Add job to the pending schedule batch.
            scheduled_jobs.append(job)

        # Send the list of scheduled jobs to Batsim.
        now = self.batsim.time
        for job in scheduled_jobs:
            execute_event = ExecuteJobEvent(timestamp=now, job=job)
            self.batsim.append_event(execute_event)
