"""
Модуль для создания и решения модели CP-SAT
"""
from ortools.sat.python import cp_model
from typing import List, Dict, Tuple

# Дни работы
def get_days(year: int = 2026, month: int = 4) -> List[str]:
    """Генерирует список дат месяца"""
    return [f"{year}-{month:02d}-{i:02d}" for i in range(1, 31)]

# Смены
SHIFTS = [
    (6, 14.5),   # с 06:00 до 14:30
    (8, 16.5),     # с 08:00 до 16:30
    (11.5, 20),    # с 11:30 до 20:00
    (12.5, 21)      # с 12:30 до 21:00
]

# Операторы
OPERATORS = [
    "Иванов",
    "Алексеев",
    "Сидоров",
    "Козобородов",
    "Трудяшкина",
    "Петрова"
]

# Предпочтения операторов
OPERATOR_PREFERENCES = {
    "Иванов": ["2026-04-10", "2026-04-13"],
    "Алексеев": [(i, 0) for i in get_days() if (int(i[-2:]) - 1) % 7 == 0],  # Понедельники
    "Сидоров": ["2026-04-25", "2026-04-26"],
    "Козобородов": [(i, 2) for i in get_days()],  # Предпочитает смену 11:30–20:00
    "Трудяшкина": ["2026-04-04"],
    "Петрова": ["2026-04-08", "2026-04-29"]
}


def create_model() -> Tuple[cp_model.CpModel, Dict, Dict]:
    """
    Создаёт и настраивает CP-модель.
    
    Возвращает:
        model: CP-модель
        variables: словарь переменных работы
        operators_set: множество операторов для удобства
    """
    model = cp_model.CpModel()
    variables = {}
    days = get_days()
    
    # Определяем переменные: works[(day, shift_idx, op)]
    for day in days:
        for shift_idx in range(len(SHIFTS)):
            for op in OPERATORS:
                var_name = f"{op}_{day}_shift_{shift_idx}"
                variables[(day, shift_idx, op)] = model.NewBoolVar(var_name)
    
    # Добавляем ограничения
    
    # 1. Учёт предпочтений операторов
    for op, prefs in OPERATOR_PREFERENCES.items():
        if isinstance(prefs, list):
            for pref in prefs:
                if isinstance(pref, tuple):
                    day, shift = pref
                    day_key = f"{day}"
                    shift_key = shift
                    # Создаем ключ для переменных, используя стандартный формат дней
                    # Нужно найти соответствующий день в списке days
                    for d in days:
                        if d.endswith(day.split('-')[2]) and d.startswith(f"{day.split('-')[0]}-{day.split('-')[1]}"):
                            if f"{d}-{shift_key}" in pref:
                                continue
                            day, shift_idx_pref = pref
                            model.AddHint(variables[(day, shift_idx_pref, op)], 1)
                            break
        else:
            day = pref
            model.Add(sum(variables[(day, shift_idx, op)] for shift_idx in range(len(SHIFTS))) == 0)
    
    # Условие 1: На каждую смену должен быть назначен оператор
    for day in days:
        for shift_idx in range(len(SHIFTS)):
            model.Add(sum(variables[(day, shift_idx, op)] for op in OPERATORS) == 1)
    
    # Условие 2: У каждого оператора должно быть 8 выходных (20 рабочих дней)
    for op in OPERATORS:
        model.Add(sum(variables[(day, shift_idx, op)] for day in days for shift_idx in range(len(SHIFTS))) == 20)
    
    # Условие 3: Оператор не должен работать больше 5 дней подряд
    for op in OPERATORS:
        for day_idx in range(len(days) - 4):
            model.Add(sum(variables[(days[day_idx + i], shift_idx, op)] 
                         for i in range(5) 
                         for shift_idx in range(len(SHIFTS))) <= 4)
    
    # Условие 4: В день оператор может работать только одну смену
    for day in days:
        for op in OPERATORS:
            model.Add(sum(variables[(day, shift_idx, op)] for shift_idx in range(len(SHIFTS))) <= 1)
    
    return model, variables, set(OPERATORS)

def solve_model():
    """
    Решает модель.
    
    Возвращает:
        True - если решение найдено
        False - если не удалось составить график
    """
    model, variables, operators_set = create_model()
    
    solver = cp_model.CpSolver()
    status = solver.Solve(model)
    
    return status in (cp_model.OPTIMAL, cp_model.FEASIBLE)

def get_solution():
    """
    Получает решение модели.
    
    Возвращает:
        Словарь с расписанием: {день: {смена: {оператор: True/False}}}
    """
    model, variables, operators_set = create_model()
    solver = cp_model.CpSolver()
    solver.Solve(model)
    
    schedule = {}
    for day in get_days():
        schedule[day] = {
            shift_idx: {op: solver.Value(variables[(day, shift_idx, op)]) 
                        for op in OPERATORS}
            for shift_idx in range(len(SHIFTS))
        }
    
    return schedule