# DO NOT TOUCH
# ============
# created time: {{ created_time }}
# nbox version: {{ nbox_version }}
# Auto generated code by 'nbox deploy/serve upload' command

import io
import os
import sys
import fire
import traceback

# the trick to importing nbox is to ensure proper loading order and setting
# of correct env vars, the first import is made in the user stub
os.environ["NBOX_JOB_FOLDER"] = os.path.split(os.path.abspath(__file__))[0] # do not change

from nbx_user import get_op # << nbox imported here, fill all envs before this

import nbox.utils as U
from nbox import logger
from nbox.nbxlib.tracer import Tracer
from nbox import Operator, nbox_grpc_stub
from nbox.hyperloop.nbox_ws_pb2 import UpdateRunRequest
from nbox.hyperloop.job_pb2 import Job
from nbox.messages import rpc
from nbox.messages import get_current_timestamp


def run():
  op: Operator = get_op()
  op.propagate(_tracer = Tracer())
  if hasattr(op._tracer, "job_proto"):
    op.thaw(op._tracer.job_proto)

  run_success = True
  try:
    op()
  except Exception as e:
    f = io.StringIO("")
    traceback.print_exception(*sys.exc_info(), file = f)
    f.seek(0)
    for _l in f.readlines():
      U.logger.error(_l.rstrip())
    os._exit(1)
  finally:
    if hasattr(op._tracer, "job_proto"):
      op._tracer.job_proto.status = Job.Status.COMPLETED
      rpc(
        nbox_grpc_stub.UpdateRun, UpdateRunRequest(
          token = op._tracer.token, job=op._tracer.job_proto, updated_at=get_current_timestamp()
        ), "Failed to end job!"
      )
    os._exit(0)
  
  # why use os._exit over sys.exit:
  # https://stackoverflow.com/questions/9591350/what-is-difference-between-sys-exit0-and-os-exit0
  # https://stackoverflow.com/questions/19747371/python-exit-commands-why-so-many-and-when-should-each-be-used
  # tl;dr: os._exit kills without cleanup and so it's okay on the Pod


def serve():
  raise NotImplementedError("serve is not implemented yet")
  
  # step 1: create the operator
  op: Operator = get_op()
  
  # step 2A: create the base models
  request_model = _get_base_model(op)

  # step 2B: create the fastAPI app with all the relevant API endpoints
  app = FastAPI()
  app.add_route(
    ("GET", "/ping"),
    ("GET", "/metadata"),
    ("POST", "/predict"),
  )

  # step 3: Populate the Operator with things like thaw, etc.
  op.propagate(_tracer = Tracer())
  if hasattr(op._tracer, "job_proto"):
    op.thaw(op._tracer.job_proto)

  return app

if __name__ == "__main__":
  fire.Fire({
    "run": run,
    "serve": serve
  })
