from __future__ import annotations

from sqlalchemy.orm import Session, joinedload

from app.models import Admin, Submission, Task
from app.services.presence import is_online
from app.web import local_now


def online_admins(db: Session) -> list[Admin]:
    now = local_now()
    admins = db.query(Admin).filter(Admin.is_active.is_(True)).order_by(Admin.id.asc()).all()
    return [admin for admin in admins if is_online(admin.last_seen_at, now)]


def rebalance_pending_reviews(db: Session, activity_id: int | None = None) -> list[Submission]:
    pending_query = (
        db.query(Submission)
        .options(
            joinedload(Submission.user),
            joinedload(Submission.group),
            joinedload(Submission.task).joinedload(Task.activity),
            joinedload(Submission.assigned_admin),
            joinedload(Submission.reviewed_by),
        )
        .filter(Submission.status == "pending")
        .order_by(Submission.created_at.asc(), Submission.id.asc())
    )
    if activity_id is not None:
        pending_query = pending_query.join(Submission.task).filter(Task.activity_id == activity_id)

    pending_submissions = pending_query.all()
    admins = online_admins(db)
    if not admins:
        changed = False
        for submission in pending_submissions:
            if submission.assigned_admin_id is None and submission.assigned_at is None:
                continue
            submission.assigned_admin_id = None
            submission.assigned_at = None
            changed = True
        if changed:
            db.commit()
            pending_submissions = pending_query.all()
        return pending_submissions

    now = local_now()
    active_ids = {admin.id for admin in admins}
    buckets: dict[int, list[Submission]] = {admin.id: [] for admin in admins}
    changed = False

    for submission in pending_submissions:
        if submission.assigned_admin_id in active_ids:
            buckets[submission.assigned_admin_id].append(submission)

    for submission in pending_submissions:
        if submission.assigned_admin_id in active_ids:
            continue
        target_admin = min(admins, key=lambda item: (len(buckets[item.id]), item.id))
        submission.assigned_admin_id = target_admin.id
        submission.assigned_at = now
        buckets[target_admin.id].append(submission)
        changed = True

    while True:
        source_admin = max(admins, key=lambda item: (len(buckets[item.id]), -item.id))
        target_admin = min(admins, key=lambda item: (len(buckets[item.id]), item.id))
        if len(buckets[source_admin.id]) - len(buckets[target_admin.id]) <= 1:
            break

        buckets[source_admin.id].sort(key=lambda item: (item.created_at, item.id))
        moved_submission = buckets[source_admin.id].pop()
        moved_submission.assigned_admin_id = target_admin.id
        moved_submission.assigned_at = now
        buckets[target_admin.id].append(moved_submission)
        changed = True

    if changed:
        db.commit()
        pending_submissions = pending_query.all()
    return pending_submissions