New Driver Training

from datetime import datetime, timedelta
import requests

# API URLs and token
LIST_DRIVERS_API = 'https://api.samsara.com/fleet/drivers'
ASSIGN_TRAINING_API = 'https://api.samsara.com/training-assignments'
STREAM_TRAINING_ASSIGNMENT_API = 'https://api.samsara.com/training-assignments/stream'
API_TOKEN = 'Insert your API token'
HEADERS = {'Accept': 'application/json', 'Authorization': f'Bearer {API_TOKEN}'}

# Constants
COURSE_ID = 'Insert course ID'  # Course ID for 'Values' course
TARGET_DRIVER_TAG_IDS = ['Insert the driver tag ID']  # Tag ID for the target driver tags
NEW_DRIVERS_LOOKBACK_DAYS = -1
COURSE_COMPLETION_DUE_DAYS = 14
TRAINING_LOOKBACK_DAYS = -1

def get_drivers(tag_ids, created_after_timestamp):
    '''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.
        created_after_timestamp (str): The timestamp (in RFC3339 format) to filter
        drivers created after this time.

    Returns:
        list: A list of driver objects that match the given tags and creation timestamp.
    '''
    drivers = []
    after = ''
    has_next_page = True
    params = {'createdAfterTime': created_after_timestamp, '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):
    '''
    Retrieves a stream of filtered training assignments within a specified time range.

    Args:
        assignment_interval_timestamp (str): Start timestamp in RFC3339 format to
        filter the training assignments.
        course_id (str): The ID of the training course to filter the assignments.

    Returns:
        list: A list of training assignments that match the provided filters.
    '''
    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(driver_id_list, course_id, due_date):
    '''Assign a training course to a list of drivers.

    Args:
        driver_id_list (list of str): A list of driver IDs to assign the training course to.
        course_id (str): The ID of the course to be assigned.
        due_date (str): The due date for the training in RFC3339 format.

    This function makes an API call to assign the training course to the specified drivers.
    '''
    params = {
        'learnerIds': ','.join([f'driver-{id}' for id in driver_id_list]),
        '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 main():
    # Generate timestamp for lookback period of new drivers
    driver_lookback_days_rfc3339 = generate_rfc3339_timestamp(NEW_DRIVERS_LOOKBACK_DAYS)

    # Get a list of newly created drivers within the lookback period
    newly_created_drivers = get_drivers(
        TARGET_DRIVER_TAG_IDS, driver_lookback_days_rfc3339
    )

    # Extract driver IDs of newly created drivers
    driver_id_list = [driver['id'] for driver in newly_created_drivers]

    drivers_assigned_training = []
    # Generate timestamp for training assignment lookback period
    training_lookback_days_rfc3339 = generate_rfc3339_timestamp(TRAINING_LOOKBACK_DAYS)

    # Get the assignment details for training
    assignment_details = get_training_assignment_details(
        training_lookback_days_rfc3339, COURSE_ID
    )

    # Loop through the assignment details to get drivers who have been assigned training
    for assignment in assignment_details:
        if assignment['learner']['type'] == 'driver':
            drivers_assigned_training.append(assignment['learner']['id'])

    # Find drivers who haven't been assigned the training
    drivers_requiring_training = [
        driver for driver in driver_id_list if driver not in drivers_assigned_training
    ]

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

    # Generate RFC3339 timestamp for course completion due date
    course_due_date_rfc3339 = generate_rfc3339_timestamp(COURSE_COMPLETION_DUE_DAYS)

    # Assign the course to newly created drivers
    assign_training(driver_id_list, COURSE_ID, course_due_date_rfc3339)

if __name__ == '__main__':
    main()