import datetime
import pytest
from pytest import approx
from rota_generator.shifts import PostShiftConstraint, RotaBuilder, SingleShift, days
from rota_generator.workers import Worker
from rota_generator.shifts import PreShiftConstraint

@pytest.fixture
def demo_rota_nights():
    weeks_to_rota = 12
    start_date = datetime.date(2022, 3, 14)

    Rota = RotaBuilder(
        start_date,
        weeks_to_rota=weeks_to_rota,
        balance_offset_modifier=1,
    )

    Rota.constraint_options_model.max_shifts_per_week = 5
    Rota.constraint_options_model.max_shifts_per_month = 31
    # legacy option `limit_to_1_st2_on_nights` removed; test omits it
    Rota.constraint_options_model.constrain_time_off_after_nights = False
    Rota.constraint_options_model.balance_nights_across_sites = False
    Rota.constraint_options_model.balance_shifts = True
    Rota.constraint_options_model.balance_shifts_over_workers = True
    Rota.constraint_options_model.balance_nights = False
    Rota.constraint_options_model.minimise_shift_diffs = False
    Rota.constraint_options_model.balance_blocks = False
    Rota.constraint_options_model.max_night_frequency = 0

    # Add workers
    workers = [
        Worker(name="worker1", site="group1", grade=1),
        Worker(name="worker2", site="group1", grade=1),
        Worker(name="worker3", site="group1", grade=1),
        Worker(name="worker4", site="group1", grade=1),
        Worker(name="worker5", site="group2", grade=1, fte=60),
        Worker(name="worker6", site="group2", grade=1, fte=40),
        Worker(name="worker7", site="group2", grade=1),
        Worker(name="worker8", site="group2", grade=1),
    ]
    Rota.add_workers(workers)

    # Add shifts
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="night_weekday",
            length=12.5,
            days=days[:4],
            workers_required=1,
            constraints=[
                PreShiftConstraint(days=1),
                PostShiftConstraint(days=2),
            ],
            force_as_block=True
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="night_weekend",
            length=12.5,
            days=days[4:],
            constraints=[
                PreShiftConstraint(days=1),
                PostShiftConstraint(days=2),
            ],
            force_as_block=True,
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="twilight",
            length=12.5,
            days=days[:5],
            workers_required=2,
        ),
    )

    Rota.build_shifts()
    Rota.build_workers()
    Rota.build_model()

    solver_options = {"ratio": 0.01, "seconds": 1000, "threads": 10}
    Rota.solve_model(options=solver_options)
    Rota.export_rota_to_html("test3")

    return Rota

def test_night_assignment(demo_rota_nights):
    Rota = demo_rota_nights
    assert Rota.results.solver.status == "ok"

    shift_summary = Rota.get_shift_summary_dict()

    for worker_name in shift_summary:
        worker_shifts = shift_summary[worker_name]
        worker = Rota.get_worker_by_name(worker_name)

        if worker.fte == 40:
            assert worker_shifts["night_weekday"] in (approx(0), approx(4))
            assert worker_shifts["night_weekend"] in (approx(0), approx(3))
        elif worker.fte == 60:
            assert worker_shifts["night_weekday"] in (approx(0), approx(4), approx(8))
            assert worker_shifts["night_weekend"] in (approx(0), approx(3), approx(6))
        else:
            assert worker_shifts["night_weekday"] in (approx(4), approx(8))
            assert worker_shifts["night_weekend"] in (approx(3), approx(6))