text: #!/usr/bin/env python3
  
  # Copyright (c) Facebook, Inc. and its affiliates.
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  
  import getpass
  import glob
  import hashlib
  import netrc
  import os
  import platform
  import sh
  import shlex
  import shutil
  import subprocess
  import sys
  import time
  import requests
  import re
  from dataclasses import dataclass, field
  from omegaconf import MISSING, DictConfig
  from mephisto.operations.utils import get_mephisto_tmp_dir
  from mephisto.abstractions.architect import Architect, ArchitectArgs
  from mephisto.abstractions.architects.router.build_router import build_router
  from mephisto.abstractions.architects.channels.websocket_channel import WebsocketChannel
  from mephisto.operations.registry import register_mephisto_abstraction
  from typing import Any, Tuple, List, Dict, Optional, TYPE_CHECKING, Callable
  
  if TYPE_CHECKING:
      from mephisto.abstractions.channel import Channel
      from mephsito.data_model.packet import Packet
      from mephisto.data_model.task_run import TaskRun
      from mephisto.abstractions.database import MephistoDB
      from mephisto.abstractions.blueprint import SharedTaskState
      from argparse import _ArgumentGroup as ArgumentGroup
  
  from mephisto.operations.logger_core import get_logger
  
  logger = get_logger(name=__name__)
  
  ARCHITECT_TYPE = "heroku"
  
  USER_NAME = getpass.getuser()
  HEROKU_SERVER_BUILD_DIRECTORY = "heroku_server"
  HEROKU_CLIENT_URL = (
      "https://cli-assets.heroku.com/heroku-cli/channels/stable/heroku-cli"
  )
  
  HEROKU_WAIT_TIME = 3
  
  HEROKU_TMP_DIR = os.path.join(get_mephisto_tmp_dir(), "heroku")
  os.makedirs(HEROKU_TMP_DIR, exist_ok=True)
  
  
  @dataclass
  class HerokuArchitectArgs(ArchitectArgs):
      """Additional arguments for configuring a heroku architect"""
  
      _architect_type: str = ARCHITECT_TYPE
      use_hobby: bool = field(
          default=False, metadata={"help": "Launch on the Heroku Hobby tier"}
      )
      heroku_team: Optional[str] = field(
          default=MISSING, metadata={"help": "Heroku team to use for this launch"}
      )
      heroku_app_name: Optional[str] = field(
          default=MISSING, metadata={"help": "Heroku app name to use for this launch"}
      )
      heroku_config_args: Dict[str, str] = field(
          default_factory=dict,
          metadata={
              "help": "str:str dict containing all heroku config variables to set for the app"
          },
      )
  
  
  @register_mephisto_abstraction()
  class HerokuArchitect(Architect):
      """
      Sets up a server on heroku and deploys the task on that server
      """
  
      ArgsClass = HerokuArchitectArgs
      ARCHITECT_TYPE = ARCHITECT_TYPE
  
      def __init__(
          self,
          db: "MephistoDB",
          args: DictConfig,
          shared_state: "SharedTaskState",
          task_run: "TaskRun",
          build_dir_root: str,
      ):
          """
          Ensure heroku credentials are setup, then prepare the necessary files
          for launching for this task.
  
          All necessary paths should be built in the init or stored in the database
          such that a re-init on the same task run can pull the server information.
  
          This means that we can shutdown a server that is still running after a
          catastrophic failure.
          """
          # TODO(#102) put the expected info into the MephistoDB rather than storing here?
          # Servers will have a status which needs to be kept track of.
          self.args = args
          # print(dict(args.architect.heroku_config_args), type(dict(args.architect.heroku_config_args)))
          self.task_run = task_run
          self.deploy_name = f"{task_run.get_task().task_name}_{task_run.db_id}"
          self.build_dir = build_dir_root
          self.server_type = args.architect.server_type
          self.server_source_path = args.architect.get("server_source_path", None)
          self.heroku_config_args = dict(args.architect.heroku_config_args)
  
          # Cache-able parameters
          self.__heroku_app_name: Optional[str] = args.architect.get(
              "heroku_app_name", None
          )
          self.__heroku_executable_path: Optional[str] = None
          self.__heroku_user_identifier: Optional[str] = None
  
          self.created = False
  
      def _get_socket_urls(self) -> List[str]:
          """Returns the path to the heroku app socket"""
          heroku_app_name = self.__get_app_name()
          return ["wss://{}.herokuapp.com/".format(heroku_app_name)]
  
      def get_channels(
          self,
          on_channel_open: Callable[[str], None],
          on_catastrophic_disconnect: Callable[[str], None],
          on_message: Callable[[str, "Packet"], None],
      ) -> List["Channel"]:
          """
          Return a list of all relevant channels that the Supervisor will
          need to register to in order to function
          """
          urls = self._get_socket_urls()
          return [
              WebsocketChannel(
                  f"heroku_channel_{self.deploy_name}_{idx}",
                  on_channel_open=on_channel_open,
                  on_catastrophic_disconnect=on_catastrophic_disconnect,
                  on_message=on_message,
                  socket_url=url,
              )
              for idx, url in enumerate(urls)
          ]
  
      def download_file(self, target_filename: str, save_dir: str) -> None:
          """
          Heroku architects need to download the file
          """
          heroku_app_name = self.__get_app_name()
          target_url = (
              f"https://{heroku_app_name}.herokuapp.com/download_file/{target_filename}"
          )
          dest_path = os.path.join(save_dir, target_filename)
          r = requests.get(target_url, stream=True)
  
          with open(dest_path, "wb") as out_file:
              for chunk in r.iter_content(chunk_size=1024):
                  if chunk:
                      out_file.write(chunk)
  
      @classmethod
      def assert_task_args(cls, args: DictConfig, shared_state: "SharedTaskState"):
          """
          Assert that the provided arguments are valid. Should
          fail if a task launched with these arguments would
          not work.
  
          This should include throwing an exception if the architect
          needs login details or something similar given the
          arguments passed in.
          """
          heroku_executable_path = HerokuArchitect.get_heroku_client_path()
          try:
              output = subprocess.check_output(
                  shlex.split(heroku_executable_path + " auth:whoami")
              )
          except subprocess.CalledProcessError:
              raise Exception(
                  "A free Heroku account is required for launching tasks via "
                  "the HerokuArchitect. Please register at "
                  "https://signup.heroku.com/ and run `{} login -i` at the terminal "
                  "to login to Heroku before trying to use HerokuArchitect."
                  "".format(heroku_executable_path)
              )
          return
  
      @staticmethod
      def get_heroku_client_path() -> str:
          """
          Get the path to the heroku executable client, download a new one if it
          doesnt exist.
          """
          print("Locating heroku...")
          # Install Heroku CLI
          os_name = None
          bit_architecture = None
  
          # Get the platform we are working on
          if sys.platform == "darwin":  # Mac OS X
              os_name = "darwin"
          elif sys.platform.startswith("linux"):  # Linux
              os_name = "linux"
          else:
              os_name = "windows"
  
          # Find our architecture
          bit_architecture_info = platform.architecture()[0]
          if "64bit" in bit_architecture_info:
              bit_architecture = "x64"
          else:
              bit_architecture = "x86"
  
          # Find existing heroku files to use
          existing_heroku_directory_names = glob.glob(
              os.path.join(HEROKU_TMP_DIR, "heroku-cli-*")
          )
          if len(existing_heroku_directory_names) == 0:
              print("Getting heroku")
              if os.path.exists(os.path.join(HEROKU_TMP_DIR, "heroku.tar.gz")):
                  os.remove(os.path.join(HEROKU_TMP_DIR, "heroku.tar.gz"))
  
              # Get the heroku client and unzip
              tar_path = os.path.join(HEROKU_TMP_DIR, "heroku.tar.gz")
              sh.wget(
                  shlex.split(
                      "{}-{}-{}.tar.gz -O {}".format(
                          HEROKU_CLIENT_URL, os_name, bit_architecture, tar_path
                      )
                  )
              )
              sh.tar(shlex.split(f"-xvzf {tar_path} -C {HEROKU_TMP_DIR}"))
  
              # Clean up the tar
              if os.path.exists(tar_path):
                  os.remove(tar_path)
  
          heroku_directory_name = os.path.basename(
              glob.glob(os.path.join(HEROKU_TMP_DIR, "heroku-cli-*"))[0]
          )
          heroku_directory_path = os.path.join(HEROKU_TMP_DIR, heroku_directory_name)
          return os.path.join(heroku_directory_path, "bin", "heroku")
  
      def __get_heroku_client(self) -> Tuple[str, str]:
          """
          Get an authorized heroku client path and authorization token
          """
          if (
              self.__heroku_executable_path is None
              or self.__heroku_user_identifier is None
          ):
              heroku_executable_path = HerokuArchitect.get_heroku_client_path()
  
              # get heroku credentials
              heroku_user_identifier = None
              while not heroku_user_identifier:
                  try:
                      output = subprocess.check_output(
                          shlex.split(heroku_executable_path + " auth:whoami")
                      )
                      output = subprocess.check_output(
                          shlex.split(heroku_executable_path + " auth:token")
                      )
                      heroku_user_identifier = netrc.netrc(
                          os.path.join(os.path.expanduser("~"), ".netrc")
                      ).hosts["api.heroku.com"][0]
                  except subprocess.CalledProcessError:
                      print(
                          "A free Heroku account is required for launching MTurk tasks. "
                          "Please register at https://signup.heroku.com/ and run `{} "
                          "login -i` at the terminal to login to Heroku, and then run this "
                          "program again.".format(heroku_executable_path)
                      )
                      raise Exception("Please login to heroku before trying again.")
              self.__heroku_executable_path = heroku_executable_path
              self.__heroku_user_identifier = heroku_user_identifier
          return self.__heroku_executable_path, self.__heroku_user_identifier
  
      def __get_build_directory(self) -> str:
          """
          Return the string where the server should be built in.
          """
          return os.path.join(
              self.build_dir,
              "{}_{}".format(HEROKU_SERVER_BUILD_DIRECTORY, self.deploy_name),
          )
  
      def __get_app_name(self) -> str:
          """
          Get the name of the heroku app associated with this task
          """
          if self.__heroku_app_name is None:
              _, heroku_user_identifier = self.__get_heroku_client()
              heroku_app_name = (
                  "{}-{}-{}".format(
                      USER_NAME,
                      self.deploy_name,
                      hashlib.md5(heroku_user_identifier.encode("utf-8")).hexdigest(),
                  )
              )[:30]
              heroku_app_name = heroku_app_name.replace("_", "-")
              while heroku_app_name[-1] == "-":
                  heroku_app_name = heroku_app_name[:-1]
              self.__heroku_app_name = re.sub(r"[^a-zA-Z0-9-]", "", heroku_app_name)
          return self.__heroku_app_name
  
      def __compile_server(self) -> str:
          """
          Move the required task files to a specific directory to be deployed to
          heroku directly. Return the location that the packaged files are
          now prepared in.
          """
          print("Building server files...")
          heroku_server_development_root = self.__get_build_directory()
          os.makedirs(heroku_server_development_root)
          heroku_server_development_path = self.server_dir = build_router(
              heroku_server_development_root,
              self.task_run,
              version=self.server_type,
              server_source_path=self.server_source_path,
          )
          return heroku_server_development_path
  
      def __setup_heroku_server(self) -> str:
          """
          Deploy the server using the setup server directory, return the URL
          """
  
          heroku_executable_path, heroku_user_identifier = self.__get_heroku_client()
          server_dir = self.__get_build_directory()
  
          print("Heroku: Starting server...")
  
          heroku_server_directory_path = os.path.join(server_dir, "router")
          sh.git(shlex.split(f"-C {heroku_server_directory_path} init"))
  
          heroku_app_name = self.__get_app_name()
  
          # Create or attach to the server
          return_dir = os.getcwd()
          os.chdir(heroku_server_directory_path)
          try:
              if self.args.architect.get("heroku_team", None) is not None:
                  subprocess.check_output(
                      shlex.split(
                          "{} create {} --team {}".format(
                              heroku_executable_path,
                              heroku_app_name,
                              self.args.architect.heroku_team,
                          )
                      )
                  )
              else:
                  subprocess.check_output(
                      shlex.split(
                          "{} create {}".format(heroku_executable_path, heroku_app_name)
                      )
                  )
                  self.created = True
          except subprocess.CalledProcessError as e:  # User has too many apps?
              # TODO(#93) check response codes to determine what actually happened
              logger.exception(e, exc_info=True)
              sh.rm(shlex.split("-rf {}".format(heroku_server_directory_path)))
              raise Exception(
                  "An exception has occurred when launching your heroku app. This "
                  "can commonly occur when you have hit your limit on concurrent "
                  "apps in heroku, especially if you are running multiple tasks "
                  "at once. It also may occur if the app-name generated for your "
                  "task is using illegal characters for heroku. Check the logs "
                  "above for confirmation.\n"
                  "If the issue is indeed the concurrent server cap, Please wait for"
                  " some of your existing tasks to complete. If you have no tasks "
                  "running, login to heroku and delete some of the running apps or "
                  "verify your account to allow more concurrent apps"
              )
  
          # Enable WebSockets
          try:
              subprocess.check_output(
                  shlex.split(
                      "{} features:enable http-session-affinity".format(
                          heroku_executable_path
                      )
                  )
              )
          except subprocess.CalledProcessError:  # Already enabled WebSockets
              pass
          os.chdir(return_dir)
  
          # push config args, as desired
          if len(self.heroku_config_args) > 0:
              config_strs = [
                  f"{config_key}={config_val}"
                  for config_key, config_val in self.heroku_config_args.items()
              ]
              full_config_str = " ".join(config_strs)
              subprocess.check_output(
                  shlex.split(
                      f"{heroku_executable_path} config:set -a {heroku_app_name} {full_config_str}"
                  )
              )
  
          # commit and push to the heroku server
          sh.git(shlex.split(f"-C {heroku_server_directory_path} add -A"))
          sh.git(shlex.split(f'-C {heroku_server_directory_path} commit -m "app"'))
          sh.git(shlex.split(f"-C {heroku_server_directory_path} push -f heroku master"))
  
          os.chdir(heroku_server_directory_path)
          subprocess.check_output(
              shlex.split("{} ps:scale web=1".format(heroku_executable_path))
          )
  
          if self.args.architect.use_hobby is True:
              try:
                  subprocess.check_output(
                      shlex.split("{} dyno:type Hobby".format(heroku_executable_path))
                  )
              except subprocess.CalledProcessError:  # User doesn't have hobby access
                  self.__delete_heroku_server()
                  sh.rm(shlex.split("-rf {}".format(heroku_server_directory_path)))
                  raise Exception(
                      "Server launched with hobby flag but account cannot create "
                      "hobby servers."
                  )
          os.chdir(return_dir)
  
          time.sleep(HEROKU_WAIT_TIME)
  
          return "https://{}.herokuapp.com".format(heroku_app_name)
  
      def __delete_heroku_server(self):
          """
          Remove the heroku server associated with this task run
          """
          heroku_executable_path, heroku_user_identifier = self.__get_heroku_client()
          heroku_app_name = self.__get_app_name()
          print("Heroku: Deleting server: {}".format(heroku_app_name))
          subprocess.check_output(
              shlex.split(
                  "{} destroy {} --confirm {}".format(
                      heroku_executable_path, heroku_app_name, heroku_app_name
                  )
              )
          )
          time.sleep(HEROKU_WAIT_TIME)
  
      def server_is_running(self) -> bool:
          """
          Utility function to check if the given heroku app (by app-name) is
          still running
          """
          heroku_executable_path, _token = self.__get_heroku_client()
          app_name = self.__get_app_name()
          output = subprocess.check_output(shlex.split(heroku_executable_path + " apps"))
          all_apps = str(output, "utf-8")
          return app_name in all_apps
  
      def build_is_clean(self) -> bool:
          """
          Utility function to see if the build has been cleaned up
          """
          server_dir = self.__get_build_directory()
          return not os.path.exists(server_dir)
  
      def prepare(self) -> str:
          """
          Produce the server files that will be deployed to the server
          """
          return self.__compile_server()
  
      def deploy(self) -> str:
          """
          Launch the server, and push the task files to the server. Return
          the server URL
          """
          return self.__setup_heroku_server()
  
      def cleanup(self) -> None:
          """
          Remove any files that were used for the deployment process that
          no longer need to be kept track of now that the task has
          been launched.
          """
          server_dir = self.__get_build_directory()
          sh.rm(shlex.split("-rf {}".format(server_dir)))
  
      def shutdown(self) -> None:
          """
          Shut down the server launched by this Architect, as stored
          in the db.
          """
          if self.created:  # only delete the server if it's created by us
              self.__delete_heroku_server()
  