#########################

# This entire file should be re-written 
# using ast or tree-hugger tree traversal
# The way it is now, can cause problems 
# in some corner cases.

#######################

import os
import shutil
from pathlib import Path
from typing import Dict
# from tabnanny import check

from docly.ioutils import CACHE_DIR, cache_exists, make_cache_dir
from docly.ioutils.convert_ipynb import convert_python_to_ipynb


def _generate_main_docstring(docstring, spaces, style, will_follow_param: bool):
    if not will_follow_param:
        if style == "google":
            return f'{spaces}"""\n{spaces}{docstring}\n\n{spaces}(Generated by docly)\n{spaces}"""\n'
        elif style == "sphinx":
            return f'{spaces}"""{docstring}\n\n{spaces}(Generated by docly)\n{spaces}"""\n'
        elif style == "numpy":
            return f"{spaces}'''\n{spaces}{docstring}\n\n{spaces}(Generated by docly)\n{spaces}'''\n"
    else:
        if style == "google":
            return f'{spaces}"""\n{spaces}{docstring}'
        elif style == "sphinx":
            return f'{spaces}"""{docstring}'
        elif style == "numpy":
            return f"{spaces}'''\n{spaces}{docstring}"


def _generate_param_list(params, style, spaces):
    line_to_write = ""
    if style == "google":
        line_to_write = line_to_write + f"\n\n{spaces}Args:\n"
        for (param_name, param_type, default_val) in params:
            param_desc = f"{param_name}"
            if param_type:
                param_desc = param_desc + f" ({param_type})"
            if default_val:
                param_desc = param_desc + f" : Defaults to {default_val}"
            if param_desc.find(":") == -1:
                param_desc = param_desc + " :"
            if len(spaces) == 4:
                line_to_write = line_to_write + f"{spaces}{spaces}{param_desc}\n"
            if len(spaces) > 4:
                line_to_write = line_to_write + f"{spaces}    {param_desc}\n"
    elif style == "numpy":
        line_to_write = line_to_write + f"\n\n{spaces}Parameters\n{spaces}----------"
        for (param_name, param_type, default_val) in params:
            param_desc = f"\n{spaces}{param_name} :"
            if param_type:
                param_desc = param_desc + f" {param_type}"
            line_to_write = line_to_write + f"{param_desc}"
    elif style == "sphinx":
        line_to_write = "\n"
        for (param_name, param_type, default_val) in params:
            param_desc = f"\n{spaces}:param {param_name}:"
            if default_val:
                param_desc = param_desc + f"\n{spaces}    defaults to {default_val}"
            if param_type:
                param_desc = param_desc + f"\n{spaces}:type {param_name}: {param_type}"
            line_to_write = line_to_write + f"{param_desc}"
    return line_to_write


def _get_line_to_write(docstrs, line_num, should_write_args_list: bool, docstring_style: str):
    docstr_line = docstrs.get(line_num+1)[1]
    num_spaces = int(docstrs.get(line_num+1)[0])
    params = docstrs.get(line_num+1)[2]
    
    line_to_write = None

    spaces = " ".join([''] * (num_spaces + 1))
    if not params:
        line_to_write = _generate_main_docstring(docstr_line, spaces, docstring_style, False)
    else:
        line_to_write = _generate_main_docstring(docstr_line, spaces, docstring_style, True)
        if should_write_args_list:
            line_to_write = line_to_write + _generate_param_list(params, docstring_style, spaces)
        if docstring_style == "google" or docstring_style == "sphinx":
            line_to_write = line_to_write + f'\n\n{spaces}(Generated by docly)\n{spaces}"""\n'
        elif docstring_style == "numpy":
            line_to_write = line_to_write + f"\n\n{spaces}(Generated by docly)\n{spaces}'''\n"
    return line_to_write


def apply_diff(docstr_loc: Dict[str, Dict[int, tuple]], should_write_args_list: bool, ipynb_files: Dict, docstring_style: str):
    try:
        for file_loc, docstrs in docstr_loc.items():
            # l = check(file_loc)
            temp_file_name = f"{str(Path(file_loc).stem)}.pytemp"
            final_file_name = f"{str(Path(file_loc).stem)}.py"
            temp_file = (Path(file_loc).parent / temp_file_name)
            final_file = (Path(file_loc).parent / final_file_name)
            write_handle = open(temp_file, "w")
            
            with open(file_loc) as f:
                for line_num, line in enumerate(f):
                    if docstrs.get(line_num+1):
                        line_to_write = _get_line_to_write(docstrs, line_num, should_write_args_list, docstring_style)
                        write_handle.write(line_to_write)
                        write_handle.write(line)
                    else:
                        write_handle.write(line)
            
            write_handle.close()
            if len(ipynb_files) == 0:
                cache_file_name = file_loc[1:].replace("/", "#")
                
                if not cache_exists():
                    make_cache_dir()
                
                if (CACHE_DIR / cache_file_name).exists():
                    (CACHE_DIR / cache_file_name).unlink()
                
                shutil.move(file_loc, str(CACHE_DIR / cache_file_name))
                shutil.move(str(temp_file), str(final_file))
            elif len(ipynb_files) > 0 and file_loc not in ipynb_files.keys():
                cache_file_name = file_loc[1:].replace("/", "#")
                
                if not cache_exists():
                    make_cache_dir()
                
                if (CACHE_DIR / cache_file_name).exists():
                    (CACHE_DIR / cache_file_name).unlink()
                
                shutil.move(file_loc, str(CACHE_DIR / cache_file_name))
                shutil.move(str(temp_file), str(final_file))
            else:
                if Path(final_file).exists():
                    Path(final_file).unlink()
                
                shutil.move(str(temp_file), str(final_file))
                convert_python_to_ipynb(Path(final_file))

                Path(final_file).unlink()
    except KeyboardInterrupt:
        temp_file.unlink()
