# Auto generated code by 'nbox job/serve upload' command
# created time: {{ created_time }}
# nbox version: {{ nbox_version }}

import os
os.environ["NBOX_JSON_LOG"] = '1' # since we finally have a pretty logging on the dashboard
os.environ["NBOX_JOB_FOLDER"] = os.path.split(os.path.abspath(__file__))[0] # do not change
os.environ["PYTHONUNBUFFERED"] = "true" # so print comes when it should come

import fire
import inspect
from functools import lru_cache, partial
from {{ file_name }} import {{ fn_name }}

from nbox import Operator, logger
import nbox.utils as U
from nbox.lib.dist import NBXLet

@lru_cache(1)
def get_op(cloud = False) -> Operator:
  # The beauty of this function is that it ensures that the operator class is loaded only once
  try:
    obj = {{ init_code }}
    if not type(obj) == Operator and {{ load_operator }}:
      # there is an initial level of precaution that we use during deployment, but we are adding simple
      # checks here as well. If it is a class, initialise it and pass it to the Operator, if function
      # just wrap it
      if inspect.isclass({{ fn_name }}):
        obj = Operator.from_class(obj)
      elif inspect.isfunction({{ fn_name }}):
        obj = Operator.from_fn(obj)
      else:
        raise ValueError("{{ fn_name }} is not an Operator or class or function")

    # at this point there is a guarantee that obj is an Operator
    op: Operator = obj
  except Exception as e:
    U.log_traceback()
    logger.error(f"Failed to load operator: {e}")
    U.hard_exit_program(1)

  if cloud:
    try:
      # loads only things that are Operators
      if {{ load_operator }}:
        op.remote_init()
    except Exception as e:
      U.log_traceback()
      logger.error(f"Failed to remote initialise operator: {e}")
      U.hard_exit_program(1)
  return op

if __name__ == "__main__":
  nbxlet = NBXLet(op = get_op(cloud = True))
  fire.Fire({
    "run": nbxlet.run,     # NBX-Jobs
    "serve": partial(
      nbxlet.serve,
      model_name = "{{ model_name }}",
      serving_type = "{{ serving_type }}"
    ), # NBX-Deploy
  })
