import pytest
import datetime
from copy import deepcopy
from rota_generator.shifts import RotaBuilder, SingleShift, days
from rota_generator.workers import Worker

@pytest.fixture
def rota_2workers():
    weeks_to_rota = 10
    start_date = datetime.date(2022, 3, 7)
    Rota = RotaBuilder(
        start_date,
        weeks_to_rota=weeks_to_rota,
        balance_offset_modifier=1,
    )
    worker1 = Worker(name="worker1", site="group1", grade=1)
    worker2 = Worker(name="worker2", site="group1", grade=1)
    Rota.add_workers((worker1, worker2))
    Rota.add_shifts(
        SingleShift(
            sites=("group1",),
            name="a",
            length=12.5,
            days=days,
            balance_offset=40,
            workers_required=1,
        ),
    )
    return Rota

@pytest.fixture
def rota_3workers(rota_2workers):
    Rota3 = deepcopy(rota_2workers)
    worker3 = Worker(name="worker3", site="group1", grade=1)
    Rota3.add_worker(worker3)
    return Rota3

def test_max_shifts(rota_2workers):
    rota_2workers.constraint_options_model.max_shifts_per_month = 14
    rota_2workers.build_and_solve()
    assert rota_2workers.results.solver.status == "ok"
    assert rota_2workers.results.solver.termination_condition == "optimal"

def test_max_shifts_fail(rota_2workers):
    rota_2workers.constraint_options_model.max_shifts_per_month = 13
    rota_2workers.build_and_solve()
    assert rota_2workers.results.solver.status in ("warning", "error")
    assert rota_2workers.results.solver.termination_condition == "infeasible"

def test_max_shifts_extra_worker(rota_3workers):
    rota_3workers.constraint_options_model.max_shifts_per_month = 12
    rota_3workers.build_and_solve()
    assert rota_3workers.results.solver.status == "ok"
    assert rota_3workers.results.solver.termination_condition == "optimal"

def test_max_shifts_per_week_fail(rota_2workers):
    rota_2workers.constraint_options_model.max_shifts_per_week = 3
    rota_2workers.build_and_solve()
    assert rota_2workers.results.solver.status in ("warning", "error")
    assert rota_2workers.results.solver.termination_condition == "infeasible"

def test_max_shifts_per_week_pass(rota_2workers):
    rota_2workers.constraint_options_model.max_shifts_per_month = 14
    rota_2workers.constraint_options_model.max_shifts_per_week = 4
    rota_2workers.build_and_solve()
    assert rota_2workers.results.solver.status == "ok"

def test_max_shifts_per_week_extra_worker_pass(rota_3workers):
    rota_3workers.constraint_options_model.max_shifts_per_week = 3
    rota_3workers.build_and_solve()
    assert rota_3workers.results.solver.status == "ok"