# weird file name, yessir. Why? because
# from nbox.framework.on_functions import PureFunctionParser
# read as "from nbox's framework on Functions import the pure-function Parser"

"""
This contains the code for our public static code parsers. We will be moving them internally as we expand our
scope. Till then here's all the code we have written till now. Keep a lookout for ``Timeless`` from NimbleBox.
"""

import ast
import base64
import inspect
from typing import Union
from uuid import uuid4

from ..utils import logger
from ..hyperloop.dag_pb2 import DAG, Flowchart, Node, Code, Edge, RunStatus

# ==================


# ==================
# the next set of functions are meant as support methods to create nbox_strings IR.

class NboxStrings:
  """Class to generate all the strings, for our representation."""
  OP_TO_STRING = {
    "function": "FUNCTION: {name} ( {inputs} ) => [ {outputs} ]",
    "define": "DEFINE: {name} ( {inputs} )",
    "for": "FOR: ( {iter} ) => [ {target} ]",
    "return": "RETURN: [ {value} ]",
  }

  def __init__(self):
    pass

  def function(self, name, inputs, outputs):
    """Render code for any function."""
    return self.OP_TO_STRING["function"].format(
      name=name,
      inputs=", ".join([f"{x['kwarg']}={x['value']}" for x in inputs]),
      outputs=", ".join(outputs)
    )

  def define(self, name, inputs):
    """Render code for any definition."""
    return self.OP_TO_STRING["define"].format(
      name=name,
      inputs=", ".join([f"{x['kwarg']}={x['value']}" for x in inputs])
    )

  def for_loop(self, iter, target):
    """Render code for any for loop."""
    return self.OP_TO_STRING["for"].format(
      iter=iter,
      target=", ".join(target)
    )

  def return_statement(self, value):
    """Render code for any return statement."""
    return self.OP_TO_STRING["return"].format(
      value=value
    )

nbxl = NboxStrings()

def write_program(nodes):
  """Render code for any program."""
  for i, n in enumerate(nodes):
    logger.debug(f"{i:03d}|{n.get('nbox_string', n.get('info').get('nbox_string'))}")


# ==================

def get_code_portion(cl, lineno, col_offset, end_lineno, end_col_offset, b64 = True, **_):
  """Returns code piece for any section given all the lines it contains"""
  sl, so, el, eo = lineno, col_offset, end_lineno, end_col_offset
  if sl == el:
    return cl[sl-1][so:eo]
  code = ""
  for i in range(sl - 1, el, 1):
    if i == sl - 1:
      code += cl[i][so:]
    elif i == el - 1:
      code += "\n" + cl[i][:eo]
    else:
      code += "\n" + cl[i]
  
  # convert to base64
  if b64:
    return base64.b64encode(code.encode()).decode()
  return code

def parse_args(node: ast.arguments):
  """Takes an ``ast.arguments`` list and reutrns kwarg/value pair from it."""
  inputs = []
  for a in node.args:
    a = a.arg if isinstance(a, ast.arg) else a
    inputs.append({
      "kwarg": None,
      "value": a,
    })
  for a in node.kwonlyargs:
    inputs.append({
      "kwarg": a[0],
      "value": a[1],
    })
  if node.vararg:
    inputs.append({
      "kwarg": "*"+node.vararg.arg,
      "value": None
    })
  if node.kwarg:
    inputs.append({
      "kwarg": "**"+node.kwarg.arg,
      "value": None
    })
  return inputs

def get_name(node):
  """Given a node, try to extract it's name."""
  if isinstance(node, ast.Name):
    return node.id
  elif isinstance(node, ast.Constant):
    return node.value
  elif isinstance(node, ast.Attribute):
    return get_name(node.value) + "." + node.attr
  elif isinstance(node, ast.Call):
    return get_name(node.func)

