from django.db import models from django.utils import timezone class RotaSchedule(models.Model): """Represents a named rota (collection of shifts over a date range). This is the top-level object that groups workers and runs. """ name = models.CharField(max_length=200, unique=True) start_date = models.DateField() end_date = models.DateField() description = models.TextField(blank=True) created_at = models.DateTimeField(auto_now_add=True) # Store the SingleShift definitions as a list of dicts (pydantic -> dict) shifts = models.JSONField(default=list, blank=True) # Additional builder options / constraints (serialisable) options = models.JSONField(default=dict, blank=True) def __str__(self): return self.name def to_rota_builder(self): """Construct a `rota.shifts.RotaBuilder` from this schedule and its shifts/workers. This will: - instantiate a RotaBuilder with `start_date` and computed `weeks_to_rota`; - convert stored `shifts` (list of dicts) to `SingleShift` pydantic models and add them; - convert assigned Django `Worker` objects to pydantic `Worker` and add them. Returns the RotaBuilder instance. """ try: from rota_generator.shifts import RotaBuilder, SingleShift except Exception: raise RuntimeError("Unable to import rota_generator.shifts; ensure the rota package is importable") # compute weeks_to_rota (integer number of weeks) delta = self.end_date - self.start_date weeks_to_rota = max(1, int(delta.days // 7)) # Allow passing options directly if they match RotaBuilder kwargs rb_kwargs = dict(self.options or {}) rb = RotaBuilder(start_date=self.start_date, weeks_to_rota=weeks_to_rota, name=self.name, **rb_kwargs) # Add shifts (expecting list of dicts compatible with SingleShift) shifts_objs = [] for s in self.shifts or []: try: obj = SingleShift(**s) except Exception as exc: raise ValueError(f"Invalid shift data for rota {self.name}: {exc}") shifts_objs.append(obj) if shifts_objs: rb.add_shifts(*shifts_objs) # Add workers assigned to this rota for w in self.workers.all(): try: pyd_worker = w.to_pydantic() except Exception: # skip workers that cannot be converted continue rb.add_worker(pyd_worker) return rb class Worker(models.Model): """Worker / staff member that can be assigned to one or more rotas.""" name = models.CharField(max_length=200) email = models.EmailField(blank=True) site = models.CharField(max_length=100, blank=True) grade = models.IntegerField(null=True, blank=True) fte = models.FloatField(default=1.0) active = models.BooleanField(default=True) rotas = models.ManyToManyField(RotaSchedule, through="Assignment", related_name="workers") def __str__(self): return self.name def to_pydantic(self): """Return a `rota.workers.Worker` pydantic model populated from this Django model. Only maps a minimal set of fields required by the scheduler; additional pydantic defaults will fill the rest. """ try: from rota_generator.workers import Worker as PydWorker except Exception: raise RuntimeError("Unable to import rota.workers.Worker; ensure project package is importable") # Map fte: Django stores percentage (100) while pydantic Worker expects int pyd_kwargs = { "name": self.name, "site": self.site or "", "grade": self.grade or 1, "id": self.pk, "fte": int(self.fte), } return PydWorker(**pyd_kwargs) class Assignment(models.Model): """Join model assigning a worker to a rota (can hold role / metadata).""" worker = models.ForeignKey(Worker, on_delete=models.CASCADE) rota = models.ForeignKey(RotaSchedule, on_delete=models.CASCADE) role = models.CharField(max_length=100, blank=True) created_at = models.DateTimeField(auto_now_add=True) class Meta: unique_together = ("worker", "rota") def __str__(self): return f"{self.worker} -> {self.rota}" class Leave(models.Model): """Leave record for a worker. Stored as inclusive start/end dates.""" worker = models.ForeignKey(Worker, on_delete=models.CASCADE, related_name="leaves") start_date = models.DateField() end_date = models.DateField() reason = models.CharField(max_length=255, blank=True) created_at = models.DateTimeField(auto_now_add=True) class Meta: ordering = ("-start_date",) def __str__(self): return f"{self.worker}: {self.start_date} → {self.end_date}" class RotaRun(models.Model): """Represents an execution of the scheduler for a RotaSchedule. The run stores status and (optionally) JSON results. A background runner or management command can update this record as the job progresses. """ STATUS_PENDING = "pending" STATUS_RUNNING = "running" STATUS_SUCCESS = "success" STATUS_FAILED = "failed" STATUS_CHOICES = [ (STATUS_PENDING, "Pending"), (STATUS_RUNNING, "Running"), (STATUS_SUCCESS, "Success"), (STATUS_FAILED, "Failed"), ] rota = models.ForeignKey(RotaSchedule, on_delete=models.CASCADE, related_name="runs") created_at = models.DateTimeField(auto_now_add=True) started_at = models.DateTimeField(null=True, blank=True) finished_at = models.DateTimeField(null=True, blank=True) status = models.CharField(max_length=20, choices=STATUS_CHOICES, default=STATUS_PENDING) result = models.JSONField(null=True, blank=True) log = models.TextField(blank=True) def mark_running(self): self.status = self.STATUS_RUNNING self.started_at = timezone.now() self.save(update_fields=["status", "started_at"]) def mark_finished(self, result=None): self.status = self.STATUS_SUCCESS self.finished_at = timezone.now() if result is not None: self.result = result self.save(update_fields=["status", "finished_at", "result"]) def mark_failed(self, message=""): self.status = self.STATUS_FAILED self.finished_at = timezone.now() self.log = (self.log or "") + "\n" + message self.save(update_fields=["status", "finished_at", "log"]) def __str__(self): return f"RotaRun({self.rota.name} - {self.status} @ {self.created_at.isoformat()})"