Scheduling Training

from datetime import datetime, timedelta
import requests

# API URLs and token
STREAM_TRAINING_ASSIGNMENT_API = 'https://api.samsara.com/training-assignments/stream'
LIST_DRIVERS_API = 'https://api.samsara.com/fleet/drivers'
GET_TRAINING_COURSES_API = 'https://api.samsara.com/training-courses'
ASSIGN_TRAINING_API = 'https://api.samsara.com/training-assignments'
API_KEY = 'Insert your API token'
HEADERS = {'Authorization': f'Bearer {API_KEY}'}

# Constants
COURSE_ID = 'Insert course ID'
TARGET_DRIVER_TAG_IDS = ['Insert the driver tag ID']
TRAINING_LOOKBACK_DAYS = -30
TRAINING_DUE_DAYS = 14
STOP_YEAR = 2030
PLANNED_EXECUTION_DATE = 'Insert date'

def get_drivers(tag_ids):
    '''Retrieve a list of drivers based on specific tags and a creation timestamp.

    Args:
        tag_ids (list of str): A list of tag IDs used to filter drivers.

    Returns:
        list: A list of driver objects that match the given tags.
    '''
    drivers = []
    after = ''
    has_next_page = True
    params = {'tagIds': tag_ids}

    while has_next_page:
        params['after'] = after
        response = requests.get(LIST_DRIVERS_API, headers=HEADERS, params=params)
        data = response.json()
        drivers.extend(data.get('data', []))

        if data.get('pagination', {}).get('hasNextPage'):
            has_next_page = True
            after = data['pagination']['endCursor']
        else:
            has_next_page = False

    return drivers

def get_training_assignment_details(assignment_interval_timestamp, course_id):
    '''
    Retrieve filtered training assignment details for a specific course and assignment timestamp range.

    Parameters:
    assignment_interval_timestamp (str): The timestamp to filter the assignment stream.
    course_id (str): The ID of the course whose assignments are being queried.

    Returns:
    list: A list of training assignment details for the specified course.
    '''
    training_assignments = []
    after = ''
    params = {'startTime': assignment_interval_timestamp, 'courseIds': [course_id]}
    has_next_page = True
    while has_next_page:
        params['after'] = after
        response = requests.get(
            STREAM_TRAINING_ASSIGNMENT_API, headers=HEADERS, params=params
        )
        data = response.json()
        training_assignments.extend(data['data'])
        if data.get('pagination', {}).get('hasNextPage', ''):
            has_next_page = True
            after = data['pagination']['endCursor']
        else:
            has_next_page = False

    return training_assignments

def assign_training(course_id, driver_ids, due_date):
    '''
    Assigns a specific training course to the provided drivers with a due date.

    Args:
        driver_ids (list): List of driver IDs to whom the training should be assigned.
        course_id (str): The ID of the course to be assigned.
        due_date (str): The due date for completing the course in RFC3339 format.

    Returns:
        None
    '''
    params = {'learnerIds': driver_ids, 'courseId': course_id, 'dueAtTime': due_date}

    requests.post(ASSIGN_TRAINING_API, headers=HEADERS, params=params)

def generate_rfc3339_timestamp(delta_days):
    '''
    Generates an RFC3339 timestamp by adding or subtracting days from the current date.

    Args:
        delta_days (int): Number of days to add (positive) or subtract (negative) from the
        current date.

    Returns:
        str: The RFC3339 formatted timestamp.
    '''
    target_date = (datetime.now() + timedelta(days=delta_days)).replace(
        hour=0, minute=0, second=0, microsecond=0
    )
    return target_date.strftime('%Y-%m-%dT%H:%M:%SZ')

def generate_milliseconds_timestamp(delta_days):
    '''
    Generates a timestamp in milliseconds format by adding or subtracting
    days from the current date.

    Args:
        delta_days (int): Number of days to add (positive) or subtract(negative)
        from the current date.

    Returns:
        int: The timestamp in milliseconds.
    '''
    target_date = (datetime.now() + timedelta(days=delta_days)).replace(
        hour=0, minute=0, second=0, microsecond=0
    )
    return int(target_date.timestamp()) * 1000

def main():
    # Convert the planned execution date from string format to a datetime object
    planned_execution_date = datetime.strptime(
        PLANNED_EXECUTION_DATE, '%d/%m/%Y'
    ).date()

    # Get the current year and the current date
    current_year = datetime.now().year
    current_date = datetime.now().date()

    # Check if the current year is greater than the stop year or the current date doesn't match the planned execution date
    if current_year > STOP_YEAR or current_date != planned_execution_date:
        return  # Exit the function if the conditions aren't met

    drivers_completed_training = []
    drivers_requiring_training = []

    # Fetch the list of drivers with the specified tags
    drivers = get_drivers(TARGET_DRIVER_TAG_IDS)

    # Generate RFC3339 timestamp for the assignment lookback period
    assignment_interval_rfc3339 = generate_rfc3339_timestamp(TRAINING_LOOKBACK_DAYS)

    # Fetch training assignment details for the lookback period
    assignment_details = get_training_assignment_details(
        assignment_interval_rfc3339, COURSE_ID
    )

    # Identify drivers who haven't completed training and add them to the list
    for assignment in assignment_details:
        if (
            assignment['learner']['type'] == 'driver'
            and assignment['status'] == 'completed'
        ):
            drivers_completed_training.append(assignment['learner']['id'])
    for driver in drivers:
        if driver['id'] not in drivers_completed_training:
            drivers_requiring_training.append(driver['id'])

    drivers_requiring_training = ','.join(
        [f'driver-{id}' for id in drivers_requiring_training]
    )

    # Generate RFC3339 timestamp for the training due date
    training_due_date_rfc3339 = generate_rfc3339_timestamp(TRAINING_DUE_DAYS)

    # Assign training to the drivers
    assign_training(COURSE_ID, drivers_requiring_training, training_due_date_rfc3339)

if __name__ == '__main__':
    main()