import logging
from typing import (
    List,
    Literal,
    Mapping,
    Optional,
    Set,
    Tuple,
    Union,
    overload,
)

from kumoapi.jobs import (
    BatchPredictionOptions,
    BatchPredictionRequest,
    JobStatus,
    MetadataField,
    PredictionOutputConfig,
    TrainingJobRequest,
    TrainingJobResource,
)
from kumoapi.model_plan import ModelPlan

from kumoai import global_state
from kumoai.artifact_export.config import OutputConfig
from kumoai.client.jobs import (
    GeneratePredictionTableJobID,
    TrainingJobAPI,
    TrainingJobID,
)
from kumoai.connector.base import Connector
from kumoai.connector.s3_connector import S3URI
from kumoai.databricks import to_db_table_name
from kumoai.graph import Graph
from kumoai.pquery.prediction_table import PredictionTable, PredictionTableJob
from kumoai.pquery.training_table import TrainingTable, TrainingTableJob
from kumoai.trainer.job import (
    BatchPredictionJob,
    BatchPredictionJobResult,
    TrainingJob,
    TrainingJobResult,
)
from kumoai.trainer.util import (
    build_prediction_output_config,
    validate_output_arguments,
)

logger = logging.getLogger(__name__)


class Trainer:
    r"""A trainer supports creating a Kumo machine learning model on a
    :class:`~kumoai.pquery.PredictiveQuery`. It is primarily oriented around
    two methods: :meth:`~kumoai.trainer.Trainer.fit`, which accepts a
    :class:`~kumoai.graph.Graph` and :class:`~kumoai.pquery.TrainingTable` and
    produces a :class:`~kumoai.trainer.TrainingJobResult`, and
    :meth:`~kumoai.trainer.Trainer.predict`, which accepts a
    :class:`~kumoai.graph.Graph` and :class:`~kumoai.pquery.PredictionTable`
    and produces a :class:`~kumoai.trainer.BatchPredictionJobResult`.

    A :class:`~kumoai.trainer.Trainer` can also be loaded from a training job,
    with :meth:`~kumoai.trainer.Trainer.load`.

    .. code-block:: python

        import kumoai

        # See `Graph` and `PredictiveQuery` documentation:
        graph = kumoai.Graph(...)
        pquery = kumoai.PredictiveQuery(graph=graph, query=...)

        # Create a `Trainer` object, using a suggested model plan given the
        # predictive query:
        model_plan = pquery.suggest_model_plan()
        trainer = kumoai.Trainer(model_plan)

        # Create a training table from the predictive query:
        training_table = pquery.generate_training_table()

        # Fit a model:
        training_job = trainer.fit(
            graph = graph,
            training_table = training_table,
        )

        # Create a prediction table from the predictive query:
        prediction_table = pquery.generate_prediction_table()

        # Generate predictions:
        prediction_job = trainer.predict(
            graph = graph,
            prediction_table = prediction_table,
            # other arguments here...
        )

        # Load a trained query to generate predictions:
        pquery = kumoai.PredictiveQuery.load_from_training_job("trainingjob-...")
        trainer = kumoai.Trainer.load("trainingjob-...")
        prediction_job = trainer.predict(
            pquery.graph,
            pquery.generate_prediction_table(),
        )

    Args:
        model_plan: A model plan that the trainer should follow when fitting a
            Kumo model to a predictive query. This model plan can either be
            generated from a predictive query, with
            :meth:`~kumoai.pquery.PredictiveQuery.suggest_model_plan`, or can
            be manually specified.
    """  # noqa: E501

    def __init__(self, model_plan: ModelPlan) -> None:
        self._model_plan: Optional[ModelPlan] = model_plan

        # Cached from backend:
        self._training_job_id: Optional[TrainingJobID] = None

    @property
    def model_plan(self) -> Optional[ModelPlan]:
        return self._model_plan

    @model_plan.setter
    def model_plan(self, model_plan: ModelPlan) -> None:
        self._model_plan = model_plan

    # Metadata ################################################################

    @property
    def is_trained(self) -> bool:
        r"""Returns ``True`` if this trainer instance has successfully been
        trained (and is therefore ready for prediction); ``False`` otherwise.
        """
        if not self._training_job_id:
            return False
        try:
            api = global_state.client.training_job_api
            res: TrainingJobResource = api.get(self._training_job_id)
        except Exception as e:  # noqa
            logger.exception(
                "Failed to fetch training status for training job with ID %s",
                self._training_job_id, exc_info=e)
            return False
        return res.job_status_report.status == JobStatus.DONE

    # Fit / predict ###########################################################

    @overload
    def fit(
        self,
        graph: Graph,
        train_table: Union[TrainingTable, TrainingTableJob],
    ) -> TrainingJobResult:
        pass

    @overload
    def fit(
        self,
        graph: Graph,
        train_table: Union[TrainingTable, TrainingTableJob],
        *,
        non_blocking: Literal[False],
    ) -> TrainingJobResult:
        pass

    @overload
    def fit(
        self,
        graph: Graph,
        train_table: Union[TrainingTable, TrainingTableJob],
        *,
        non_blocking: Literal[True],
    ) -> TrainingJob:
        pass

    @overload
    def fit(
        self,
        graph: Graph,
        train_table: Union[TrainingTable, TrainingTableJob],
        *,
        non_blocking: bool,
    ) -> Union[TrainingJob, TrainingJobResult]:
        pass

    def fit(
        self,
        graph: Graph,
        train_table: Union[TrainingTable, TrainingTableJob],
        *,
        non_blocking: bool = False,
        custom_tags: Mapping[str, str] = {},
    ) -> Union[TrainingJob, TrainingJobResult]:
        r"""Fits a model to the specified graph and training table, with the
        strategy defined by this :class:`Trainer`'s :obj:`model_plan`.

        Args:
            graph: The :class:`~kumoai.graph.Graph` object that represents the
                tables and relationships that Kumo will learn from.
            train_table: The :class:`~kumoai.pquery.TrainingTable`, or
                in-progress :class:`~kumoai.pquery.TrainingTableJob`, that
                represents the training data produced by a
                :class:`~kumoai.pquery.PredictiveQuery` on :obj:`graph`.
            non_blocking: Whether this operation should return immediately
                after launching the training job, or await completion of the
                training job.
            custom_tags: Additional, customer defined k-v tags to be associated
                with the job to be launched. Job tags are useful for grouping
                and searching jobs.

        Returns:
            Union[TrainingJobResult, TrainingJob]:
                If ``non_blocking=False``, returns a training job object. If
                ``non_blocking=True``, returns a training job future object.
        """
        # TODO(manan, siyang): remove soon:
        job_id = train_table.job_id
        assert job_id is not None

        train_table_job_api = global_state.client.generate_train_table_job_api
        pq_id = train_table_job_api.get(job_id).config.pquery_id
        assert pq_id is not None

        custom_table = None
        if isinstance(train_table, TrainingTable):
            custom_table = train_table._custom_train_table

        # NOTE the backend implementation currently handles sequentialization
        # between a training table future and a training job; that is, if the
        # training table future is still executing, the backend will wait on
        # the job ID completion before executing a training job. This preserves
        # semantics for both futures, ensures that Kumo works as expected if
        # used only via REST API, and allows us to avoid chaining calllbacks
        # in an ugly way here:
        api = global_state.client.training_job_api
        self._training_job_id = api.create(
            TrainingJobRequest(
                dict(custom_tags), pquery_id=pq_id,
                model_plan=self._model_plan,
                graph_snapshot_id=graph.snapshot(non_blocking=non_blocking),
                train_table_job_id=job_id,
                custom_train_table=custom_table,))

        out = TrainingJob(job_id=self._training_job_id)
        if non_blocking:
            return out
        return out.attach()

    def predict(
        self,
        graph: Graph,
        prediction_table: Union[PredictionTable, PredictionTableJob],
        output_types: Optional[Set[str]] = None,
        output_connector: Optional[Connector] = None,
        output_table_name: Optional[Union[str, Tuple]] = None,
        output_metadata_fields: Optional[List[MetadataField]] = None,
        output_config: Optional[OutputConfig] = None,
        training_job_id: Optional[TrainingJobID] = None,
        binary_classification_threshold: Optional[float] = None,
        num_classes_to_return: Optional[int] = None,
        num_workers: int = 1,
        non_blocking: bool = False,
        custom_tags: Mapping[str, str] = {},
    ) -> Union[BatchPredictionJob, BatchPredictionJobResult]:
        """Using the trained model specified by :obj:`training_job_id` (or
        the model corresponding to the last invocation of
        :meth:`~kumoai.trainer.Trainer.fit`, if not present), predicts the
        future values of the entities in :obj:`prediction_table` leveraging
        current information from :obj:`graph`.

        Args:
            graph: The :class:`~kumoai.graph.Graph` object that represents the
                tables and relationships that Kumo will use to make
                predictions.
            prediction_table: The :class:`~kumoai.pquery.PredictionTable`, or
                in-progress :class:`~kumoai.pquery.PredictionTableJob`, that
                represents the prediction data produced by a
                :class:`~kumoai.pquery.PredictiveQuery` on :obj:`graph`. This
                table may also be custom-generated.
            output_config: Output configuration defining properties of the
                generated batch prediction outputs. Typically aspects like
                destination for the output, any additional columns needed,
                whether to override or append to an existing table etc. is
                supplied via this.
            output_types: The types of outputs that should be produced by
                the prediction job. Can include either ``'predictions'``,
                ``'embeddings'``, or both.
            output_connector: The output data source that Kumo should write
                batch predictions to, if it is None, produce local download
                output only.
            output_table_name: The name of the table in the output data source
                that Kumo should write batch predictions to. In the case of
                a Databricks connector, this should be a tuple of two strings:
                the schema name and the output prediction table name.
            output_metadata_fields: Any additional metadata fields to include
                as new columns in the produced ``'predictions'`` output.
                Currently, allowed options are ``JOB_TIMESTAMP`` and
                ``ANCHOR_TIMESTAMP``.
            training_job_id: The job ID of the training job whose model will be
                used for making predictions.
            binary_classification_threshold: If this model corresponds to
                a binary classification task, the threshold for which higher
                values correspond to ``1``, and lower values correspond to
                ``0``.
            num_classes_to_return: If this model corresponds to a ranking task,
                the number of classes to return in the prediction output.
            num_workers: Number of workers to use when generating batch
                predictions. When set to a value greater than 1, the prediction
                table is partitioned into smaller parts and processed in
                parallel. The default is 1, which implies sequential inference
                over the prediction table.
            non_blocking: Whether this operation should return immediately
                after launching the batch prediction job, or await
                completion of the batch prediction job.
            custom_tags: Additional, customer defined k-v tags to be associated
                with the job to be launched. Job tags are useful for grouping
                and searching jobs.

        Returns:
            Union[BatchPredictionJob, BatchPredictionJobResult]:
                If ``non_blocking=False``, returns a training job object. If
                ``non_blocking=True``, returns a training job future object.
        """
        if (output_types is not None or output_connector is not None
                or output_table_name is not None
                or output_metadata_fields is not None):
            raise ValueError(
                'Specifying output_types, output_connector, '
                'output_metadata_fields '
                'and output_table_name as direct inputs to predict() is '
                'deprecated. Please use output_config to specify these '
                'parameters.')
        assert output_config is not None
        output_table_name = to_db_table_name(output_config.output_table_name)
        validate_output_arguments(
            output_config.output_types,
            output_config.output_connector,
            output_table_name,
        )

        # Create outputs:
        outputs: List[PredictionOutputConfig] = []
        for output_type in output_config.output_types:
            if output_config.output_connector is None:
                # Predictions are generated to the Kumo dataplane, and can
                # only be exported via the UI for now:
                pass
            else:
                outputs.append(
                    build_prediction_output_config(
                        output_type,
                        output_config.output_connector,
                        output_table_name,
                        output_config.output_metadata_fields,
                        output_config,
                    ))

        training_job_id = training_job_id or self._training_job_id
        if training_job_id is None:
            raise ValueError(
                "Cannot run batch prediction without a specified or saved "
                "training job ID. Please either call `fit(...)` or pass a "
                "job ID of a completed training job to proceed.")

        pred_table_job_id: Optional[GeneratePredictionTableJobID] = \
            prediction_table.job_id
        pred_table_data_path = None
        if pred_table_job_id is None:
            assert isinstance(prediction_table, PredictionTable)
            if isinstance(prediction_table.table_data_uri, S3URI):
                pred_table_data_path = prediction_table.table_data_uri.uri
            else:
                pred_table_data_path = prediction_table.table_data_uri

        api = global_state.client.batch_prediction_job_api

        job_id, response = api.maybe_create(
            BatchPredictionRequest(
                dict(custom_tags),
                model_training_job_id=training_job_id,
                predict_options=BatchPredictionOptions(
                    binary_classification_threshold=(
                        binary_classification_threshold),
                    num_classes_to_return=num_classes_to_return,
                    num_workers=num_workers,
                ),
                outputs=outputs,
                graph_snapshot_id=graph.snapshot(non_blocking=non_blocking),
                pred_table_job_id=pred_table_job_id,
                pred_table_path=pred_table_data_path,
            ))

        message = response.message()
        if not response.ok:
            raise RuntimeError(f"Prediction failed. {message}")
        elif not response.empty():
            logger.warning("Prediction produced the following warnings: %s",
                           message)
        assert job_id is not None

        self._batch_prediction_job_id = job_id
        out = BatchPredictionJob(job_id=self._batch_prediction_job_id)
        if non_blocking:
            return out
        return out.attach()

    # Persistence #############################################################

    @classmethod
    def _load_from_job(cls, job: TrainingJobResource) -> 'Trainer':
        trainer = cls(job.config.model_plan)
        trainer._training_job_id = job.job_id
        return trainer

    @classmethod
    def load(cls, job_id: TrainingJobID) -> 'Trainer':
        r"""Creates a :class:`~kumoai.trainer.Trainer` instance from a training
        job ID.
        """
        api: TrainingJobAPI = global_state.client.training_job_api
        job = api.get(job_id)
        return cls._load_from_job(job)

    # TODO(siyang): load trainer by searching training job via tags.
    @classmethod
    def load_from_tags(cls, tags: Mapping[str, str]) -> 'Trainer':
        r"""Creates a :class:`~kumoai.trainer.Trainer` instance from a set of
        job tags. If multiple jobs match the list of tags, only one will be
        selected.
        """
        api = global_state.client.training_job_api
        jobs = api.list(limit=1, additional_tags=tags)
        if not jobs:
            raise RuntimeError(f'No successful training job found for {tags}')
        assert len(jobs) == 1
        return cls._load_from_job(jobs[0])
