#
# SPDX-FileCopyrightText: Hadad <hadad@linuxmail.org>
# SPDX-License-Identifier: Apache-2.0
#

import random  # Import the random module to enable random selection of elements from a list, which helps in load balancing host assignments
from datetime import datetime  # Import the datetime class to work with timestamps, particularly to check and compare current UTC time against host busy status
from typing import Dict, List  # Import type hinting classes for dictionaries and lists to improve code readability and static analysis, though not explicitly used here
from config import auth  # Import the 'auth' configuration, which is expected to be a list of host dictionaries containing credentials and identifiers for available hosts
from src.utils.helper import busy, mark  # Import 'busy', a dictionary tracking host busy states with expiration timestamps, and 'mark', a function to update the busy status of hosts when assigned

# Initialize a global dictionary named 'mapping' that will maintain a persistent association between session IDs and their assigned hosts
mapping = {}  # This dictionary stores session_id keys mapped to host dictionaries, ensuring consistent host allocation for repeated requests within the same session

# Define a function to get an available hosts for a given session
def get_host(session_id: str, exclude_hosts: List[str] = None) -> dict:
    """
    Obtain a host for the specified session ID, guaranteeing that the function never raises exceptions or returns None.
    This function implements a robust mechanism to ensure a host is always returned by dynamically excluding busy or failed hosts,
    and retrying until a suitable host becomes available.

    Args:
        session_id (str): A unique string identifier representing the current user or process session requesting a host.
        exclude_hosts (List[str], optional): An optional list of host identifiers to be excluded from selection, useful for avoiding problematic or previously failed hosts. Defaults to None.

    Returns:
        dict: A dictionary representing the selected host from the 'auth' configuration, including its credentials and identifier.

    Detailed Explanation:
        The function first checks if the session already has an assigned host in the 'mapping' dictionary. If so, it verifies whether
        the assigned host is neither excluded nor currently busy. If the assigned host is available, it refreshes its busy status to
        extend its reservation and returns it immediately, ensuring session affinity.

        If the assigned host is excluded or busy, the mapping is deleted to allow reassignment. The function then filters the list of
        all hosts from 'auth' to exclude busy hosts, explicitly excluded hosts, and hosts already tried in the current function call.
        From the filtered list, it randomly selects a host to distribute load evenly.

        If no suitable hosts are found, the function updates the set of tried hosts to include all currently busy and excluded hosts,
        and clears the exclude list to retry all hosts again. This loop continues indefinitely until a host becomes available,
        ensuring that the function never fails or returns None.

        This design supports high availability and fault tolerance by dynamically adapting to host availability and avoiding deadlocks
        or infinite retries on unavailable hosts.
    """

    # Initialize exclude_hosts as an empty list if no value is provided to avoid errors during host filtering
    if exclude_hosts is None:
        exclude_hosts = []  # Assign an empty list to exclude_hosts to safely perform membership checks and list operations later

    # Create a set to track hosts that have been attempted and found unsuitable during this invocation of the function
    tried_hosts = set()  # Using a set for efficient membership testing and to prevent repeated attempts on the same hosts within this call

    # Enter an infinite loop that will only exit when a valid host is found and returned, ensuring robustness and continuous operation
    while True:
        # Check if the current session ID already has a host assigned in the mapping dictionary
        if session_id in mapping:
            assigned_host = mapping[session_id]  # Retrieve the previously assigned host dictionary for this session to maintain session consistency

            # Determine if the assigned host is not in the exclude list and is either not busy or its busy period has expired
            if (
                assigned_host["jarvis"] not in exclude_hosts  # Confirm the assigned host is not explicitly excluded for this request
                and (
                    assigned_host["jarvis"] not in busy  # Check if the host is not currently marked as busy
                    or busy[assigned_host["jarvis"]] <= datetime.utcnow()  # Alternatively, check if the host's busy timestamp has expired, meaning it is free
                )
            ):
                # Since the assigned host is available, update its busy timestamp to extend its reservation for this session
                mark(assigned_host["jarvis"])  # Call the helper function to refresh the busy status, preventing other sessions from using it simultaneously
                return assigned_host  # Return the assigned host dictionary immediately to maintain session affinity and reduce latency

            else:
                # If the assigned host is either excluded or currently busy, remove the mapping to allow reassignment of a new host
                del mapping[session_id]  # Delete the session-to-host association to enable the selection of a different host in subsequent steps

        # Capture the current UTC time once to use for filtering hosts based on their busy status
        now = datetime.utcnow()  # Store the current time to compare against busy timestamps for all hosts

        # Generate a list of hosts that are eligible for selection by applying multiple filters:
        # 1. Hosts not currently busy or whose busy period has expired
        # 2. Hosts not included in the exclude_hosts list provided by the caller
        # 3. Hosts not already attempted in this function call to avoid redundant retries
        available_hosts = [
            h for h in auth  # Iterate over all hosts defined in the authentication configuration list
            if h["jarvis"] not in busy or busy[h["jarvis"]] <= now  # Include only hosts that are free or whose busy reservation has expired
            if h["jarvis"] not in exclude_hosts  # Exclude hosts explicitly marked to be avoided for this selection
            if h["jarvis"] not in tried_hosts  # Exclude hosts that have already been tried and failed during this function call to prevent infinite loops
        ]

        # If there are any hosts available after filtering, proceed to select one randomly
        if available_hosts:
            selected = random.choice(available_hosts)  # Randomly pick one host from the available list to distribute load fairly among hosts

            # Store the selected host in the global mapping dictionary to maintain session affinity for future requests with the same session ID
            mapping[session_id] = selected  # Cache the selected host so subsequent calls with this session ID return the same host

            # Mark the selected host as busy by updating its busy timestamp, indicating it is currently reserved for this session
            mark(selected["jarvis"])  # Update the busy dictionary to reflect that this host is now occupied, preventing concurrent use

            # Return the selected host dictionary to the caller, completing the host assignment process for the session
            return selected

        else:
            # If no hosts are available that have not already been tried, update the tried_hosts set to include all hosts currently busy or excluded
            # This prevents immediate reattempts on these hosts in the next iteration, allowing time for busy hosts to become free
            tried_hosts.update(
                h["jarvis"] for h in auth  # Iterate over all hosts in the auth list to identify those currently busy
                if h["jarvis"] in busy and busy[h["jarvis"]] > now  # Add hosts whose busy timestamp is still in the future, indicating they are occupied
            )
            tried_hosts.update(exclude_hosts)  # Also add all explicitly excluded hosts to the tried_hosts set to avoid retrying them immediately

            # Create a set of all host identifiers from the auth configuration to compare against tried_hosts
            all_host_ids = {h["jarvis"] for h in auth}  # Extract all host IDs to check if all hosts have been attempted

            # If the set of tried hosts now includes every host in the configuration, clear the tried_hosts set to reset the retry mechanism
            if tried_hosts >= all_host_ids:
                tried_hosts.clear()  # Clear the tried_hosts set to allow all hosts to be considered again, enabling retries after some time

            # Clear the exclude_hosts list as well to allow all hosts to be eligible for selection in the next iteration
            exclude_hosts.clear()  # Reset the exclude list so that no hosts are excluded in the upcoming retry cycle, maximizing availability