///|
/// A person who can be assigned to shifts.
pub(all) struct Worker {
  id : Int
  name : String
  skills : Array[String]
  available_slots : Array[Int]
} derive(Eq, Debug, ToJson, FromJson)

///|
/// A unit of work that needs exactly one eligible worker.
pub(all) struct Shift {
  id : Int
  name : String
  slot : Int
  required_skill : String
} derive(Eq, Debug, ToJson, FromJson)

///|
/// One resolved worker-to-shift assignment.
pub(all) struct RosterEntry {
  shift : Shift
  worker : Worker
} derive(Eq, Debug, ToJson)

///|
/// A feasible roster together with search evidence.
pub(all) struct Roster {
  entries : Array[RosterEntry]
  stats : SolveStats
} derive(Eq, Debug, ToJson)

///|
/// A workload balance score. Lower spread and lower maximum load are better.
pub(all) struct BalanceScore {
  minimum_load : Int
  maximum_load : Int
  spread : Int
} derive(Eq, Debug, ToJson)

///|
/// The best roster found within a bounded candidate search.
pub(all) struct OptimizedRoster {
  roster : Roster
  balance : BalanceScore
  candidates_evaluated : Int
} derive(Eq, Debug)

///|
/// Hard rules applied while building a roster.
///
/// Slot values are treated as ordered, equally sized time units. A minimum
/// gap of two therefore prevents one worker from taking shifts in adjacent
/// slots while allowing shifts two or more units apart.
pub(all) struct RosterPolicy {
  max_shifts_per_worker : Int?
  minimum_slot_gap : Int
} derive(Eq, Debug, ToJson, FromJson)

///|
/// An inclusive workload range for one worker.
pub(all) struct WorkerQuota {
  worker_id : Int
  minimum_shifts : Int
  maximum_shifts : Int
} derive(Eq, Debug, ToJson, FromJson)

///|
/// A policy that adds no restrictions beyond skills, availability, and
/// same-slot exclusivity.
pub fn RosterPolicy::default() -> RosterPolicy {
  { max_shifts_per_worker: None, minimum_slot_gap: 0, }
}

///|
/// Domain-specific scheduling failures.
pub(all) enum RosterError {
  DuplicateWorkerId(Int)
  DuplicateShiftId(Int)
  DuplicateCoverageRequirementId(Int)
  InvalidCoverageWorkerCount(Int, Int)
  NoEligibleWorker(Int, String)
  InsufficientSlotCapacity(Int, Int, Int)
  InvalidCandidateLimit(Int)
  InvalidWorkloadLimit(Int)
  InvalidRestGap(Int)
  UnknownPenaltyWorker(Int)
  UnknownPenaltyShift(Int)
  DuplicateAssignmentPenalty(Int, Int)
  InvalidAssignmentPenalty(Int, Int, Int)
  InvalidConsecutiveSlotPenalty(Int)
  InvalidConsecutiveWorkLimit(Int)
  UnknownQuotaWorker(Int)
  DuplicateWorkerQuota(Int)
  InvalidWorkerQuota(Int, Int, Int)
  MinimumWorkerQuotaExceedsShiftCount(Int, Int, Int)
  InvalidModel(ModelError)
  Unsatisfiable(SolveStats)
} derive(Eq, Debug)

