import json

def get_code_from_notebook(filename):
    f = open(filename, "r")
    data = json.loads(f.read())
    code = ''
    for cell in data['cells']:
        for line in cell['source']:
            code = code + line
        code = code + '\n'
    return code

def add_headers(codebase, function_name, add_code_for_caching = False, is_query=False):
    head_before = f"CREATE OR REPLACE FUNCTION {function_name}() " + '''
    RETURNS TEXT
    AS $$ '''

    code_for_caching = f'''
# later replace this with one of the cache_back.py functions
# this code will get all the pandas variables from the running code and cache them into DBMS
# user should have already imported using `from cacheback2 import cache_back`
cache_tables()
    '''

    head_after = '''$$ LANGUAGE plpython3u;
    '''
    final_query = ''
    if add_code_for_caching: final_query = head_before + '\n' + codebase + '\n' + code_for_caching + '\n' + head_after
    else: final_query = head_before + '\n' + codebase + '\n' + head_after

    if is_query:
        final_query = final_query.replace("'", "''")
    return final_query

def generate_query(notebook, function_name, add_code_for_caching = False):
    return add_headers(get_code_from_notebook(notebook), function_name, add_code_for_caching)

def comment_line_by_var_usage(var_name, codebase):
    '''
    This needs to comment out all the variable initialisations and in-place methods
    BEFORE caching of that variable occurs and NOT AFTER.
    '''
    loc = codebase.split('\n')
    updated_code = ''
    for line in loc:
        temp_line = line.replace(" ", "")
        if temp_line.startswith(var_name+'=') or temp_line.startswith(var_name+'.'):
            line = '# ' + line
        updated_code = updated_code + line + '\n'
    return updated_code

def get_notebook_name(notebook_path):
    return notebook_path.split('.ipynb')[0]