def parse_kwargs(node, lines):
  """Given a node, try to extract it's kwargs."""
  if isinstance(node, ast.Name):
    return node.id
  elif isinstance(node, ast.Constant):
    val = node.value
    return val
  elif isinstance(node, ast.keyword):
    arg = node.arg
    value = node.value
    if 'id' in value.__dict__:
      # arg = my_model
      return (arg, value.id)
    elif 'value' in value.__dict__:
      if isinstance(value.value, ast.Call):
        # arg = my_model(...)
        return (arg, get_code_portion(lines, b64 = False, **value.value.__dict__))
      else:
        # arg = 20
        return (arg, value.value)
    elif 'func' in value.__dict__:
      #   arg = some_function(with, some=args)
      #   ^^^   ^^^^^
      # kwarg   value
      return {"kwarg": arg, "value": get_code_portion(lines, b64 = False, **value.__dict__)}
  elif isinstance(node, ast.Call):
    return get_code_portion(lines, **node.func.__dict__)

def node_assign_or_expr(node, lines, node_proto: Node) -> Union[Node, None]:
  """Return ``nbox.hyperloop`` Node for assign or expressions"""
  # print(get_code_portion(lines, b64 = False, **node.__dict__))
  value = node.value
  try:
    name = get_name(value.func)
  except AttributeError:
    return None
  args = [parse_kwargs(x, lines) for x in value.args + value.keywords]
  inputs = []
  for a in args:
    if isinstance(a, dict):
      inputs.append(a)
      continue
    inputs.append({
      "kwarg": a[0] if isinstance(a, tuple) else None,
      "value": a[1] if isinstance(a, tuple) else a,
    })

  outputs = []
  if isinstance(node, ast.Assign):
    targets = node.targets[0]
    outputs = [parse_kwargs(x, lines) for x in targets.elts] \
      if isinstance(targets, ast.Tuple) \
      else [parse_kwargs(targets, lines)
    ]

  nbox_string = nbxl.function(name, inputs, outputs)

  node_proto.MergeFrom(Node(
    name = name,
    type = Node.NodeType.OP,
    info = Code(
      name = name,
      nbox_string = nbox_string,
      code = get_code_portion(lines, **node.__dict__),
      lineno = node.lineno,
      col_offset = node.col_offset,
      end_lineno = node.end_lineno,
      end_col_offset = node.end_col_offset,
      inputs = {str(x["kwarg"]): str(x["value"]) for x in inputs},
      outputs = {x: "" for x in outputs},
    )
  ))
  return node_proto

def node_if_expr(node: ast.IfExp, lines: list, node_proto: Node) -> Node:
  """Return ``nbox.hyperloop`` Node for if conditional"""
  def get_conditions(node, lines, conds = []):
    if not hasattr(node, "test"):
      else_cond = list(filter(lambda x: x["condition"] == "else", conds))
      if not else_cond:
        conds.append({
          "condition": "else",
          "code": dict(
            lineno = node.lineno,
            col_offset = node.col_offset,
            end_lineno = node.end_lineno,
            end_col_offset = node.end_col_offset,
          )
        })
      else:
        cond = else_cond[0]
        cond["code"]["end_lineno"] = node.end_lineno
        cond["code"]["end_col_offset"] = node.end_col_offset
    else:
      condition = get_code_portion(lines, **node.test.__dict__)

      # need to run this last or "else" comes up first
      conds.append({
        "condition": condition,
        "code": {
          "lineno": node.lineno,
          "col_offset": node.col_offset,
          "end_lineno": node.end_lineno,
          "end_col_offset": node.end_col_offset,
        }
      })
      for x in node.orelse:
        get_conditions(x, lines, conds)

    return conds

  all_conditions = get_conditions(node, lines, conds = [])
  ends = []
  for b0, b1  in zip(all_conditions[:-1], all_conditions[1:]):
    ends.append([b0["code"], b1["code"]])
  for i in range(len(ends)):
    ends[i] = {
      "lineno": ends[i][0]["lineno"],
      "col_offset": ends[i][0]["col_offset"],
      "end_lineno": ends[i][1]["lineno"],
      "end_col_offset": ends[i][1]["col_offset"],
    }
  ends += [all_conditions[-1]["code"]]

  conditions = []
  for i, c in enumerate(all_conditions):
    box = ends[i]
    _node = Code(
      name = f"if-{i}",
      nbox_string = c["condition"],
      lineno = box['lineno'],
      col_offset = box['col_offset'],
      end_lineno = box['end_lineno'],
      end_col_offset = box['end_col_offset'],
      inputs = [],
      outputs = [],
    )
    conditions.append(_node)

  # update the node_proto
  node_proto.MergeFrom(Node(
    name = "if",
    type = Node.NodeType.BRANCHING,
    info = Code(
      name = "if",
      nbox_string = "IF: { " + ", ".join(x.nbox_string for x in conditions) + " }",
      lineno = node.lineno,
      col_offset = node.col_offset,
      end_lineno = node.end_lineno,
      end_col_offset = node.end_col_offset,
      inputs = {},
      outputs = {},
      conditions = conditions,
    )
  ))
  return node_proto