///|
/// Explain a scheduling failure without exposing solver internals.
pub fn RosterError::message(self : RosterError) -> String {
  match self {
    DuplicateWorkerId(id) => "duplicate worker id: \{id}"
    DuplicateShiftId(id) => "duplicate shift id: \{id}"
    DuplicateCoverageRequirementId(id) =>
      "duplicate coverage requirement id: \{id}"
    InvalidCoverageWorkerCount(id, count) =>
      "coverage requirement \{id} must need at least one worker, got \{count}"
    NoEligibleWorker(id, name) =>
      "shift \{id} (\{name}) has no worker with the required skill and availability"
    InsufficientSlotCapacity(slot, required, assignable) =>
      "slot \{slot} requires \{required} workers, but at most \{assignable} shifts can be covered"
    InvalidCandidateLimit(limit) =>
      "candidate limit must be positive, got \{limit}"
    InvalidWorkloadLimit(limit) =>
      "workload limit must not be negative, got \{limit}"
    InvalidRestGap(gap) => "minimum slot gap must not be negative, got \{gap}"
    UnknownPenaltyWorker(id) =>
      "assignment penalty refers to unknown worker id: \{id}"
    UnknownPenaltyShift(id) =>
      "assignment penalty refers to unknown shift id: \{id}"
    DuplicateAssignmentPenalty(worker_id, shift_id) =>
      "duplicate assignment penalty for worker \{worker_id} and shift \{shift_id}"
    InvalidAssignmentPenalty(worker_id, shift_id, penalty) =>
      "assignment penalty for worker \{worker_id} and shift \{shift_id} must not be negative, got \{penalty}"
    InvalidConsecutiveSlotPenalty(penalty) =>
      "consecutive slot penalty must not be negative, got \{penalty}"
    InvalidConsecutiveWorkLimit(limit) =>
      "maximum consecutive work slots must be positive, got \{limit}"
    UnknownQuotaWorker(id) => "worker quota refers to unknown worker id: \{id}"
    DuplicateWorkerQuota(id) => "duplicate worker quota for worker id: \{id}"
    InvalidWorkerQuota(id, minimum, maximum) =>
      "worker \{id} quota must satisfy 0 <= minimum <= maximum, got \{minimum}..\{maximum}"
    MinimumWorkerQuotaExceedsShiftCount(id, minimum, shift_count) =>
      "worker \{id} minimum workload \{minimum} exceeds the \{shift_count} available shifts"
    InvalidModel(error) => "invalid constraint model: \{error.message()}"
    Unsatisfiable(stats) =>
      "no feasible roster after \{stats.nodes} assignments and \{stats.backtracks} backtracks"
  }
}

///|
priv struct RosterModel {
  problem : Problem
  variables : Array[Var]
}

///|
fn add_workload_constraints(
  model : RosterModel,
  workers : Array[Worker],
  max_shifts_per_worker : Int,
) -> Result[Unit, RosterError] {
  if model.variables.length() == 0 {
    return Ok(())
  }
  for worker_index = 0
      worker_index < workers.length()
      worker_index = worker_index + 1 {
    match
      model.problem.add_named_constraint(
        "worker-\{workers[worker_index].id}:workload-cap",
        CountAtMost(model.variables.copy(), worker_index, max_shifts_per_worker),
      ) {
      Ok(_) => ()
      Err(error) => return Err(InvalidModel(error))
    }
  }
  Ok(())
}

///|
fn worker_index_for_id(workers : Array[Worker], worker_id : Int) -> Int? {
  for index = 0; index < workers.length(); index = index + 1 {
    if workers[index].id == worker_id {
      return Some(index)
    }
  }
  None
}

///|
fn add_worker_quotas(
  model : RosterModel,
  workers : Array[Worker],
  shifts : Array[Shift],
  quotas : Array[WorkerQuota],
) -> Result[Unit, RosterError] {
  for index = 0; index < quotas.length(); index = index + 1 {
    let quota = quotas[index]
    if quota.minimum_shifts < 0 || quota.maximum_shifts < quota.minimum_shifts {
      return Err(
        InvalidWorkerQuota(
          quota.worker_id,
          quota.minimum_shifts,
          quota.maximum_shifts,
        ),
      )
    }
    if quota.minimum_shifts > shifts.length() {
      return Err(
        MinimumWorkerQuotaExceedsShiftCount(
          quota.worker_id,
          quota.minimum_shifts,
          shifts.length(),
        ),
      )
    }
    for previous = 0; previous < index; previous = previous + 1 {
      if quotas[previous].worker_id == quota.worker_id {
        return Err(DuplicateWorkerQuota(quota.worker_id))
      }
    }
    let worker_index = match worker_index_for_id(workers, quota.worker_id) {
      None => return Err(UnknownQuotaWorker(quota.worker_id))
      Some(value) => value
    }
    if model.variables.length() > 0 {
      match
        model.problem.add_named_constraint(
          "worker-\{quota.worker_id}:workload-quota",
          CountBetween(
            model.variables.copy(),
            worker_index,
            quota.minimum_shifts,
            quota.maximum_shifts,
          ),
        ) {
        Ok(_) => ()
        Err(error) => return Err(InvalidModel(error))
      }
    }
  }
  Ok(())
}

