"""Batch enrichment used by document read serializers."""

from django.contrib.auth import get_user_model
from django.contrib.contenttypes.models import ContentType
from django.db.models import Count

from communications.models import Message
from documents.models import Document, DocumentAcknowledgement
from employees.models import EmployeeDepartment, RoleAssignment

User = get_user_model()


def attach_linked_task_payloads(documents: list[Document], *, user) -> None:
    """Attach accessible task badges to a materialized document collection."""
    if not documents:
        return

    document_ids = [document.id for document in documents]
    mapping = {document_id: [] for document_id in document_ids}
    if user and user.is_authenticated:
        from tasks.access import task_board_access_q
        from tasks.models import TaskBoard, TaskLinkedObject, TaskLinkedObjectKind

        document_content_type = ContentType.objects.get_for_model(Document)
        accessible_boards = TaskBoard.objects.filter(is_archived=False).filter(
            task_board_access_q(user)
        )
        links = (
            TaskLinkedObject.objects.filter(
                kind=TaskLinkedObjectKind.DOCUMENT,
                content_type=document_content_type,
                object_id__in=document_ids,
                task__board__in=accessible_boards,
            )
            .select_related("task", "task__board", "task__column")
            .order_by("object_id", "task__title", "task_id")
        )
        for link in links:
            mapping.setdefault(link.object_id, []).append(
                {
                    "link_id": link.id,
                    "id": link.task_id,
                    "title": link.task.title,
                    "board_id": link.task.board_id,
                    "board_name": link.task.board.name,
                    "column_id": link.task.column_id,
                    "column_name": link.task.column.name,
                    "column_color": link.task.column.color,
                    "priority": link.task.priority,
                    "priority_display": link.task.get_priority_display(),
                }
            )

    for document in documents:
        document._linked_task_payloads = mapping.get(document.id, [])


def _department_ids(documents: list[Document]) -> set[int]:
    return {
        department.id
        for document in documents
        for department in (
            list(document.departments.all())
            + list(document.acknowledgement_departments.all())
        )
    }


def _employees_by_department(department_ids: set[int]) -> dict[int, set[int]]:
    mapping = {department_id: set() for department_id in department_ids}
    if not department_ids:
        return mapping

    membership_rows = EmployeeDepartment.objects.filter(
        department_id__in=department_ids,
        is_active=True,
        employee__is_active=True,
    ).values_list("department_id", "employee_id")
    role_rows = RoleAssignment.objects.filter(
        role__department_id__in=department_ids,
        is_active=True,
        employee__is_active=True,
    ).values_list("role__department_id", "employee_id")
    for department_id, employee_id in (*membership_rows, *role_rows):
        mapping.setdefault(department_id, set()).add(employee_id)
    return mapping


def _acknowledged_users(document_ids: list[int]) -> dict[int, set[int]]:
    mapping = {document_id: set() for document_id in document_ids}
    rows = DocumentAcknowledgement.objects.filter(
        document_id__in=document_ids,
        user__is_active=True,
    ).values_list("document_id", "user_id")
    for document_id, user_id in rows:
        mapping.setdefault(document_id, set()).add(user_id)
    return mapping


def _comment_counts(document_ids: list[int]) -> dict[int, int]:
    document_content_type = ContentType.objects.get_for_model(Document)
    return {
        row["chat__context_object_id"]: row["total"]
        for row in Message.objects.filter(
            chat__type="comments",
            chat__context_content_type=document_content_type,
            chat__context_object_id__in=document_ids,
            is_deleted=False,
        )
        .values("chat__context_object_id")
        .annotate(total=Count("id"))
    }


def _document_access_ids(
    document: Document,
    active_user_ids: set[int],
    employees_by_department: dict[int, set[int]],
) -> set[int]:
    if document.sent_to_all:
        return set(active_user_ids)
    access_ids = {
        recipient.id
        for recipient in document.recipients.all()
        if recipient.id in active_user_ids
    }
    for department in document.departments.all():
        access_ids.update(employees_by_department.get(department.id, set()))
    return access_ids


def _acknowledgement_ids(
    document: Document,
    access_ids: set[int],
    active_user_ids: set[int],
    employees_by_department: dict[int, set[int]],
) -> set[int]:
    if not document.acknowledgement_required:
        return set()
    if document.acknowledgement_for_all:
        return set(access_ids)

    required_ids = {
        recipient.id
        for recipient in document.acknowledgement_recipients.all()
        if recipient.id in active_user_ids
    }
    for department in document.acknowledgement_departments.all():
        required_ids.update(employees_by_department.get(department.id, set()))
    return required_ids & access_ids


def attach_document_metrics(documents: list[Document]) -> None:
    """Attach comment and acknowledgement aggregates without N+1 queries."""
    if not documents:
        return

    document_ids = [document.id for document in documents]
    active_user_ids = set(
        User.objects.filter(is_active=True).values_list("id", flat=True)
    )
    employees_by_department = _employees_by_department(_department_ids(documents))
    acknowledged_users = _acknowledged_users(document_ids)
    comment_counts = _comment_counts(document_ids)

    for document in documents:
        access_ids = _document_access_ids(
            document,
            active_user_ids,
            employees_by_department,
        )
        required_ids = _acknowledgement_ids(
            document,
            access_ids,
            active_user_ids,
            employees_by_department,
        )
        document._acknowledgement_user_ids = required_ids
        document._acknowledgement_total = len(required_ids)
        document._acknowledged_count = len(
            required_ids & acknowledged_users.get(document.id, set())
        )
        document._comments_count = comment_counts.get(document.id, 0)