content: #  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
  #
  #  Licensed under the Apache License, Version 2.0 (the "License");
  #  you may not use this file except in compliance with the License.
  #  You may obtain a copy of the License at:
  #
  #       https://www.apache.org/licenses/LICENSE-2.0
  #
  #  Unless required by applicable law or agreed to in writing, software
  #  distributed under the License is distributed on an "AS IS" BASIS,
  #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
  #  or implied. See the License for the specific language governing
  #  permissions and limitations under the License.
  """Implementation of the Prodigy annotation integration."""
  
  import json
  from typing import Any, List, Optional, Tuple, cast
  
  import prodigy
  from peewee import Database as PeeweeDatabase
  from prodigy.components.db import (
      Database as ProdigyDatabase,
  )
  from prodigy.components.db import connect
  from prodigy.errors import ProdigyError
  
  from zenml.annotators.base_annotator import BaseAnnotator
  from zenml.integrations.prodigy.flavors.prodigy_annotator_flavor import (
      ProdigyAnnotatorConfig,
  )
  from zenml.logger import get_logger
  from zenml.stack.authentication_mixin import AuthenticationMixin
  
  logger = get_logger(__name__)
  
  DEFAULT_LOCAL_INSTANCE_HOST = "localhost"
  DEFAULT_LOCAL_PRODIGY_PORT = 8080
  
  
  class ProdigyAnnotator(BaseAnnotator, AuthenticationMixin):
      """Class to interact with the Prodigy annotation interface."""
  
      @property
      def config(self) -> ProdigyAnnotatorConfig:
          """Returns the `ProdigyAnnotatorConfig` config.
  
          Returns:
              The configuration.
          """
          return cast(ProdigyAnnotatorConfig, self._config)
  
      def get_url(self) -> str:
          """Gets the top-level URL of the annotation interface.
  
          Returns:
              The URL of the annotation interface.
          """
          instance_url = DEFAULT_LOCAL_INSTANCE_HOST
          port = DEFAULT_LOCAL_PRODIGY_PORT
          if self.config.custom_config_path:
              with open(self.config.custom_config_path, "r") as f:
                  config = json.load(f)
              instance_url = config.get("instance_url", instance_url)
              port = config.get("port", port)
          return f"http://{instance_url}:{port}"
  
      def get_url_for_dataset(self, dataset_name: str) -> str:
          """Gets the URL of the annotation interface for the given dataset.
  
          Prodigy does not support dataset-specific URLs, so this method returns
          the top-level URL since that's what will be served for the user.
  
          Args:
              dataset_name: The name of the dataset. (Unuse)
  
          Returns:
              The URL of the annotation interface.
          """
          return self.get_url()
  
      def get_datasets(self) -> List[Any]:
          """Gets the datasets currently available for annotation.
  
          Returns:
              A list of datasets (str).
          """
          datasets = self._get_db().datasets
          return cast(List[Any], datasets)
  
      def get_dataset_names(self) -> List[str]:
          """Gets the names of the datasets.
  
          Returns:
              A list of dataset names.
          """
          return self.get_datasets()
  
      def get_dataset_stats(self, dataset_name: str) -> Tuple[int, int]:
          """Gets the statistics of the given dataset.
  
          Args:
              dataset_name: The name of the dataset.
  
          Returns:
              A tuple containing (labeled_task_count, unlabeled_task_count) for
                  the dataset.
  
          Raises:
              IndexError: If the dataset does not exist.
          """
          db = self._get_db()
          try:
              labeled_data_count = db.count_dataset(name=dataset_name)
          except ValueError as e:
              raise IndexError(
                  f"Dataset {dataset_name} does not exist. Please use `zenml "
                  f"annotator dataset list` to list the available datasets."
              ) from e
          return (labeled_data_count, 0)
  
      def launch(self, **kwargs: Any) -> None:
          """Launches the annotation interface.
  
          This method extracts the 'command' and additional config
              parameters from kwargs.
  
          Args:
              **kwargs: Should include:
                  - command: The full recipe command without "prodigy".
                  - Any additional config parameters to overwrite the
                      project-specific, global, and recipe config.
  
          Raises:
              ValueError: If the 'command' keyword argument is not provided.
          """
          command = kwargs.get("command")
          if not command:
              raise ValueError(
                  "The 'command' keyword argument is required for launching Prodigy."
              )
  
          # Remove 'command' from kwargs to pass the rest as config parameters
          config = {
              key: value for key, value in kwargs.items() if key != "command"
          }
          prodigy.serve(command=command, **config)
  
      def _get_db(
          self,
          custom_database: PeeweeDatabase = None,
          display_id: Optional[str] = None,
          display_name: Optional[str] = None,
      ) -> ProdigyDatabase:
          """Gets Prodigy database / client.
  
          Args:
              custom_database: Custom database to use.
              display_id: The display id of the database.
              display_name: The display name of the database.
  
          Returns:
              Prodigy database client.
          """
          db_kwargs = {}
          if custom_database:
              db_kwargs["db"] = custom_database
          if display_id:
              db_kwargs["display_id"] = display_id
          if display_name:
              db_kwargs["display_name"] = display_name
  
          # database is passed in without the keyword argument
          if custom_database:
              return connect(custom_database, **db_kwargs)
          return connect(**db_kwargs)
  
      def add_dataset(self, **kwargs: Any) -> Any:
          """Registers a dataset for annotation.
  
          Args:
              **kwargs: Additional keyword arguments to pass to the Prodigy client.
  
          Returns:
              A Prodigy list representing the dataset.
  
          Raises:
              ValueError: if 'dataset_name' and 'label_config' aren't provided.
          """
          db = self._get_db()
          dataset_kwargs = {"dataset_name": kwargs.get("dataset_name")}
          if not dataset_kwargs["dataset_name"]:
              raise ValueError("`dataset_name` keyword argument is required.")
  
          if kwargs.get("dataset_meta"):
              dataset_kwargs["dataset_meta"] = kwargs.get("dataset_meta")
          return db.add_dataset(**dataset_kwargs)
  
      def delete_dataset(self, **kwargs: Any) -> None:
          """Deletes a dataset from the annotation interface.
  
          Args:
              **kwargs: Additional keyword arguments to pass to the Prodigy
                  client.
  
          Raises:
              ValueError: If the dataset name is not provided or if the dataset
                  does not exist.
          """
          db = self._get_db()
          if not (dataset_name := kwargs.get("dataset_name")):
              raise ValueError("`dataset_name` keyword argument is required.")
          try:
              db.drop_dataset(name=dataset_name)
          except ProdigyError as e:
              # see https://support.prodi.gy/t/how-to-import-datasetdoesnotexist-error/7205
              if type(e).__name__ == "DatasetNotFound":
                  raise ValueError(
                      f"Dataset name '{dataset_name}' does not exist."
                  ) from e
  
      def get_dataset(self, **kwargs: Any) -> Any:
          """Gets the dataset metadata for the given name.
  
          If you would like the labeled data, use `get_labeled_data` instead.
  
          Args:
              **kwargs: Additional keyword arguments to pass to the Prodigy client.
  
          Returns:
              The metadata associated with a Prodigy dataset
  
          Raises:
              ValueError: If the dataset name is not provided or if the dataset
                  does not exist.
          """
          db = self._get_db()
          if dataset_name := kwargs.get("dataset_name"):
              try:
                  return db.get_meta(name=dataset_name)
              except Exception as e:
                  raise ValueError(
                      f"Dataset name '{dataset_name}' does not exist."
                  ) from e
  
      def get_labeled_data(self, **kwargs: Any) -> Any:
          """Gets the labeled data for the given dataset.
  
          Args:
              **kwargs: Additional keyword arguments to pass to the Prodigy client.
  
          Returns:
              A list of all examples in the dataset serialized to the
                  Prodigy Task format.
  
          Raises:
              ValueError: If the dataset name is not provided or if the dataset
                  does not exist.
          """
          if dataset_name := kwargs.get("dataset_name"):
              return self._get_db().get_dataset_examples(dataset_name)
          else:
              raise ValueError("`dataset_name` keyword argument is required.")
  
      def get_unlabeled_data(self, **kwargs: str) -> Any:
          """Gets the unlabeled data for the given dataset.
  
          Args:
              **kwargs: Additional keyword arguments to pass to the Prodigy client.
  
          Raises:
              NotImplementedError: Prodigy doesn't allow fetching unlabeled data.
          """
          raise NotImplementedError(
              "Prodigy doesn't allow fetching unlabeled data."
          )
  