///|
fn slot_distance(left : Int, right : Int) -> Int {
  if left >= right {
    left - right
  } else {
    right - left
  }
}

///|
fn add_rest_constraints(
  model : RosterModel,
  shifts : Array[Shift],
  minimum_slot_gap : Int,
) -> Result[Unit, RosterError] {
  for left = 0; left < shifts.length(); left = left + 1 {
    for right = left + 1; right < shifts.length(); right = right + 1 {
      let distance = slot_distance(shifts[left].slot, shifts[right].slot)
      if distance > 0 && distance < minimum_slot_gap {
        match
          model.problem.add_named_constraint(
            "rest-gap-\{minimum_slot_gap}:shift-\{shifts[left].id}-vs-\{shifts[right].id}",
            Different(model.variables[left], model.variables[right]),
          ) {
          Ok(_) => ()
          Err(error) => return Err(InvalidModel(error))
        }
      }
    }
  }
  Ok(())
}

///|
fn apply_roster_policy(
  model : RosterModel,
  workers : Array[Worker],
  shifts : Array[Shift],
  policy : RosterPolicy,
) -> Result[Unit, RosterError] {
  match policy.max_shifts_per_worker {
    Some(limit) =>
      match add_workload_constraints(model, workers, limit) {
        Err(error) => return Err(error)
        Ok(_) => ()
      }
    None => ()
  }
  add_rest_constraints(model, shifts, policy.minimum_slot_gap)
}

///|
fn validate_roster_policy(policy : RosterPolicy) -> Result[Unit, RosterError] {
  if policy.minimum_slot_gap < 0 {
    return Err(InvalidRestGap(policy.minimum_slot_gap))
  }
  match policy.max_shifts_per_worker {
    Some(limit) if limit < 0 => Err(InvalidWorkloadLimit(limit))
    _ => Ok(())
  }
}

///|
fn string_member(values : Array[String], expected : String) -> Bool {
  for value in values {
    if value == expected {
      return true
    }
  }
  false
}

///|
fn int_member(values : Array[Int], expected : Int) -> Bool {
  for value in values {
    if value == expected {
      return true
    }
  }
  false
}

///|
fn worker_is_eligible(worker : Worker, shift : Shift) -> Bool {
  int_member(worker.available_slots, shift.slot) &&
  (
    shift.required_skill.length() == 0 ||
    string_member(worker.skills, shift.required_skill)
  )
}

///|
fn duplicate_worker_id(workers : Array[Worker]) -> Int? {
  for i = 0; i < workers.length(); i = i + 1 {
    for j = 0; j < i; j = j + 1 {
      if workers[i].id == workers[j].id {
        return Some(workers[i].id)
      }
    }
  }
  None
}

///|
fn duplicate_shift_id(shifts : Array[Shift]) -> Int? {
  for i = 0; i < shifts.length(); i = i + 1 {
    for j = 0; j < i; j = j + 1 {
      if shifts[i].id == shifts[j].id {
        return Some(shifts[i].id)
      }
    }
  }
  None
}

///|
fn build_roster_model(
  workers : Array[Worker],
  shifts : Array[Shift],
) -> Result[RosterModel, RosterError] {
  let analysis = analyze_roster(workers, shifts)
  if analysis.issues.length() > 0 {
    match analysis.issues[0] {
      DuplicateWorkerIdFound(id) => return Err(DuplicateWorkerId(id))
      DuplicateShiftIdFound(id) => return Err(DuplicateShiftId(id))
      NoEligibleWorkerFound(id, name) => return Err(NoEligibleWorker(id, name))
      SlotCoverageShortfall(slot, required, assignable) =>
        return Err(InsufficientSlotCapacity(slot, required, assignable))
    }
  }
  let problem = Problem::new()
  let variables : Array[Var] = []
  for shift in shifts {
    let eligible : Array[Int] = []
    for worker_index = 0
        worker_index < workers.length()
        worker_index = worker_index + 1 {
      if worker_is_eligible(workers[worker_index], shift) {
        eligible.push(worker_index)
      }
    }
    if eligible.length() == 0 {
      return Err(NoEligibleWorker(shift.id, shift.name))
    }
    match problem.add_variable("shift-\{shift.id}", eligible) {
      Ok(variable) => variables.push(variable)
      Err(error) => return Err(InvalidModel(error))
    }
  }
  for left = 0; left < shifts.length(); left = left + 1 {
    for right = left + 1; right < shifts.length(); right = right + 1 {
      if shifts[left].slot == shifts[right].slot {
        match
          problem.add_named_constraint(
            "slot-\{shifts[left].slot}:shift-\{shifts[left].id}-vs-\{shifts[right].id}",
            Different(variables[left], variables[right]),
          ) {
          Ok(_) => ()
          Err(error) => return Err(InvalidModel(error))
        }
      }
    }
  }
  Ok({ problem, variables, })
}