def node_for_expr(node: ast.For, lines: list, node_proto: Node) -> Node:
  """Return ``nbox.hyperloop`` Node for as for loop."""
  iter_str = get_code_portion(lines, **node.iter.__dict__)
  if isinstance(node.target, ast.Tuple):
    targets = [x.id for x in node.target.elts]
  else:
    targets = [get_name(node.target)]
  code_ = get_code_portion(lines, **node.__dict__)
  nbox_string = nbxl.for_loop(iter_str, targets)

  # update node_proto
  node_proto.MergeFrom(Node(
    name = "for",
    type = Node.NodeType.LOOP,
    info = Code(
      name = "for",
      code = code_,
      nbox_string = nbox_string,
      lineno = node.lineno,
      col_offset = node.col_offset,
      end_lineno = node.end_lineno,
      end_col_offset = node.end_col_offset,
      inputs = {},
      outputs = {'None':x for x in targets},
    )
  ))
  return node_proto

def node_return(node: ast.Return, lines, node_proto: Node) -> Node:
  """Return ``nbox.hyperloop`` Node for a return method."""
  if isinstance(node.value, ast.Tuple):
    returns = [x.id for x in node.value.elts]
  else:
    returns = [get_name(node.value)]
  node_proto.MergeFrom(Node(
    name = "return",
    # type = Node.NodeType.RETURN,
    info = Code(
      name = "return",
      code = get_code_portion(lines, **node.__dict__),
      nbox_string = nbxl.return_statement(returns),
      lineno = node.lineno,
      col_offset = node.col_offset,
      end_lineno = node.end_lineno,
      end_col_offset = node.end_col_offset,
      inputs = {},
      outputs = {'None':'None'},
    )
  ))
  return node_proto

def def_func_or_class(node, lines, node_proto: Node) -> Node:
  """Return ``nbox.hyperloop`` Node for any function or class definition."""
  node_proto.MergeFrom(Node(
    name = f"symbol-{node.name}",
    type = Node.NodeType.OP, # TODO: @yashbonde expand node types to include class/func definitions
    info = Code(
      name = "return",
      code = get_code_portion(lines, **node.__dict__),
      nbox_string = "",
      lineno = node.lineno,
      col_offset = node.col_offset,
      end_lineno = node.end_lineno,
      end_col_offset = node.end_col_offset,
      inputs = {},
      outputs = {'None':'None'},
    )
  ))
  return node_proto

