from arango import ArangoClient, exceptions
import pandas as pd
import yaml

def get_documents(query=False, collection=False, fields=[], filter = '', df=False, index=False, field_names=False, limit=False):
    """
    This function retrieves documents from a specified collection or based on a query in ArangoDB.

    Parameters:
    query (str): If specified this will be the query. Defaults to False.
    collection (str): The name of the collection from which to retrieve documents. Defaults to False.
    fields (list): The fields of the documents to retrieve. If empty, all fields are retrieved.
    filter (str): AQL filter to apply to the retrieval. Defaults to no filter.
    df (bool): If True, the result is returned as a pandas DataFrame. Defaults to False.
    index (bool): If True and df is True, the DataFrame is indexed. Defaults to False.
    field_names (dict): If provided, these field names will replace the original field names in the result.

    Returns:
    list or DataFrame: The retrieved documents as a list of dictionaries or a DataFrame.
    """
    if query:
        pass

    else:
        if fields == []:
            return_fields = 'doc'
        else:
            fields_dict = {}
            for field in fields:
                fields_dict[field] = field

            if field_names:
                for k, v in field_names.items():
                    fields_dict[k] = v

            fields_list = [f'{v}: doc.{k}' for k, v in fields_dict.items()]
            fields_string = ', '.join(fields_list)
            return_fields = f"{{{fields_string}}}"
        
        if filter != None and filter != '':
            filter = f'filter {filter}'
        else:
            filter = ''

        if limit:
            limit = f'limit {limit}'
        else:           
            limit = ''

        query = f'''
            for doc
            in {collection}
            {filter}
            {limit}
            return {return_fields}
            '''
    try:
        cursor = arango_db.aql.execute(query)
    except exceptions.AQLQueryExecuteError:
        print('ERROR:\n', query)
        exit()
    result = [i for i in cursor]
    if df:
        result = pd.DataFrame(result)
        if index:
            result.set_index(index, inplace=True)
    return result


with open('config.yml', 'r') as f:
    config = yaml.safe_load(f)

db = config['arango']['db']
username = config['arango']['username']
pwd = config['arango']['pwd_lasse']

# Initialize the database for ArangoDB.
client = ArangoClient(hosts=config['arango']['hosts'])
arango_db = client.db(db, username=username, password=pwd)