///|
fn roster_from_solution(
  workers : Array[Worker],
  shifts : Array[Shift],
  variables : Array[Var],
  solution : Solution,
  stats : SolveStats,
) -> Roster {
  let entries : Array[RosterEntry] = []
  for index = 0; index < shifts.length(); index = index + 1 {
    match solution.get(variables[index].name) {
      Some(worker_index) =>
        entries.push({ shift: shifts[index], worker: workers[worker_index], })
      None => ()
    }
  }
  { entries, stats, }
}

///|
/// Assign one eligible worker to every shift.
///
/// Workers may serve again in a different slot, but cannot cover two shifts in
/// the same slot. Skill and availability rules are compiled into variable
/// domains before search starts.
pub fn build_roster(
  workers : Array[Worker],
  shifts : Array[Shift],
) -> Result[Roster, RosterError] {
  match build_roster_model(workers, shifts) {
    Err(error) => Err(error)
    Ok(model) =>
      match model.problem.solve() {
        Unsatisfied(stats) => Err(Unsatisfiable(stats))
        Satisfied(solution, stats) =>
          Ok(
            roster_from_solution(
              workers,
              shifts,
              model.variables,
              solution,
              stats,
            ),
          )
      }
  }
}

///|
/// Assign every shift while limiting each worker's total workload.
///
/// A zero limit is valid for an empty schedule. If the limit makes a non-empty
/// schedule impossible, the result reports an unsatisfiable model with search
/// statistics.
pub fn build_capped_roster(
  workers : Array[Worker],
  shifts : Array[Shift],
  max_shifts_per_worker : Int,
) -> Result[Roster, RosterError] {
  build_roster_with_policy(workers, shifts, {
    max_shifts_per_worker: Some(max_shifts_per_worker),
    minimum_slot_gap: 0,
  })
}

///|
/// Build a roster under composable workload and rest rules.
pub fn build_roster_with_policy(
  workers : Array[Worker],
  shifts : Array[Shift],
  policy : RosterPolicy,
) -> Result[Roster, RosterError] {
  match validate_roster_policy(policy) {
    Err(error) => return Err(error)
    Ok(_) => ()
  }
  match build_roster_model(workers, shifts) {
    Err(error) => Err(error)
    Ok(model) => {
      match apply_roster_policy(model, workers, shifts, policy) {
        Err(error) => return Err(error)
        Ok(_) => ()
      }
      match model.problem.solve() {
        Unsatisfied(stats) => Err(Unsatisfiable(stats))
        Satisfied(solution, stats) =>
          Ok(
            roster_from_solution(
              workers,
              shifts,
              model.variables,
              solution,
              stats,
            ),
          )
      }
    }
  }
}

///|
/// Build a roster under global policy rules and per-worker workload quotas.
///
/// Quota bounds are inclusive. They compose with the uniform workload cap and
/// minimum rest gap already provided by `RosterPolicy`.
pub fn build_roster_with_quotas(
  workers : Array[Worker],
  shifts : Array[Shift],
  policy : RosterPolicy,
  quotas : Array[WorkerQuota],
) -> Result[Roster, RosterError] {
  match validate_roster_policy(policy) {
    Err(error) => return Err(error)
    Ok(_) => ()
  }
  match build_roster_model(workers, shifts) {
    Err(error) => Err(error)
    Ok(model) => {
      match apply_roster_policy(model, workers, shifts, policy) {
        Err(error) => return Err(error)
        Ok(_) => ()
      }
      match add_worker_quotas(model, workers, shifts, quotas) {
        Err(error) => return Err(error)
        Ok(_) => ()
      }
      match model.problem.solve() {
        Unsatisfied(stats) => Err(Unsatisfiable(stats))
        Satisfied(solution, stats) =>
          Ok(
            roster_from_solution(
              workers,
              shifts,
              model.variables,
              solution,
              stats,
            ),
          )
      }
    }
  }
}

