proc-rota/djangorota/rota/models.py

from django.db import models
from django.utils import timezone
import uuid


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 importing the project's rota generator implementation. Some setups
		# provide the implementation as `rota_generator.shifts` (script/package),
		# while others provide it as `rota.shifts` (library folder). Try both
		# so the Django app works in either layout.
		RotaBuilder = None
		SingleShift = None
		import importlib
		for mod_name in ("rota_generator.shifts", "rota.shifts"):
			try:
				mod = importlib.import_module(mod_name)
				RotaBuilder = getattr(mod, "RotaBuilder")
				SingleShift = getattr(mod, "SingleShift")
				break
			except Exception:
				continue
		if RotaBuilder is None or SingleShift is None:
			raise RuntimeError(
				"Unable to import rota generator (tried 'rota_generator.shifts' and 'rota.shifts'); "
				"ensure the rota package is importable and contains RotaBuilder/SingleShift"
			)

		# compute weeks_to_rota (integer number of weeks)
		delta = self.end_date - self.start_date
		weeks_to_rota = max(1, int(delta.days // 7))

		# Build kwargs for RotaBuilder from any explicit builder args stored in
		# `self.options` (kept for backward compatibility), and treat the
		# remaining keys as constraint options which should map to
		# `RotaConstraintOptions`.
		opts = dict(self.options or {})
		builder_keys = [
			"balance_offset_modifier",
			"ltft_balance_offset",
			"use_previous_shifts",
			"use_shift_balance_extra",
			"use_bank_holiday_extra",
			"allow_force_assignment_with_leave_conflict",
		]
		rb_kwargs = {k: opts.pop(k) for k in builder_keys if k in opts}
		rb = RotaBuilder(start_date=self.start_date, weeks_to_rota=weeks_to_rota, name=self.name, **rb_kwargs)

		# If the rota package exposes a typed RotaConstraintOptions model, use it
		# to validate/coerce the remaining `opts` and attach it to the builder.
		try:
			RotaConstraintOptions = getattr(mod, "RotaConstraintOptions")
			if opts:
				# validate raw JSON into the typed model
				opts_model = RotaConstraintOptions.model_validate(opts)
			else:
				opts_model = RotaConstraintOptions()
			# attach the typed model to the builder (overrides defaults)
			setattr(rb, "constraint_options_model", opts_model)
		except Exception:
			# If the model isn't available or validation fails, keep builder
			# defaults and proceed (backwards compatible behaviour).
			pass

		# 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)

	# Store additional worker settings / metadata that map to the
	# pydantic `Worker` model used by the rota generator. This allows the
	# web UI to expose richer worker options without changing the core
	# Django schema for frequently used fields.
	options = models.JSONField(default=dict, blank=True)

	rotas = models.ManyToManyField(RotaSchedule, through="Assignment", related_name="workers")

	# Token used to allow workers (without Django accounts) to access
	# self-service links (request leave, view their calendar, etc.). This
	# is a UUID4 and should be treated as a secret. Regenerate to revoke.
	# During migration we add this field as nullable and non-unique first,
	# then populate unique values in a data migration and finally make it
	# non-null and unique. See migrations/ for the migration sequence.
	request_token = models.UUIDField(default=uuid.uuid4, editable=False, null=True, unique=False)

	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 `fte` as a float (fractional e.g. 1.0 == 100%),
		# while the pydantic Worker expects an integer percentage (100).
		# Convert conservatively:
		# - if stored value looks like a fraction (<= 10) treat it as fraction
		#   and multiply by 100 (e.g. 1.0 -> 100).
		# - otherwise treat it as already a percentage and coerce to int.
		raw_fte = getattr(self, "fte", None)
		try:
			if raw_fte is None:
				fte_pct = 100
			else:
				# numeric values
				f = float(raw_fte)
				if f <= 10:
					fte_pct = int(round(f * 100))
				else:
					fte_pct = int(round(f))
		except Exception:
			fte_pct = 100

		pyd_kwargs = {
			"name": self.name,
			"site": self.site or "",
			"grade": self.grade or 1,
			"id": self.pk,
			"fte": fte_pct,
		}

		# Merge any additional options stored on the Django model into the
		# kwargs passed to the pydantic Worker model. The pydantic model will
		# honour its own defaults and types; storing these as a dict keeps the
		# Django DB schema simple while exposing full worker settings in the UI.
		try:
			opts = getattr(self, "options", None) or {}
			if isinstance(opts, dict):
				# shallow merge; keys in `options` override defaults above when present
				pyd_kwargs.update(opts)

				# Defensive normalization: the DB may contain `null`/`[]` for fields
				# that the pydantic model expects to be ints/dicts/lists. Coerce a
				# few common keys to safe defaults so Worker(**kwargs) doesn't raise
				# on older or malformed records.
				int_keys = [
					"locum_max_shifts",
					"locum_max_shifts_per_week",
					"bank_holiday_extra",
					"prefer_multi_shift_together",
					"weekend_shift_target_number",
				]
				dict_keys = [
					"previous_shifts",
					"shift_balance_extra",
					"force_assign_with",
					"max_shifts_per_week_by_shift_name",
					"assign_as_block_preferences",
					"shift_fte_overrides",
				]
				list_keys = [
					"oop",
					"not_available_to_work",
					"pref_not_to_work",
					"work_requests",
					"locum_availability",
					"allowed_multi_shift_sets",
					"hard_day_dependencies",
					"hard_day_exclusions",
					"forced_assignments",
					"forced_assignments_by_date",
				]

				for k in int_keys:
					v = pyd_kwargs.get(k, None)
					try:
						if v is None:
							pyd_kwargs[k] = 0
						else:
							pyd_kwargs[k] = int(v)
					except Exception:
						pyd_kwargs[k] = 0

				for k in dict_keys:
					v = pyd_kwargs.get(k, None)
					if not isinstance(v, dict):
						# If a list was stored accidentally (common migration/serialization
						# mistakes) coerce to an empty dict rather than letting pydantic
						# fail.
						pyd_kwargs[k] = {}

				for k in list_keys:
					v = pyd_kwargs.get(k, None)
					if v is None:
						pyd_kwargs[k] = []
					elif isinstance(v, list):
						# leave as-is
						pass
					elif isinstance(v, (set, tuple)):
						pyd_kwargs[k] = list(v)
					else:
						pyd_kwargs[k] = []

				# allowed_multi_shift_sets is expected to be a list of sets; if the
				# DB contains lists-of-lists convert inner lists to sets.
				ams = pyd_kwargs.get("allowed_multi_shift_sets")
				if isinstance(ams, list):
					new = []
					for item in ams:
						if isinstance(item, list):
							new.append(set(item))
						elif isinstance(item, (set, tuple)):
							new.append(set(item))
						else:
							# leave unknown types to pydantic to either accept or coerce
							new.append(item)
					pyd_kwargs["allowed_multi_shift_sets"] = new
		except Exception:
			# ignore any issues reading options and proceed with core fields
			pass

		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)

	# Optional HTML export produced by a run (stored inline for convenience).
	# May be large; keep nullable to avoid breaking existing runs.
	export_html = models.TextField(null=True, 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()})"