"""
Shell
=====

This module contains the Operators that can be used to interact with the shell.
"""

from nbox import Operator

class ShellCommand(Operator):
  def __init__(self, *commands):
    """Run multiple shell commands, uses ``shelex`` to prevent injection"""
    super().__init__()
    import string

    self.commands = commands
    all_in = []
    for c in self.commands:
      all_in.extend([tup[1] for tup in string.Formatter().parse(c) if tup[1] is not None])

  def forward(self, *args, **kwargs):
    import shlex
    import subprocess

    for comm in self.commands:
      comm = comm.format(*args, **kwargs)
      comm = shlex.split(comm)
      called_process: subprocess.CompletedProcess = subprocess.run(comm, check = True)
      if called_process.returncode != 0:
        raise Exception(f"Command {comm} failed with return code {called_process.stdout}")


class Python(Operator):
  def __init__(self, func, *args, **kwargs):
    """Convert a python function into an operator, everything has to be passed at runtime"""
    super().__init__()
    self.fak = (func, args, kwargs)

  def forward(self):
    return self.fak[0](*self.fak[1], **self.fak[2])


class PythonScript(Operator):
  def __init__(self, fpath, python_path = "/usr/bin/python3", **kwargs):
    """Run any file in python as an operator"""
    super().__init__()
    self.fp = fpath
    self.python_path = python_path
    self.kwargs = kwargs

    self.kwargs_strings = []
    for k,v in kwargs.items():
      if v != None:
        self.kwargs_strings.append(f"--{k}={v}")
      else:
        self.kwargs_strings.append(f"--{k}")
    kwargs = "\ \n".join(self.kwargs_strings)
    self.com = ShellCommand(f"{python_path} {fpath} {kwargs}")

  def forward(self):
    return self.com()

class NBToScript(Operator):
  def __init__(self, notebook_path: str, py_file: str) -> None:
    super().__init__()
    self.nb = notebook_path
    self.trg = py_file

  def forward(self):
    import json

    with open(self.nb, "r") as f:
      data = json.load(f)
    
    code_script = "\n# This is an Autogenerated file by NBToScript Operators\n# Find out more: nimblebox.ai\n"
    # here we can start iterating over the code cells and fill up a script
    for cell in data["cells"]:
      if cell["cell_type"] == "code":
        code_script += "".join(cell["source"]) + "\n"

    with open(self.trg, "w") as f:
      f.write(code_script)

class GitClone(Operator):
  def __init__(self, url, path = None, branch = None):
    """Clone a git repository into a path"""
    super().__init__()
    self.url = url
    self.path = path
    self.branch = branch

  def forward(self):
    # git clone -b <branch> <repo> <path>
    import subprocess
    command = ["git", "clone"]
    if self.branch:
      command.append("-b")
      command.append(self.branch)
    if self.path:
      command.append(self.path)
    command.append(self.url)
    subprocess.run(command, check = True)

