# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../pts/api/export/01_export_as_func.pct.py.

# %% auto 0
__all__ = ['export_to_lib_as_func']

# %% ../../../pts/api/export/01_export_as_func.pct.py 3
import re
import tempfile
import nbformat
from pathlib import Path
import ast

from . import convert_nb, get_nb_directives, lookup_directive, export_to_lib, get_nb_module_export_path

# %% ../../../pts/api/export/01_export_as_func.pct.py 6
def get_nb_as_py_file(nb_path: str, lib_name: str, nb_format=None):
    with tempfile.TemporaryDirectory() as tmp_dir:  # Create a temporary directory instead
        temp_lib_path = Path(tmp_dir) / lib_name
        temp_lib_path.mkdir(parents=True, exist_ok=True)
        export_to_lib(nb_path, temp_lib_path, nb_format)
        py_file_path = get_nb_module_export_path(nb_path, temp_lib_path)
        py_file_content = Path(py_file_path).read_text()
        return py_file_content
    raise Exception('Failed to get the content of the notebook as a python file')

# %% ../../../pts/api/export/01_export_as_func.pct.py 9
def get_top_exports(nb_path: str, nb_format=None):
    """
    Get the content of the notebook as a python file
    """
    directives = get_nb_directives(nb_path, nb_format)
    header_codes = [f"# %% top_export\n{d['cell']['source_without_directives']}" for d in directives if d['directive'] == 'top_export']
    return "\n\n".join(header_codes)

# %% ../../../pts/api/export/01_export_as_func.pct.py 14
def export_to_lib_as_func(nb_path: str, lib_path: str, nb_format: str = None):
    # Get the function signature from the notebook    
    directive = lookup_directive(get_nb_directives(nb_path, nb_format), 'set_func_signature')
    func_sig = directive['cell']['source_without_directives'].strip()    
    if func_sig.endswith('...'): func_sig = func_sig[:-3]
    elif func_sig.endswith('pass'): func_sig = func_sig[:-4]
    else: raise Exception('Invalid function signature')
    
    # Get the return cell from the notebook
    directive = lookup_directive(get_nb_directives(nb_path, nb_format), 'func_return')
    if directive is not None:
        return_cell = directive['cell']['source_without_directives'].strip()
        if len(return_cell.split('\n')) > 1:
            raise Exception('Return cell must contain only one line')
        return_statement = f'return {return_cell}'
    else:
        return_statement = 'return'
     
    # Get the content of the notebook as a python file
    lib_name = Path(lib_path).stem
    py_file_content = get_nb_as_py_file(nb_path, lib_name, nb_format)
    
    # Get the function header content
    header_content = get_top_exports(nb_path, nb_format)
    
    # Construct the function
    first_line = py_file_content.split('\n')[0] # The first line contains the 'AUTOGENERATED!...' comment
    func_body = "\n".join([f"    {l}" for l in py_file_content.splitlines()])
    
    py_func_file_content = f"""
{first_line}

{header_content}

{func_sig}
{func_body}
    {return_statement}
    """.strip()

    # Check the syntax
    ast.parse(py_func_file_content)

    # Export to the library
    py_file_path = get_nb_module_export_path(nb_path, lib_path)
    Path(py_file_path).parent.mkdir(parents=True, exist_ok=True)
    with open(py_file_path, 'w') as f:
        f.write(py_func_file_content)
