"""Provides `psh` persistent bash magics in Jupyter and IPython"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../00_core.ipynb.

# %% auto 0
__all__ = ['ShellInterpreter', 'shell_replace', 'PshMagic', 'create_magic', 'load_ipython_extension', 'create_ipython_config']

# %% ../00_core.ipynb
from fastcore.utils import *
import pexpect, re, os
from pexpect import TIMEOUT
from pathlib import Path
from getpass import getpass
from IPython.core.magic import register_cell_magic, no_var_expand
from IPython.display import display, Javascript
from IPython.paths import get_ipython_dir
from IPython.core.interactiveshell import InteractiveShell
from IPython.core.magic_arguments import magic_arguments, argument

# %% ../00_core.ipynb
class ShellInterpreter:
    def __init__(self, debug=False, timeout=2, shell_path=None, sudo=False, dumb=False):
        self.debug,self.timeout = debug,timeout
        if shell_path is None: shell_path = os.environ.get('SHELL', '/bin/bash')
        if sudo: shell_path = 'sudo -i ' + shell_path
        env = dict(os.environ, TERM='dumb' if dumb else 'xterm')
        self.sh = pexpect.spawn(shell_path, encoding='utf-8', env=env)
        self.sh.sendline('stty -echo')
        self.sh.readline()
        self.echo = os.urandom(8).hex()
        self.echo_re = re.compile(fr'^{self.echo}\s*$', flags=re.MULTILINE)
        self.sh.sendline(f'export PS1=""')
        self.sh.sendline(f'export PS2=""')
        self.sh.sendline('set +o vi +o emacs')
        self.wait_echo()

    def wait_echo(self, timeout=None):
        self.sh.sendline('echo')
        self.sh.sendline('echo '+self.echo)
        self.sh.expect(self.echo_re, timeout=timeout)
        return self.sh.before.rstrip()

    def _ex(self, s, timeout=None):
        if timeout is None: timeout=self.timeout
        if self.debug: print('#', s)
        self.sh.sendline(s)
        res = self.wait_echo(timeout=timeout)
        return res
        
    def __call__(self, cmd, timeout=None):
        output = self._ex(cmd.rstrip(), timeout=timeout)
        return output.replace(cmd + '\r\n', '', 1).rstrip()

# %% ../00_core.ipynb
def shell_replace(s, shell=None):
    "Replace `@{var}` refs in `s` with their variable values, if they exist"
    if not shell: shell = get_ipython()
    def f(m): return str(shell.user_ns.get(m[1], m[0]))
    return re.sub(r'\@{(\w+?)}', f, s)

# %% ../00_core.ipynb
class PshMagic:
    def __init__(self, shell, sudo=False, timeout=2, expand=True, o=None): store_attr()
    def reset(self): self.o = ShellInterpreter(sudo=self.sudo, timeout=self.timeout)
    def help (self): self.psh.parser.print_help()

    def _xpand(self, expand=False): self.expand = expand
    def _sudo(self, sudo=False):
        self.sudo = sudo
        self.o = None
    def _timeout(self, timeout=2):
        self.timeout = timeout
        self.o = None

    @magic_arguments()
    @argument('-h', '--help',      action='store_true', help='Show this help')
    @argument('-r', '--reset',     action='store_true', help='Reset the shell interpreter')
    @argument('-o', '--obj',       action='store_true', help='Return this magic object')
    @argument('-x', '--expand',    action='store_true', help='Enable variable expansion')
    @argument('-X', '--no-expand', action='store_true', help='Disable variable expansion')
    @argument('-s', '--sudo',      action='store_true', help='Enable sudo')
    @argument('-S', '--no-sudo',   action='store_true', help='Disable sudo')
    @argument('-t', '--timeout', type=int, help='Set timeout in seconds')
    @argument('command', nargs='*', help='The command to run')
    @no_var_expand
    def psh(self, line, cell=None):
        "Run line or cell in persistent shell"
        if cell: cell = shell_replace(cell, self.shell)
        if line: line = shell_replace(line, self.shell)
        args = self.psh.parser.parse_args(line.split())
        if args.expand:    return self._xpand(True)
        if args.no_expand: return self._xpand(False)
        if args.sudo:      return self._sudo (True)
        if args.no_sudo:   return self._sudo (False)
        if args.timeout:   return self._timeout(args.timeout)
        if args.reset:     return self.reset()
        if args.help:      return self.help()
        if args.obj:       return self
        if args.command: cell = ' '.join(args.command)
        if not cell and line: cell=line
        disp = True
        if cell.endswith(';'): disp,cell = False,cell[:-1]
        if not self.o: self.reset()
        try: res = self.o(cell) or None
        except Exception as e:
            self.o = None
            raise e from None
        if disp and res: print(res)

# %% ../00_core.ipynb
def create_magic(shell=None):
    if not shell: shell = get_ipython()
    magic = PshMagic(shell)
    shell.register_magic_function(magic.psh, magic_name='psh', magic_kind='line_cell')

# %% ../00_core.ipynb
def load_ipython_extension(ipython):
    "Required function for creating magic"
    create_magic(shell=ipython)

# %% ../00_core.ipynb
def create_ipython_config():
    "Called by `pshnb_install` to install magic"
    ipython_dir = Path(get_ipython_dir())
    cf = ipython_dir/'profile_default'/'ipython_config.py'
    cf.parent.mkdir(parents=True, exist_ok=True)
    if cf.exists() and 'pshnb' in cf.read_text(): return print('pshnb already installed!')
    with cf.open(mode='a') as f: f.write("\nc.InteractiveShellApp.extensions.append('pshnb.core')\n\n")
    print(f"Jupyter config updated at {cf}")
