"""Fairness algorithms for fair task queue."""

from datetime import UTC, datetime, timedelta

from pydantic import BaseModel

from haolib.background.fair.abstract import GroupQueueState, GroupStates


class WeightedFairQueuing(BaseModel):
    """Weighted Fair Queuing algorithm."""

    def select_next_group(self, group_states: GroupStates) -> str | None:
        """Select next group to process.

        Selects the group with the minimum virtual finish time.

        Args:
            group_states: GroupStates model containing group states

        Returns:
            Group ID to process next, or None if no groups available

        """
        if not group_states.groups:
            return None

        # Filter groups with tasks
        groups_with_tasks = {
            group_id: state for group_id, state in group_states.groups.items() if state.queue_length > 0
        }

        if not groups_with_tasks:
            return None

        # Select group with minimum virtual finish time
        min_group = min(groups_with_tasks.items(), key=lambda x: x[1].virtual_finish_time)
        return min_group[0]

    def update_after_processing(
        self,
        group_state: GroupQueueState,
        processing_time: timedelta,
    ) -> None:
        """Update virtual finish time after processing task.

        VFT = max(current_time, previous_VFT) + processing_time / weight

        Args:
            group_state: The group state to update
            processing_time: Time taken to process the task

        """
        current_time = datetime.now(UTC).timestamp()
        processing_seconds = processing_time.total_seconds()
        group_state.virtual_finish_time = max(
            current_time,
            group_state.virtual_finish_time,
        ) + (processing_seconds / group_state.weight)
        group_state.last_processed_at = datetime.now(UTC)


class RoundRobin(BaseModel):
    """Round-Robin fairness algorithm.

    Uses storage to track last selected group, making it safe
    for use with multiple scheduler instances. The last selected group
    is stored in group state metadata.
    """

    def select_next_group(self, group_states: GroupStates) -> str | None:
        """Select next group to process in round-robin fashion.

        Uses metadata in group states to track last selected group.
        The metadata is persisted in storage, making it safe for distributed use.

        Args:
            group_states: GroupStates model containing group states

        Returns:
            Group ID to process next, or None if no groups available

        """
        if not group_states.groups:
            return None

        # Filter groups with tasks
        groups_with_tasks = {
            group_id: state for group_id, state in group_states.groups.items() if state.queue_length > 0
        }

        if not groups_with_tasks:
            return None

        # Get sorted list of group IDs for consistent ordering
        sorted_groups = sorted(groups_with_tasks.keys())

        # Get last selected group from any group's metadata (they should be synced)
        # We check the first group's metadata as a convention
        last_selected_group: str | None = None
        round_robin_key = "round_robin_last_group"
        if sorted_groups:
            first_group_state = group_states.groups[sorted_groups[0]]
            if round_robin_key in first_group_state.metadata.custom_data:
                last_selected_group = first_group_state.metadata.custom_data[round_robin_key]

        # If no group was selected before, start with the first one
        if last_selected_group is None or last_selected_group not in sorted_groups:
            last_selected_group = sorted_groups[0]
            # Update metadata only in the selected group's state (more efficient)
            selected_state = groups_with_tasks[last_selected_group]
            selected_state.metadata.custom_data[round_robin_key] = last_selected_group
            return last_selected_group

        # Find the index of the last selected group
        try:
            last_index = sorted_groups.index(last_selected_group)
        except ValueError:
            # Last selected group no longer exists, start from beginning
            last_selected_group = sorted_groups[0]
        else:
            # Select the next group in the list (wrap around)
            next_index = (last_index + 1) % len(sorted_groups)
            last_selected_group = sorted_groups[next_index]

        # Update metadata only in the selected group's state (more efficient)
        # Other groups will read from this when needed
        selected_state = groups_with_tasks[last_selected_group]
        selected_state.metadata.custom_data[round_robin_key] = last_selected_group

        return last_selected_group

    def update_after_processing(
        self,
        group_state: GroupQueueState,
        processing_time: timedelta,  # noqa: ARG002
    ) -> None:
        """Update group state after processing task.

        Args:
            group_state: The group state to update
            processing_time: Time taken to process the task (unused in round-robin)

        """
        group_state.last_processed_at = datetime.now(UTC)
