#!/usr/bin/env python3

import sys
import logging
import json
import boto3

import argh

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.WARNING)

def set_environment_variables(func, vars, quiet=False):
    """
    Sets or deletes environment variables for an AWS Lambda function.

    Args:
        func (str): Name of the AWS Lambda function.
        vars (dict): Environment variables to set or delete.
    """
    client = boto3.client('lambda')

    # Get the existing function configuration
    response = client.get_function_configuration(FunctionName=func)
    current_vars = response.get('Environment', {}).get('Variables', {})

    if vars:
        logger.info(f"old vars: {json.dumps(current_vars, indent=4)}")

        # Set or delete the environment variables
        for var, value in vars.items():
            if value is None:
                current_vars.pop(var, None)
            else:
                current_vars[var] = value

        logger.info(f"new vars: {json.dumps(current_vars, indent=4)}")

        # Update the function configuration
        client.update_function_configuration(
            FunctionName=func,
            Environment={
                'Variables': current_vars
            }
        )

    if not quiet:
        print(json.dumps(current_vars))

@argh.arg('func', help='name of the AWS Lambda function')
@argh.arg('envvars', nargs='*', help='environment variables in key=value format or key to delete')
@argh.arg('-v', '--verbose', help='enable verbose logging')
@argh.arg('-q', '--quiet', help='do not output the new variables')
def main(func, envvars, verbose=False, quiet=False):
    """
    lambda_env.py - List, set or delete environment variables for an AWS Lambda function.

    Usage:
      lambda_env.py [OPTIONS] <func> <var=value>...
    """
    if verbose:
        logging.getLogger().setLevel(logging.INFO)

    try:
        vars = {}
        for envvar in envvars:
            if '=' in envvar:
                var, value = envvar.split('=', 1)
                vars[var] = value
            else:
                vars[envvar] = None
        set_environment_variables(func, vars, quiet=quiet)
    except Exception as e:
        logger.error(e)

if __name__ == '__main__':
    try:
        argh.dispatch_command(main)
    except Exception as e:
        logger.error(e)