def async_def_func_or_class(node, lines, node_proto: Node) -> Node:
  """Return ``nbox.hyperloop`` Node for any asyn function."""
  print(node.__dict__)
  node_proto.MergeFrom(Node(
    name = f"symbol-{node.name}",
    type = Node.NodeType.OP, # TODO: @yashbonde expand node types to include class/func definitions
    info = Code(
      name = "return",
      code = get_code_portion(lines, **node.__dict__),
      nbox_string = "",
      lineno = node.lineno,
      col_offset = node.col_offset,
      end_lineno = node.end_lineno,
      end_col_offset = node.end_col_offset,
      inputs = {},
      outputs = {'None':'None'},
    )
  ))
  return node_proto


# ==================

type_wise_logic = {
  # defns ------
  ast.FunctionDef: def_func_or_class,
  ast.ClassDef: def_func_or_class,
  ast.AsyncFunctionDef: def_func_or_class,

  # nodes ------
  ast.Assign: node_assign_or_expr,
  ast.Expr: node_assign_or_expr,
  ast.If: node_if_expr,
  ast.For: node_for_expr,

  # Return ------
  ast.Return: node_return,

  # todos ------
  # ast.AsyncFunctionDef: async_func_def,
  # ast.Await: node_assign_or_expr,
}

# ==================


def code_node(execution_index, expr, code_lines) -> Node:
  """Return default ``nbox.hyperloop`` Node."""
  # code pieces that are not yet supported should still see the code
  return Node(
    id = str(uuid4()),
    execution_index = execution_index,
    name = f"codeblock-{execution_index}",
    type = Node.NodeType.OP,
    operator = "CodeBlock",
    nbox_string = f"CODE: {str(type(expr))}", # :str
    run_status = RunStatus(), # no need to initialise this object
    info = Code(
      name = f"codeblock-{execution_index}", # :str
      code = get_code_portion(code_lines, bs64 = True, **expr.__dict__), # :str (base64)
      nbox_string = None, # :str
      lineno = expr.lineno, # :int
      col_offset = expr.col_offset, # :int
      end_lineno = expr.end_lineno, # :int
      end_col_offset = expr.end_col_offset, # :int
      inputs = [], # :list[Dict[Any, Any]]
      outputs = [], # :list[str]
    )
  )


def get_nbx_flow(forward) -> DAG:
  """Get NBX flowchart. Read python grammar `here <https://docs.python.org/3/reference/grammar.html>`_

  Args:
      forward (callable): the function whose flowchart is to be generated
  """
  if type(forward) == str:
    code = forward
  else:
    code = inspect.getsource(forward).strip() # get code string from operator
  code_lines = code.splitlines()
  ast_node = ast.parse(code)

  edges = {} # this is the flow
  nodes = {} # this is the operators
  symbols_to_nodes = {} # this is things that are defined at runtime

  try:
    iter_obj = ast_node.body[0].body if hasattr(ast_node.body[0], 'body') else ast_node.body
    for i, expr in enumerate(iter_obj):
      # create the empty node that will be used everywhere
      if not type(expr) in type_wise_logic:
        node: Node = code_node(i, expr, code_lines)
        nodes[node.id] = node
        continue

      # define an initial proto and then then other functions will fill it up
      node_proto = Node(id = str(uuid4()), execution_index = i)
      output = type_wise_logic[type(expr)](expr, code_lines, node_proto)
      if output is None:
        continue
      nodes[node_proto.id] = node_proto

    # edges for execution order can be added
    _node_ids = tuple(nodes.keys())
    for op0, op1 in zip(_node_ids[:-1], _node_ids[1:]):
      _id = f"edge-{op0}-X-{op1}"
      edges[_id] = Edge(
        id = _id,
        source = op0,
        target = op1,
        type = Edge.EdgeType.EXECUTION_ORDER,
        nbox_string = None
      )
  except Exception as e:
    logger.error("Error in parsing the job-flow")
    logger.error(e)
    nodes = {}
    edges = {}
    symbols_to_nodes = {}

  return DAG(flowchart = Flowchart(nodes = nodes, edges = edges), symbols=symbols_to_nodes)
