import pytest
from rota_generator.shifts import NoWorkers, RequireRemoteSitePresenceConstraint, RotaBuilder, SingleShift, days

import datetime
from rota_generator.workers import Worker

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

    return RotaBuilder(
        start_date,
        weeks_to_rota=weeks_to_rota,
    )

class TestRemoteRotas:

    def test_remote_fail(self):
        Rota = setup_basic_rota()

        Rota.add_workers(
            [Worker(
                name="worker1", site="group1", grade=1, remote_site="group1"
            ),
            Worker(
                name="worker2", site="group2", grade=1, remote_site="group2"
            ),]
        )

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"), name="a", length= 12.5, days=days,
                workers_required=1,
                balance_offset=100,
                force_as_block=False,
                constraints=[RequireRemoteSitePresenceConstraint(site="group2", required_number=2)],
            ),
        )

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

        assert Rota.results.solver.status in ("warning", "error")

    def test_remote_pass_week(self):
        Rota = setup_basic_rota()
        Rota.constraint_options["balance_bank_holidays"] = False
        Rota.constraint_options["balance_shifts"] = False
        Rota.constraint_options["balance_shifts_quadratic"] = False
        Rota.constraint_options["balance_shifts_over_workers"] =True
        Rota.constraint_options["balance_blocks"] = False
        Rota.constraint_options["balance_weekends"] = False
        Rota.constraint_options["balance_nights"] = False
        Rota.constraint_options["minimse_shift_diffs"] = False

        Rota.add_workers(
            [Worker(
                name="worker1", site="group2", grade=1, remote_site="group1"
            ),
            Worker(
                name="worker2", site="group2", grade=1, remote_site="group2"
            ),Worker(
                name="worker3", site="group2", grade=1, remote_site="group2"
            ),]
        )

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"), name="a", length= 12.5, days=days[:3],
                workers_required=1,
                balance_offset=10,
                force_as_block=True,
                constraints=[RequireRemoteSitePresenceConstraint(site="group2", required_number=1)],
            ),
            SingleShift(
                sites=("group1", "group2"), name="b", length= 12.5, days=days[3:],
                workers_required=1,
                balance_offset=10,
                force_as_block=False,
                #constraint=[("require_remote_site_presence_week", ("group2", 1))],
                #constraint=["balance_shifts_over_workers"]
            ),
        )

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

        Rota.export_rota_to_html("remote_pass_week")

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

        for worker in Rota.workers:
            shifts = Rota.get_worker_shift_list(worker)

            # Convert shift to a string representation
            shifts_string = "".join([i if i != "" else "-" for i in shifts])

            if worker.site == "group1":
                assert "a" not in shifts_string

        total_shifts = Rota.get_workers_total_shifts()
        
        shift_range = max(total_shifts.values()) - min(total_shifts.values())
        assert shift_range <= 1

    def test_remote_pass(self):
        Rota = setup_basic_rota()
        Rota.constraint_options["balance_bank_holidays"] = False
        Rota.constraint_options["balance_shifts"] = False
        Rota.constraint_options["balance_shifts_quadratic"] = False
        Rota.constraint_options["balance_shifts_over_workers"] =True
        Rota.constraint_options["balance_blocks"] = False
        Rota.constraint_options["balance_weekends"] = False
        Rota.constraint_options["balance_nights"] = False
        Rota.constraint_options["minimse_shift_diffs"] = False

        Rota.add_workers(
            [Worker(
                name="worker1", site="group2", grade=1, remote_site="group1"
            ),
            Worker(
                name="worker2", site="group2", grade=1, remote_site="group2"
            ),Worker(
                name="worker3", site="group2", grade=1, remote_site="group2"
            ),]
        )

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"), name="a", length= 12.5, days=days[:3],
                workers_required=1,
                balance_offset=10,
                force_as_block=False,
                constraints=[RequireRemoteSitePresenceConstraint(site="group2", required_number=1)],
            ),
            SingleShift(
                sites=("group1", "group2"), name="b", length= 12.5, days=days[3:],
                workers_required=1,
                balance_offset=10,
                force_as_block=False,
                #constraint=[("require_remote_site_presence_week", ("group2", 1))],
                #constraint=["balance_shifts_over_workers"]
            ),
        )

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

        Rota.export_rota_to_html("remote_pass")

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

        for worker in Rota.workers:
            shifts = Rota.get_worker_shift_list(worker)

            # Convert shift to a string representation
            shifts_string = "".join([i if i != "" else "-" for i in shifts])

            if worker.site == "group1":
                assert "a" not in shifts_string

        total_shifts = Rota.get_workers_total_shifts()
        
        shift_range = max(total_shifts.values()) - min(total_shifts.values())
        assert shift_range <= 1