import itertools
import logging
from collections import defaultdict

from sqlalchemy import func
from sqlalchemy.orm import Session, Query
from typing import List, Iterable, Optional

from watcher.crud.composition import (
    query_compositions_by_services,
)
from watcher.db import Rating, Schedule
from watcher.logic.exceptions import InvalidData, RatingNotFound

logger = logging.getLogger(__name__)


def query_ratings_by_schedule(db: Session, schedule_id: int, query: Query = None) -> Query:
    if not query:
        query = db.query(Rating)

    return query.filter(Rating.schedule_id == schedule_id)


def get_rating_by_schedule_staff(db: Session, schedule_id: int, staff_id: int) -> Rating:
    ratings = query_ratings_by_schedule(db, schedule_id).filter(Rating.staff_id == staff_id).all()

    if not len(ratings):
        raise RatingNotFound(message={
            'ru': f'У расписания (id={schedule_id}) еще нет рейтинга для сотрудника (id={staff_id})',
            'en': f'The schedule (id={schedule_id}) does not yet have a rating for the employee (id={staff_id})'
        })

    if len(ratings) > 1:
        raise InvalidData(
            message={
                'ru': f'Вернулось больше одного рейтинга сотрудника (id={staff_id}) внутри графика (id={schedule_id})',
                'en': f'More than one employee rating returned (id={staff_id}) for schedule (id={schedule_id})'
            }
        )

    return ratings[0]


def get_existing_ratings_by_schedule_staff_ids(db: Session, schedule_ids: Iterable[int],
                                               staff_ids: Iterable[int]) -> List[Rating]:
    """ Возвращает рейтинги сотрудников внутри расписаний. """
    return db.query(Rating).filter(
        Rating.schedule_id.in_(schedule_ids),
        Rating.staff_id.in_(staff_ids)
    ).all()


def get_new_rating_values_by_schedules(db: Session, schedule_ids: Iterable[int]) -> dict[int, float]:
    query = db.query(Rating.schedule_id, func.max(Rating.rating)).filter(
        Rating.schedule_id.in_(schedule_ids)
    ).group_by(Rating.schedule_id)
    max_ratings = {result[0]: result[1] for result in query.all()}

    new_ratings = {}
    for schedule_id in schedule_ids:
        new_ratings[schedule_id] = float(0)
        # Если все рейтинги в расписании = 0, то для нового участника тоже 0
        if schedule_id in max_ratings and max_ratings[schedule_id] != float(0):
            new_ratings[schedule_id] = float(max_ratings[schedule_id]) + 1

    return new_ratings


def _create_ratings_for_staff_without_ratings(
    db: Session,
    schedule_participants: Iterable[tuple[int, int]],  # (staff_id, schedule_id)
    default_rating_value: int | None = None,

) -> list[Rating]:
    if not schedule_participants:
        return []

    schedule_ids = {schedule_id for _, schedule_id in schedule_participants}
    new_rating_values = get_new_rating_values_by_schedules(db=db, schedule_ids=schedule_ids)

    new_ratings = list()
    for staff_id, schedule_id in schedule_participants:
        rating_value = default_rating_value if default_rating_value is not None else new_rating_values[schedule_id]
        new_ratings.append(
            Rating(
                schedule_id=schedule_id,
                staff_id=staff_id,
                rating=rating_value,
            )
        )
        logger.info(
            f'Staff {staff_id} is new in ({schedule_id=}) composition, creating rating with value {rating_value}'
        )

    if new_ratings:
        logger.info(f'Created {len(new_ratings)} new ratings for staff without ratings')
        db.add_all(new_ratings)
        db.commit()

    return new_ratings


def get_or_create_ratings_by_schedule_staff_ids(
    db: Session, schedule_ids: Iterable[int], staff_ids: Iterable[int],
    composition_id: Optional[int] = None, default_rating_value: Optional[int] = None
) -> List[Rating]:
    """
    Возвращает рейтинги сотрудников внутри расписания. Если таковых нет - создает.
    :param composition_id: Если передано, то создает рейтинг только при условии что сотрудник есть в композиции
    :param default_rating_value: Инициализирует рейтинг заданным значением, иначе *get_new_rating_value_by_schedule*
    """
    exist_ratings = get_existing_ratings_by_schedule_staff_ids(
        db=db, schedule_ids=schedule_ids, staff_ids=staff_ids,
    )
    ratings_by_schedule = defaultdict(set)
    for rating in exist_ratings:
        ratings_by_schedule[rating.schedule_id].add(rating.staff_id)

    schedule_participants_without_ratings = {
        (staff_id, schedule_id)
        for staff_id, schedule_id in itertools.product(staff_ids, schedule_ids)
        if staff_id not in ratings_by_schedule.get(schedule_id, [])
    }

    schedules_map = {
        schedule.id: schedule for schedule in
        db.query(Schedule).filter(Schedule.id.in_(schedule_ids))
    }

    compositions = query_compositions_by_services(
        db=db,
        service_ids={schedule.service_id for schedule in schedules_map.values()},
    )
    composition_participants_map = defaultdict(set)
    for composition in compositions:
        if composition_id and composition.id != composition_id:
            continue
        composition_participants_map[composition.service_id].update(
            participant.id for participant in composition.participants
        )

    # Создаем новые рейтинги только для участников композиций (сервиса или composition_id)
    participants_to_create_new_ratings = {
        (staff_id, schedule_id)
        for staff_id, schedule_id in schedule_participants_without_ratings
        if staff_id in composition_participants_map[schedules_map[schedule_id].service_id]
    }
    exist_ratings.extend(
        _create_ratings_for_staff_without_ratings(
            db=db,
            schedule_participants=participants_to_create_new_ratings,
            default_rating_value=default_rating_value,
        )
    )

    return exist_ratings