///|
/// Find the worker assigned to a shift id.
pub fn Roster::worker_for_shift(self : Roster, shift_id : Int) -> Worker? {
  for entry in self.entries {
    if entry.shift.id == shift_id {
      return Some(entry.worker)
    }
  }
  None
}

///|
/// Count how many shifts were assigned to a worker id.
pub fn Roster::assignment_count(self : Roster, worker_id : Int) -> Int {
  let mut count = 0
  for entry in self.entries {
    if entry.worker.id == worker_id {
      count = count + 1
    }
  }
  count
}

///|
fn balance_score(workers : Array[Worker], roster : Roster) -> BalanceScore {
  if workers.length() == 0 {
    return { minimum_load: 0, maximum_load: 0, spread: 0, }
  }
  let mut minimum_load = roster.entries.length()
  let mut maximum_load = 0
  for worker in workers {
    let load = roster.assignment_count(worker.id)
    if load < minimum_load {
      minimum_load = load
    }
    if load > maximum_load {
      maximum_load = load
    }
  }
  { minimum_load, maximum_load, spread: maximum_load - minimum_load, }
}

///|
fn better_balance(candidate : BalanceScore, current : BalanceScore) -> Bool {
  candidate.spread < current.spread ||
  (
    candidate.spread == current.spread &&
    candidate.maximum_load < current.maximum_load
  )
}

///|
/// Find a feasible roster with the smallest workload spread among a bounded
/// number of deterministic candidates.
///
/// The limit makes optimization cost explicit. Increasing it may discover a
/// fairer roster, while the returned candidate count records the evidence used
/// for the decision.
pub fn build_balanced_roster(
  workers : Array[Worker],
  shifts : Array[Shift],
  candidate_limit : Int,
) -> Result[OptimizedRoster, RosterError] {
  build_balanced_roster_with_policy(
    workers,
    shifts,
    RosterPolicy::default(),
    candidate_limit,
  )
}

///|
/// Find the fairest roster among at most `candidate_limit` feasible rosters
/// while enforcing workload and rest policy as hard constraints.
pub fn build_balanced_roster_with_policy(
  workers : Array[Worker],
  shifts : Array[Shift],
  policy : RosterPolicy,
  candidate_limit : Int,
) -> Result[OptimizedRoster, RosterError] {
  if candidate_limit <= 0 {
    return Err(InvalidCandidateLimit(candidate_limit))
  }
  match validate_roster_policy(policy) {
    Err(error) => return Err(error)
    Ok(_) => ()
  }
  match build_roster_model(workers, shifts) {
    Err(error) => Err(error)
    Ok(model) => {
      match apply_roster_policy(model, workers, shifts, policy) {
        Err(error) => return Err(error)
        Ok(_) => ()
      }
      match model.problem.solve_all(candidate_limit) {
        Err(error) => Err(InvalidModel(error))
        Ok((solutions, stats)) => {
          if solutions.length() == 0 {
            return Err(Unsatisfiable(stats))
          }
          let mut best_roster = roster_from_solution(
            workers,
            shifts,
            model.variables,
            solutions[0],
            stats,
          )
          let mut best_balance = balance_score(workers, best_roster)
          for index = 1; index < solutions.length(); index = index + 1 {
            let candidate = roster_from_solution(
              workers,
              shifts,
              model.variables,
              solutions[index],
              stats,
            )
            let candidate_balance = balance_score(workers, candidate)
            if better_balance(candidate_balance, best_balance) {
              best_roster = candidate
              best_balance = candidate_balance
            }
          }
          Ok({
            roster: best_roster,
            balance: best_balance,
            candidates_evaluated: solutions.length(),
          })
        }
      }
    }
  }
}