import datetime
import pytest
from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, days

from rota.workers import Worker


def generate_basic_rota(weeks_to_rota=10):
    start_date = datetime.date(2022, 3, 7)

    Rota = RotaBuilder(
        start_date,
        weeks_to_rota=weeks_to_rota,
    )

    # Add a few workers
    Rota.add_workers(
        [
            Worker(name="worker1", site="group1", grade=1),
            Worker(name="worker2", site="group1", grade=1),
        ]
    )

    return Rota


def date_generator(from_date, days):
    n = 0
    while True:
        yield from_date

        n = n + 1
        if n >= days:
            break
        from_date = from_date + datetime.timedelta(days=1)


class TestNightShifts:
    def test_nights(self):
        Rota = generate_basic_rota()

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"),
                name="n",
                length=12.5,
                days=days[5:],
                force_as_block=True,
                constraint=[{"name": "night"}, {"name": "pre", "options": 2}],
                workers_required=2,
            ),
            SingleShift(
                sites=("group1", "group2"),
                name="b",
                length=12.5,
                days=days[:3],
                force_as_block=False,
                # constraint=[{"name": "night"}],
            ),
        )

        Rota.build_and_solve(options={"ratio": 0.000})
        Rota.export_rota_to_html("night")

        assert Rota.results.solver.status == "ok"
        assert Rota.results.solver.termination_condition == "optimal"

    def test_nights_pre3(self):
        Rota = generate_basic_rota()

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"),
                name="n",
                length=12.5,
                days=days[5:],
                force_as_block=True,
                constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
                workers_required=2,
            ),
            SingleShift(
                sites=("group1", "group2"),
                name="b",
                length=12.5,
                days=days[:2],
                force_as_block=False,
                # constraint=[{"name": "night"}],
            ),
        )

        Rota.build_and_solve(options={"ratio": 0.000})
        Rota.export_rota_to_html("night_pre3")

        assert Rota.results.solver.status == "ok"

    def test_nights_pre3_2(self):
        Rota = generate_basic_rota()

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"),
                name="n",
                length=12.5,
                days=days[5:],
                force_as_block=True,
                constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
                workers_required=1,
            ),
            SingleShift(
                sites=("group1", "group2"),
                name="b",
                length=12.5,
                days=days[:3],
                force_as_block=False,
                # constraint=[{"name": "night"}],
            ),
        )

        Rota.build_and_solve(options={"ratio": 0.000})
        Rota.export_rota_to_html("night_pre3")

        assert Rota.results.solver.status == "ok"

    def test_nights_fail(self):
        Rota = generate_basic_rota()

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"),
                name="n",
                length=12.5,
                days=days[5:],
                force_as_block=True,
                constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
                workers_required=2,
            ),
            SingleShift(
                sites=("group1", "group2"),
                name="b",
                length=12.5,
                days=days[:3],
                force_as_block=False,
                # constraint=[{"name": "night"}],
            ),
        )

        Rota.build_and_solve(options={"ratio": 0.000})

        assert Rota.results.solver.status == "warning"

    def test_nights2(self):
        Rota = generate_basic_rota()

        Rota.constraint_options["max_night_frequency"] = 0

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"),
                name="n",
                length=12.5,
                days=days[5:],
                force_as_block=True,
                constraint=[{"name": "night"}],
                workers_required=2,
            ),
            SingleShift(
                sites=("group1", "group2"),
                name="m",
                length=12.5,
                days=days[:5],
                force_as_block=True,
                # constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
                # constraint=[{"name": "night"}],
                workers_required=2,
            ),
            # SingleShift(
            #    sites=("group1", "group2"), name="b", length= 12.5, days=days[:3],
            #    force_as_block=False,
            #    #constraint=[{"name": "night"}],
            # ),
        )

        Rota.build_and_solve(options={"ratio": 0.000})
        Rota.export_rota_to_html("night")

        assert Rota.results.solver.status == "ok"

    def test_nights_pre_wrap_around(self):
        Rota = generate_basic_rota(weeks_to_rota=8)

        Rota.constraint_options["max_night_frequency"] = 0

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"),
                name="n",
                length=12.5,
                days=days[:2],
                force_as_block=True,
                constraint=[{"name": "night"}, 
                {"name": "pre", "options": 5},{"name": "post", "options": 5},
                ],
                workers_required=3,
            ),
            # SingleShift(
            #    sites=("group1", "group2"), name="m", length= 12.5, days=days[:5],
            #    force_as_block=True,
            #    #constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
            #    #constraint=[{"name": "night"}],
            #    workers_required=2,
            # ),
            # SingleShift(
            #    sites=("group1", "group2"), name="b", length= 12.5, days=days[:3],
            #    force_as_block=False,
            #    #constraint=[{"name": "night"}],
            # ),
        )

        Rota.add_workers(
            [
                Worker(
                    name="worker3",
                    site="group1",
                    grade=1,
                    not_available_to_work=[
                        {
                            "date": Rota.rota_end_date - datetime.timedelta(days=1),
                            "reason": "no wanna",
                        },
                        #{
                        #    "date": Rota.rota_end_date - datetime.timedelta(days=5),
                        #    "reason": "no wanna",
                        #},
                    ],
                ),
            ]
        )

        Rota.build_and_solve(options={"ratio": 0.000})
        Rota.export_rota_to_html("night")

        assert Rota.results.solver.status == "ok"