text: # -*- coding: utf-8 -*-
  """Tasks controller."""
  import base64
  import json
  import os
  import pkgutil
  import re
  import tempfile
  from datetime import datetime
  from typing import Optional
  
  from fastapi_mail import FastMail, MessageSchema
  from jinja2 import Template
  from sqlalchemy import asc, desc, func
  
  from projects import __version__, models, schemas
  from projects.controllers.utils import uuid_alpha
  from projects.exceptions import BadRequest, Forbidden, InternalServerError, NotFound
  from projects.kubernetes.notebook import (
      copy_file_to_pod,
      get_files_from_task,
      update_persistent_volume_claim,
      update_task_config_map,
  )
  from projects.kfp.tasks import make_task_creation_job, make_task_deletion_job
  from projects.kfp import KF_PIPELINES_NAMESPACE
  
  
  PREFIX = "tasks"
  
  CATEGORY_DEFAULT = "DEFAULT"
  CATEGORY_DATASETS = "DATASETS"
  CATEGORY_DESCRIPTIVE_STATISTICS = "DESCRIPTIVE_STATISTICS"
  CATEGORY_FEATURE_ENGINEERING = "FEATURE_ENGINEERING"
  CATEGORY_PREDICTOR = "PREDICTOR"
  CATEGORY_COMPUTER_VISION = "COMPUTER_VISION"
  CATEGORY_NLP = "NLP"
  CATEGORY_MONITORING = "MONITORING"
  
  VALID_CATEGORIES = [
      CATEGORY_DEFAULT,
      CATEGORY_DATASETS,
      CATEGORY_DESCRIPTIVE_STATISTICS,
      CATEGORY_FEATURE_ENGINEERING,
      CATEGORY_PREDICTOR,
      CATEGORY_COMPUTER_VISION,
      CATEGORY_NLP,
      CATEGORY_MONITORING,
  ]
  DEPLOYMENT_NOTEBOOK = json.loads(
      pkgutil.get_data("projects", "config/Deployment.ipynb")
  )
  EXPERIMENT_NOTEBOOK = json.loads(
      pkgutil.get_data("projects", "config/Experiment.ipynb")
  )
  
  TASK_DEFAULT_EXPERIMENT_IMAGE = os.getenv(
      "TASK_DEFAULT_EXPERIMENT_IMAGE",
      f"platiagro/platiagro-experiment-image:{__version__}",
  )
  TASK_DEFAULT_CPU_LIMIT = os.getenv("TASK_DEFAULT_CPU_LIMIT", "2000m")
  TASK_DEFAULT_CPU_REQUEST = os.getenv("TASK_DEFAULT_CPU_REQUEST", "100m")
  TASK_DEFAULT_MEMORY_LIMIT = os.getenv("TASK_DEFAULT_MEMORY_LIMIT", "10Gi")
  TASK_DEFAULT_MEMORY_REQUEST = os.getenv("TASK_DEFAULT_MEMORY_REQUEST", "2Gi")
  TASK_DEFAULT_READINESS_INITIAL_DELAY_SECONDS = int(
      os.getenv(
          "TASK_DEFAULT_READINESS_INITIAL_DELAY_SECONDS",
          "60",
      )
  )
  
  EMAIL_MESSAGE_TEMPLATE = pkgutil.get_data("projects", "config/email-template.html")
  
  NOT_FOUND = NotFound(code="TaskNotFound", message="The specified task does not exist")
  
  
  class TaskController:
      def __init__(self, session, background_tasks=None):
          self.session = session
          self.background_tasks = background_tasks
  
      def raise_if_task_does_not_exist(self, task_id: str):
          """
          Raises an exception if the specified task does not exist.
  
          Parameters
          ----------
          task_id : str
  
          Raises
          ------
          NotFound
          """
          exists = (
              self.session.query(models.Task.uuid).filter_by(uuid=task_id).scalar()
              is not None
          )
  
          if not exists:
              raise NOT_FOUND
  
      def list_tasks(
          self,
          page: Optional[int] = None,
          page_size: Optional[int] = None,
          order_by: str = Optional[str],
          **filters,
      ):
          """
          Lists tasks. Supports pagination, and sorting.
          Parameters
          ----------
          page : int
              The page number. First page is 1.
          page_size : int
              The page size.
          order_by : str
              Order by instruction. Format is "column [asc|desc]".
          **filters : dict
          Returns
          -------
          projects.schemas.task.TaskList
          Raises
          ------
          BadRequest
              When order_by is invalid.
          """
          query = self.session.query(models.Task)
          query_total = self.session.query(func.count(models.Task.uuid))
  
          for column, value in filters.items():
              query = query.filter(getattr(models.Task, column).ilike(f"%{value}%"))
              query_total = query_total.filter(
                  getattr(models.Task, column).ilike(f"%{value}%")
              )
  
          total = query_total.scalar()
  
          # Default sort is name in ascending order
          if not order_by:
              order_by = "name asc"
  
          # Sorts records
          try:
              (column, sort) = order_by.replace("+", " ").strip().split()
              assert sort.lower() in ["asc", "desc"]
              assert column in models.Task.__table__.columns.keys()
          except (AssertionError, ValueError):
              raise BadRequest(code="InvalidOrderBy", message="Invalid order argument")
  
          if sort.lower() == "asc":
              query = query.order_by(asc(getattr(models.Task, column)))
          elif sort.lower() == "desc":
              query = query.order_by(desc(getattr(models.Task, column)))
  
          if page and page_size:
              # Applies pagination
              query = query.limit(page_size).offset((page - 1) * page_size)
  
          tasks = query.all()
          return schemas.TaskList.from_orm(tasks, total)
  
      def generate_name_task(self, name, attempt=1):
          name_task = f"{name} - {attempt}"
          check_comp_name = (
              self.session.query(models.Task).filter_by(name=name_task).first()
          )
          if check_comp_name:
              return self.generate_name_task(name, attempt + 1)
          return name_task
  
      def task_category_is_not_none(self, task_cat):
          if task_cat.category is not None and task_cat.category not in VALID_CATEGORIES:
              valid_str = ",".join(VALID_CATEGORIES)
              raise BadRequest(
                  code="InvalidCategory",
                  message=f"Invalid category. Choose any of {valid_str}",
              )
  
      def create_task(self, task: schemas.TaskCreate):
          """
          Creates a new task in our database and a volume claim in the cluster.
  
          Parameters
          ----------
          task: projects.schemas.task.TaskCreate
  
          Returns
          -------
          projects.schemas.task.Task
  
          Raises
          ------
          BadRequest
              When task attributes are invalid.
          """
          has_notebook = task.experiment_notebook or task.deployment_notebook
  
          if task.copy_from and has_notebook:
              raise BadRequest(
                  code="MissingRequiredNotebookOrTaskId",
                  message="Either provide notebooks or a task to copy from",
              )
  
          if not task.tags:
              task.tags = [CATEGORY_DEFAULT]
  
          self.task_category_is_not_none(task)
  
          if not task.category:
              task.category = CATEGORY_DEFAULT
  
          # check if image is a valid docker image
          self.raise_if_invalid_docker_image(task.image)
  
          check_comp_name = (
              self.session.query(models.Task).filter_by(name=task.name).first()
          )
          if check_comp_name:
              raise BadRequest(
                  code="TaskNameExists", message="a task with that name already exists"
              )
  
          # creates a task with specified name,
          # but copies notebooks from a source task
          stored_task = None
          if task.copy_from:
              stored_task = self.session.query(models.Task).get(task.copy_from)
              if stored_task is None:
                  raise BadRequest(
                      code="InvalidTaskId", message="source task does not exist"
                  )
  
              task.image = stored_task.image
              task.commands = stored_task.commands
              task.arguments = stored_task.arguments
              task.parameters = stored_task.parameters
              experiment_notebook_path = stored_task.experiment_notebook_path
              deployment_notebook_path = stored_task.deployment_notebook_path
              task.cpu_limit = stored_task.cpu_limit
              task.cpu_request = stored_task.cpu_request
              task.memory_limit = stored_task.memory_limit
              task.memory_request = stored_task.memory_request
              # Adding the task name if it is a copy.
              task.name = self.generate_name_task(f"{stored_task.name} - Cópia")
  
          else:
              if not isinstance(task.name, str):
                  task.name = self.generate_name_task("Tarefa em branco")
              if task.image is not None:
                  experiment_notebook_path = None
                  deployment_notebook_path = None
                  task.experiment_notebook = None
                  task.deployment_notebook = None
              else:
                  # relative path to the mount_path
                  experiment_notebook_path = "Experiment.ipynb"
                  deployment_notebook_path = "Deployment.ipynb"
                  task.experiment_notebook = EXPERIMENT_NOTEBOOK
                  task.deployment_notebook = DEPLOYMENT_NOTEBOOK
  
          task_id = str(uuid_alpha())
  
          task_dict = task.dict(exclude_unset=True)
          task_dict.pop("copy_from", None)
          task_dict.pop("experiment_notebook", None)
          task_dict.pop("deployment_notebook", None)
          task_dict["uuid"] = task_id
          task_dict["experiment_notebook_path"] = experiment_notebook_path
          task_dict["deployment_notebook_path"] = deployment_notebook_path
  
          # saves task info to the database
          task = models.Task(**task_dict)
  
          self.session.add(task)
          self.session.commit()
          self.session.refresh(task)
  
          all_tasks = self.session.query(models.Task).all()
  
          make_task_creation_job(
              task=task,
              namespace=KF_PIPELINES_NAMESPACE,
              all_tasks=all_tasks,
              copy_from=stored_task,
          )
  
          return schemas.Task.from_orm(task)
  
      def get_task(self, task_id):
          """
          Details a task from our database.
  
          Parameters
          ----------
          task_id : str
  
          Returns
          -------
          projects.schemas.task.Task
  
          Raises
          ------
          NotFound
              When task_id does not exist.
          """
          task = self.session.query(models.Task).get(task_id)
  
          if task is None:
              raise NOT_FOUND
  
          return schemas.Task.from_orm(task)
  
      def update_task(self, task: schemas.TaskUpdate, task_id: str):
          """
          Updates a task in our database/object storage.
  
          Parameters
          ----------
          task: projects.schemas.task.TaskUpdate
          task_id : str
  
          Returns
          -------
          task: projects.schemas.task.Task
  
          Raises
          ------
          NotFound
              When task_id does not exist.
          BadRequest
              When task attributes are invalid.
          """
          self.raise_if_task_does_not_exist(task_id)
  
          stored_task = self.session.query(models.Task).filter_by(name=task.name).first()
          if stored_task and stored_task.uuid != task_id:
              raise BadRequest(
                  code="TaskNameExists", message="a task with that name already exists"
              )
  
          if task.category is not None and task.category not in VALID_CATEGORIES:
              valid_str = ",".join(VALID_CATEGORIES)
              raise BadRequest(
                  code="InvalidCategory",
                  message=f"Invalid category. Choose any of {valid_str}",
              )
  
          stored_task = self.session.query(models.Task).get(task_id)
  
          # If the contents of experiment/deployment notebook were sent,
          # saves them to the notebook server pod
          self.copy_notebooks_to_pod(task, stored_task)
  
          # checks whether task.name has changed
          if stored_task.name != task.name and task.name:
              # update the volume for the task in the notebook server
              self.background_tasks.add_task(
                  update_persistent_volume_claim,
                  name=f"vol-task-{task_id}",
                  mount_path=f"/home/jovyan/tasks/{task.name}",
              )
  
          # update ConfigMap for monitoring tasks
          if (task.parameters and "MONITORING" in stored_task.tags) or (
              "MONITORING" in task.tags if task.tags else False
          ):
              self.background_tasks.add_task(
                  update_task_config_map,
                  task_name=stored_task.name,
                  task_id=task_id,
                  experiment_notebook_path=stored_task.experiment_notebook_path,
              )
  
          update_data = task.dict(exclude_unset=True)
          update_data.pop("experiment_notebook", None)
          update_data.pop("deployment_notebook", None)
          update_data.update({"updated_at": datetime.utcnow()})
  
          self.session.query(models.Task).filter_by(uuid=task_id).update(update_data)
          self.session.commit()
  
          task = self.session.query(models.Task).get(task_id)
  
          return schemas.Task.from_orm(task)
  
      def get_or_create_dataset_task_if_not_exist(self):
          """
          Get or create a dataset  task if the operator has none.
  
          Returns
          -------
          dataset_task.uuid: str
          """
  
          dataset_task = (
              self.session.query(models.Task).filter_by(category="DATASETS").first()
          )
          if dataset_task is None:
              dataset_task_schema = schemas.TaskCreate(
                  name="Upload de arquivo",
                  description="Importa arquivos para utilização em experimentos.",
                  category="DATASETS",
                  tags=["DATASETS"],
                  commands=None,
                  arguments=None,
                  image=TASK_DEFAULT_EXPERIMENT_IMAGE,
                  cpu_limit=TASK_DEFAULT_CPU_LIMIT,
                  cpu_request=TASK_DEFAULT_CPU_REQUEST,
                  memory_limit=TASK_DEFAULT_MEMORY_LIMIT,
                  memory_request=TASK_DEFAULT_MEMORY_REQUEST,
              )
  
              dataset_task = self.create_task(task=dataset_task_schema)
          return dataset_task.uuid
  
      def delete_task(self, task_id: str):
          """
          Delete a task in our database.
          Parameters
          ----------
          task_id : str
          Returns
          -------
          projects.schemas.message.Message
          Raises
          ------
          NotFound
              When task_id does not exist.
          """
          task = self.session.query(models.Task).get(task_id)
  
          if task is None:
              raise NOT_FOUND
  
          if task.operator:
              raise Forbidden(
                  code="TaskProtectedFromDeletion", message="Task related to an operator"
              )
  
          # remove the volume for the task in the notebook server
          self.session.delete(task)
          self.session.commit()
          all_tasks = self.session.query(models.Task).all()
          try:
              make_task_deletion_job(
                  task=task,
                  all_tasks=all_tasks,
                  namespace=KF_PIPELINES_NAMESPACE,
              )
          except Exception as e:
              raise InternalServerError(
                  code="DeletionJobError",
                  message=f"Error while trying to make deletion container job: {e}",
              )
  
          return schemas.Message(message="Task deleted")
  
      def raise_if_invalid_docker_image(self, image):
          """
          Raise an error if a str does not meet the standards for a docker image name.
  
          Example: (username/organization)/name-of-the-image:tag
  
          Parameters
          ----------
          image : str or None
              The image name.
  
          Raises
          ------
          BadRequest
              When a given image is a invalid one.
          """
          pattern = re.compile("[a-z0-9.-]+([/]{1}[a-z0-9.-]+)+([:]{1}[a-z0-9.-]+){0,1}$")
  
          if image and pattern.match(image) is None:
              raise BadRequest(
                  code="InvalidDockerImageName", message="invalid docker image name"
              )
  
      def copy_notebooks_to_pod(self, task, stored_task):
          """
          Copies the notebook contents to the pod (if it was sent on TaskUpdate).
  
          Parameters
          ----------
          task : projects.schemas.task.TaskUpdate
          stored_task : projects.models.task.Task
          """
          if task.experiment_notebook:
              with tempfile.NamedTemporaryFile("w", delete=False) as f:
                  json.dump(task.experiment_notebook, f)
  
              task.experiment_notebook_path = stored_task.experiment_notebook_path
  
              if task.experiment_notebook_path is None:
                  task.experiment_notebook_path = "Experiment.ipynb"
  
              filepath = f.name
              destination_path = f"{stored_task.name}/{task.experiment_notebook_path}"
              copy_file_to_pod(filepath, destination_path)
              os.remove(filepath)
  
          if task.deployment_notebook:
              with tempfile.NamedTemporaryFile("w", delete=False) as f:
                  json.dump(task.deployment_notebook, f)
  
              task.deployment_notebook_path = stored_task.deployment_notebook_path
  
              if task.deployment_notebook_path is None:
                  task.deployment_notebook_path = "Deployment.ipynb"
  
              filepath = f.name
              destination_path = f"{stored_task.name}/{task.deployment_notebook_path}"
              copy_file_to_pod(filepath, destination_path)
              os.remove(filepath)
  
      def make_email_message(self, html_file_content, task_name):
          """
          Build an email body message for a specific task
  
          Parameters
          ----------
          html_file_content: bytes
          task_name : str
  
          Returns
          -------
          template: str
  
          """
  
          # byte to string
          html_string = str(html_file_content, "utf-8")
  
          # body message html string to Jinja2 template
          jinja2_like_template = Template(html_string)
  
          template = jinja2_like_template.render(task_name=task_name)
          return template
  
      def send_emails(self, email_schema, task_id):
          """
          Handles mailing of contents of task
  
          Parameters
          ----------
          task_id : str
          email_schema: projects.schemas.mailing.EmailSchema
  
          Returns
          -------
          message: str
  
          """
  
          task = self.session.query(models.Task).get(task_id)
          if task is None:
              raise NOT_FOUND
  
          # getting file content, which is by the way in base64
          file_as_b64 = get_files_from_task(task.name)
  
          # decoding as byte
          base64_bytes = file_as_b64.encode("ascii")
          file_as_bytes = base64.b64decode(base64_bytes)
  
          # using bytes to build the zipfile
          with tempfile.NamedTemporaryFile(
              "wb", delete=False, dir=os.path.dirname(__file__), suffix=".zip"
          ) as f:
              f.write(file_as_bytes)
  
          message = MessageSchema(
              subject=f"Arquivos da tarefa '{task.name}'",
              recipients=email_schema.dict().get(
                  "emails"
              ),  # List of recipients, as many as you can pass
              body=self.make_email_message(EMAIL_MESSAGE_TEMPLATE, task.name),
              attachments=[f.name],
              subtype="html",
          )
          fm = FastMail(email_schema.conf)
          self.background_tasks.add_task(fm.send_message, message)
  
          # removing file after send email
          os.remove(f.name)
          return {"message": "email has been sent"}
  