import logging
import operator
from typing import Dict, Optional

from sqlalchemy.orm import Session, attributes

from watcher.crud.rating import get_or_create_ratings_by_schedule_staff_ids
from watcher.db import Shift, Rating, Schedule
from watcher.logic.timezone import now
from watcher import enums

logger = logging.getLogger(__name__)


def get_predicted_ratings_after_shift(shift: Shift):
    from watcher.logic.shift import shift_total_points

    predicted_ratings = shift.predicted_ratings

    if shift.staff_id:
        predicted_ratings[shift.staff.login] = (
            predicted_ratings.get(shift.staff.login, 0) +
            shift_total_points(shift)
        )

    return predicted_ratings


def update_ratings_dict(rating_diff1: Dict[str, float], rating_diff2: Dict[str, float]) -> Dict[str, float]:
    """ Changes the state of first dict """
    for key, value in rating_diff2.items():
        if key in rating_diff1:
            rating_diff1[key] += value
            # тут проверять что не меньше 0
        else:
            rating_diff1[key] = max(value, 0)

    return rating_diff1


def update_shift_participant_rating(db: Session, shift: Shift) -> Optional[Rating]:
    """
    Возвращает итоговый рейтинг дежурного после указанной смены
    """
    from watcher.logic.shift import shift_total_points

    if shift.empty or not shift.staff_id:
        return

    ratings = get_or_create_ratings_by_schedule_staff_ids(
        db=db,
        default_rating_value=0,
        composition_id=shift.slot.composition_id if shift.slot else None,
        schedule_ids=[shift.schedule_id],
        staff_ids=[shift.staff_id],
    )
    if ratings:
        rating = ratings[0]
        rating.rating = float(rating.rating) + float(shift_total_points(shift))
        return rating


def update_staff_ratings_for_further_shifts(
    session: Session,
    from_shift: Shift,
    ratings_difference: Dict[str, float]
) -> Shift:
    """
    Обновляем рейтинги сотрудников для всех дальнейших смен в последовательности.
    Возвращаем следующую смену последнего подтвержденного шифта

    :param session: сессия бд
    :param from_shift: шифт с которого нужно начать
    :param ratings_difference: словарь разницы между настоящим и предсказанным рейтингом для смены threshold_shift
    :return: следующий шифт после последнего подтвержденного шифта
    """

    all_further_shifts = (
        session.query(Shift)
        .join(Schedule, Schedule.id == Shift.schedule_id)
        .filter(
            Schedule.id == from_shift.schedule_id,
            Shift.start >= from_shift.start,
        )
    ).all()
    further_shifts = {shift.id: shift for shift in all_further_shifts}

    start_from_shift = None
    cur_shift_id = from_shift.next_id
    while cur_shift_id:
        cur_shift = further_shifts.get(cur_shift_id)
        if not cur_shift:
            logger.warning(f'Shift with id: {cur_shift_id} is in broken shift sequence')
            break
        for staff_login in ratings_difference:
            if staff_login in cur_shift.predicted_ratings:
                new_rating = cur_shift.predicted_ratings[staff_login] + ratings_difference[staff_login]
                cur_shift.predicted_ratings[staff_login] = max(new_rating, 0)
                attributes.flag_modified(cur_shift, 'predicted_ratings')  # У SQLAlchemy траблы с сохранением JSON

        if not cur_shift.approved and not start_from_shift:
            start_from_shift = cur_shift
        cur_shift_id = cur_shift.next_id

    return start_from_shift


def update_ratings_for_past_shift(db: Session, shift: Shift, new_staff_id: Optional[int]):
    """
    Забираем рейтинг у прошлого дежурного и добавляем нового
    Делаем это только для смен, которые уже прошли (иначе finish_shift это само сделает)
    """
    from watcher.logic.shift import shift_total_points

    assert shift.end < now(), 'expected finished shift'
    points = float(shift_total_points(shift=shift))
    prev_staff_id = shift.staff_id
    if shift.status != enums.ShiftStatus.completed:
        #  если смену не завершали, то и отнимать рейтинг
        #  не нужно так как мы его и не добавляли
        prev_staff_id = None
    for staff_id, func in (
        (new_staff_id, operator.add),
        (prev_staff_id, operator.sub),
    ):
        if staff_id:
            ratings = get_or_create_ratings_by_schedule_staff_ids(
                db=db,
                default_rating_value=0,
                composition_id=shift.slot.composition_id if shift.slot else None,
                schedule_ids=[shift.schedule_id],
                staff_ids=[staff_id],
            )
            if ratings:
                rating = ratings[0]
                logger.info(f'Updating rating for {staff_id} in {shift.id} with {func}')
                rating.rating = max(func(float(rating.rating), points), 0)
