# 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 inspect
import traceback
from time import sleep
from threading import Thread
from functools import lru_cache

# 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 as get_op_user # << nbox imported here, fill all envs before this

@lru_cache(1)
def get_op():
  return get_op_user()

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 get_fastapi_fwd():
  from pydantic import BaseModel
  
  # first time loading the operator, will automatically create a cache miss and then all the rest will
  # be cache hits
  op = get_op() 
  
  # we use inspect signature instead of writing our own ast thing
  signature = inspect.signature(op.forward)
  forward_params = [
    str(param) for param in signature.parameters.values()
  ]

  # if your function takes in inputs then it is expected to be sent as query params, so create a pydantic
  # model and FastAPI will take care of the rest
  params_str = "\n".join([f"  {x}" for x in forward_params])
  model_name = f"{op.__class__.__qualname__}_Request"
  base_model = f'''
class {model_name}(BaseModel):
{params_str}'''
  exec(base_model)         # load the code in stack
  model = eval(model_name) # initialize from stack

  # pretty simple forward function, note that it gets operator using get_op which will be a cache hit
  code_fn = f'''
async def asgi_forward_fn(req: {model_name}):
  # the method used b
  data = req.dict()
  op = get_op()
  return op(**data)
'''
  exec(code_fn)                      # load the code in stack
  forward = eval("asgi_forward_fn")  # initialize from stack
  return forward


def print_cache_status():
  while True:
    print(get_op.cache_info())
    sleep(1)


def serve(
  host: str,
  port: int,
  debug: bool = False,
):
  from uvicorn import run
  from fastapi import FastAPI
  from fastapi.responses import JSONResponse

  # step 1: create the operator and get the needed forward method
  _forward = get_fastapi_fwd()

  async def ping_fn():
    return {"message": "pong"}

  app = FastAPI()
  app.add_api_route("/", ping_fn, methods=["GET"], response_class=JSONResponse)
  app.add_api_route("/forward", _forward, methods=["POST"], response_class=JSONResponse)  

  if debug:
    # this is for debugging purposes, it will print the cache status every second
    Thread(target = print_cache_status).start()

  run(app, host = host, port = port,)


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