import datetime
from rota_generator.shifts import RotaBuilder, SingleShift, days
from rota_generator.workers import Worker, AvoidShiftOnDates


def build_simple_rota():
    start = datetime.date(2026, 4, 6)  # Monday
    Rota = RotaBuilder(start, weeks_to_rota=4, name="test_rota")
    Rota.add_shifts(
        SingleShift(sites=("exeter",), name="night_weekday", length=12.25, days=days[:4], workers_required=1),
        SingleShift(sites=("exeter",), name="plymouth_twilight", length=8, days=days[5:], workers_required=1),
    )
    Rota.build_shifts()
    return Rota


def test_avoid_single_date_for_worker():
    Rota = build_simple_rota()
    # create a worker who should avoid night_weekday on 2026-04-07
    avoid = AvoidShiftOnDates(shifts=["night_weekday"], dates=[datetime.date(2026, 4, 7)])
    w = Worker(name="Test Worker", site="exeter", grade=3, fte=100, avoid_shifts_on_dates=[avoid])
    Rota.add_worker(w)
    Rota.build_workers()

    # find week/day corresponding to 2026-04-07
    weeks_days = [(wk, d) for wk, d in Rota.weeks_days_product if Rota.week_day_date_map[(wk, d)] == datetime.date(2026, 4, 7)]
    assert len(weeks_days) == 1
    wk, day = weeks_days[0]

    # The model should mark works[worker.id, wk, day, 'night_weekday'] impossible -> check unavailable_to_work
    assert (w.id, wk, day) not in Rota.unavailable_to_work

    # Build and solve minimal model (but we can inspect constraints). Instead, check that the avoid rule was registered
    found = False
    for c in Rota.constraint_options_model.avoid_shifts_by_worker_dates:
        if "names" in c and w.name in c["names"]:
            found = True
            # ensure the date is represented
            assert c.get("dates") is not None and datetime.date(2026,4,7) in c.get("dates")
    assert found


def test_avoid_date_range_for_worker_applies_to_multiple_days():
    Rota = build_simple_rota()
    # Avoid night_weekday for 2026-04-06 through 2026-04-09
    avoid = AvoidShiftOnDates(shifts=["night_weekday"], start_date="06/04/2026", end_date="09/04/2026")
    w = Worker(name="Range Worker", site="exeter", grade=3, fte=100, avoid_shifts_on_dates=[avoid])
    Rota.add_worker(w)
    Rota.build_workers()

    # Collect dates between start and end
    sd = datetime.date(2026,4,6)
    ed = datetime.date(2026,4,9)
    dates_in_range = [d for d in Rota.get_date_range(sd, ed)]
    assert len(dates_in_range) == 4

    # Ensure the rota-level constraint was added
    found = False
    for c in Rota.constraint_options_model.avoid_shifts_by_worker_dates:
        if w.name in c.get("names", []):
            found = True
            assert c.get("start_date") == sd
            assert c.get("end_date") == ed
    assert found


def test_export_rota_for_visual_inspection():
    Rota = build_simple_rota()
    # add two workers
    avoid = AvoidShiftOnDates(shifts=["night_weekday"], start_date="06/04/2026", end_date="26/04/2026")
    w1 = Worker(name="Export A", site="exeter", grade=3, fte=100, avoid_shifts_on_dates=[avoid])
    w2 = Worker(name="Export B", site="exeter", grade=3, fte=100)



    Rota.add_workers((w1, w2))

    # Build model but do not solve, then export
    Rota.name = "export_test_rota"
    Rota.build_and_solve(solve=True)

    assert Rota.results.solver.status == "ok", "Rota should be feasible"

    # Next we check that worker w1 has no night_weekday shifts assigned between 06/04/2026 and 26/04/2026
    sd = datetime.date(2026,4,6)
    ed = datetime.date(2026,4,26)
    for single_date in Rota.get_date_range(sd, ed):
        wk_day = Rota.get_week_day_by_date(single_date)
        assigned_shifts = Rota.get_assigned_shifts_for_worker_on_day(w1.id, wk_day[0], wk_day[1])
        for sh in assigned_shifts:
            assert sh.name != "night_weekday", f"Worker {w1.name} was assigned forbidden shift on {single_date}"

    # We also check that w1 has 4 night shifts in total (the max possible in 4 weeks)
    night_shift_count = 0
    for wk, day in Rota.weeks_days_product:
        assigned_shifts = Rota.get_assigned_shifts_for_worker_on_day(w1.id, wk, day)
        for sh in assigned_shifts:
            if sh.name == "night_weekday":
                night_shift_count += 1
    assert night_shift_count == 4, f"Worker {w1.name} should have 4 night shifts, has {night_shift_count}"

    # And that w2 has 12 night shifts (no restrictions)
    night_shift_count_w2 = 0
    for wk, day in Rota.weeks_days_product:
        assigned_shifts = Rota.get_assigned_shifts_for_worker_on_day(w2.id, wk, day)
        for sh in assigned_shifts:
            if sh.name == "night_weekday":
                night_shift_count_w2 += 1
    assert night_shift_count_w2 == 12, f"Worker {w2.name} should have 12 night shifts, has {night_shift_count_w2}"

    # Next check that adding a short avoid to w2 results in an infeasible model
    avoid2 = AvoidShiftOnDates(shifts=["night_weekday"], start_date="06/04/2026", end_date="07/04/2026")
    w2.avoid_shifts_on_dates.append(avoid2)
    Rota.build_and_solve(solve=True)

    assert Rota.results.solver.termination_condition == "infeasible", "Rota should be infeasible due to avoid shifts on dates constraints"


    Rota.export_rota_to_html("avoid_shifts_on_dates_test.html", folder="tests")