# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_stata_session.ipynb.

# %% auto 0
__all__ = ['StataSession', 'variable_names']

# %% ../nbs/03_stata_session.ipynb 4
from .helpers import diverted_stata_output, run_as_program
from fastcore.basics import patch_to
from textwrap import dedent
import re

# %% ../nbs/03_stata_session.ipynb 5
class StataSession():
    def __init__(self):
        """"""
        self.matchall = re.compile(
            r"\A.*?"
            r"^%varlist%(?P<varlist>.*?)"
            r"%globals%(?P<globals>.*?)"
            r"%locals%(?P<locals>.*?)"
            r"%scalars%(?P<scalars>.*?)"
            r"%matrices%(?P<matrices>.*?)%end%", #"(\Z|---+\s*end)",
            flags=re.DOTALL + re.MULTILINE).match

        # Varlist-style matching; applies to most
        self.varlist = re.compile(r"(?:\s+)(\S+)", flags=re.MULTILINE)

        # file-style matching
        self.filelist = re.compile(r"[\r\n]{1,2}", flags=re.MULTILINE)

        # Clean line-breaks.
        self.varclean = re.compile(
            r"(?=\s*)[\r\n]{1,2}?^>\s", flags=re.MULTILINE).sub

        #         # Match output from mata mata desc
#         self.matadesc = re.compile(
#             r"(\A.*?---+|---+[\r\n]*\Z)", flags=re.MULTILINE + re.DOTALL)

#         self.matalist = re.compile(
#             r"(?:.*?)\s(\S+)\s*$", flags=re.MULTILINE + re.DOTALL)

#         self.mataclean = re.compile(r"\W.*?(\b|$)")
#         self.matasearch = re.compile(r"(?P<kw>\w.*?(?=\W|\b|$))").search

        self.clear_suggestions()
#         self.suggestions = self.get_suggestions(kernel)
#         self.suggestions['magics'] = kernel.magics.available_magics
#         self.suggestions['magics_set'] = config.all_settings

    def clear_suggestions(self):
        self.suggestions = None

# %% ../nbs/03_stata_session.ipynb 6
@patch_to(StataSession)
def refresh_suggestions(self):
    self.suggestions = self.get_suggestions()
#     self.suggestions['magics'] = kernel.magics.available_magics
#     self.suggestions['magics_set'] = config.all_settings
#     self.globals = self.get_globals(kernel)

# %% ../nbs/03_stata_session.ipynb 7
def variable_names():
    from sfi import Data
    return [Data.getVarName(i) for i in range(Data.getVarCount())]

# %% ../nbs/03_stata_session.ipynb 12
@patch_to(StataSession)
def _completions(self):
#     return dedent(f"""\
#     %varlist%
#     {' '.join(variable_names())}
#     %globals%
#     {' '.join(global_names())}
#     """
    return diverted_stata_output(dedent("""\
        local _temp_completions_while_local_ = 1
        while `_temp_completions_while_local_' {
        set more off
        set trace off
        if `"`varlist'"' != "" {
        local _temp_completions_varlist_loc_ `"`varlist'"'
        }
        syntax [varlist]
        disp "%varlist%"
        disp `"`varlist'"'
        macro drop _varlist __temp_completions_while_local_
        if `"`_temp_completions_varlist_loc_'"' != "" {
        local varlist `"`_temp_completions_varlist_loc_'"'
        macro drop __temp_completions_varlist_loc_
        }
        disp "%globals%"
        disp `"`:all globals'"'
        disp "%locals%"
        mata : invtokens(st_dir("local", "macro", "*")')
        disp "%scalars%"
        disp `"`:all scalars'"'
        disp "%matrices%"
        disp `"`:all matrices'"'
        disp "%end%"
        local _temp_completions_while_local_ = 0
        }
        macro drop _temp_completions_while_local_
    """), noecho=False)

# %% ../nbs/03_stata_session.ipynb 16
@patch_to(StataSession)
def get_suggestions(self):
    match = self.matchall(self._completions())
    suggestions = match.groupdict()
#         suggestions['mata'] = self._parse_mata_desc(suggestions['mata'])
#         suggestions['programs'] = self._parse_programs_desc(
#             suggestions['programs'])
    for k, v in suggestions.items():
#             if k in ['mata', 'programs']:
#                 continue
#             elif k in ['logfiles']:
#                 suggestions[k] = [
#                     f for f in self.filelist.split(v.strip()) if f]
#             else:
        suggestions[k] = self.varlist.findall(self.varclean('', v))
    #suggestions['locals'] = self.get_locals()
    return suggestions

# %% ../nbs/03_stata_session.ipynb 18
@patch_to(StataSession)
def get_locals(self):
    suggestions = self.get_suggestions() if self.suggestions is None else self.suggestions
    return suggestions['locals']
#     all_locals = """mata : invtokens(st_dir("local", "macro", "*")')"""
#     res = '\r\n'.join(
#         re.split(r'[\r\n]{1,2}', diverted_stata_output(all_locals)))
#     if res.strip():
#         return self.varlist.findall(self.varclean('', res))
#     else:
#         return []

# %% ../nbs/03_stata_session.ipynb 22
@patch_to(StataSession)
def run_as_prog_with_locals(self, std_code):
    """After `break_out_prog_blocks`, run noecho, inserting locals when needed"""
    from sfi import Macro
    local_defs = (f"""local {name} `"{Macro.getLocal(name)}"'"""
                  for name in self.get_locals())
    locals_code = "\n".join(local_defs)
    run_as_program(f"""{locals_code}\n{std_code}""")
