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

from kumoapi.jobs import (
    GeneratePredictionTableRequest,
    GenerateTrainTableRequest,
)
from kumoapi.model_plan import (
    InferredType,
    PredictionTableGenerationPlan,
    RunMode,
    SuggestModelPlanRequest,
    TrainingTableGenerationPlan,
)
from kumoapi.pquery import PQueryResource
from kumoapi.task import TaskType
from kumoapi.train import TrainingTableSpec
from typing_extensions import Self

from kumoai import global_state
from kumoai.client.jobs import (
    GeneratePredictionTableJobID,
    GenerateTrainTableJobID,
    TrainingJobAPI,
)
from kumoai.graph import Graph
from kumoai.pquery.prediction_table import PredictionTable, PredictionTableJob
from kumoai.pquery.training_table import TrainingTable, TrainingTableJob
from kumoai.trainer import (
    BaselineTrainer,
    ModelPlan,
    Trainer,
    TrainingJob,
    TrainingJobResult,
)
from kumoai.trainer.job import BaselineJob, BaselineJobResult

logger = logging.getLogger(__name__)

PredictiveQueryID = str


class PredictiveQuery:
    r"""The Kumo predictive query is a declarative syntax for describing a
    machine learning task. Predictive queries are written using the predictive
    query language (PQL), a concise SQL-like syntax that allows you to define a
    model for a new business problem.

    A predictive query object can be created from a
    :class:`~kumoai.graph.Graph` and a query string. For information on the
    construction of a query string, please visit the Kumo
    `documentation <https://docs.kumo.ai/docs/pquery-structure/>`__.

    .. code-block:: python

        import kumoai

        # See `Graph` documentation for more information:
        graph = kumoai.Graph(...)

        # Create a predictive query representing a machine learning problem
        # over this Graph:
        pquery = kumoai.PredictiveQuery(
            graph=graph,
            query=(
                "PREDICT MAX(transaction.Quantity, 0, 30) "
                "FOR EACH customer.CustomerID"
            ),
        )

        # Validate the predictive query configuration, for syntax and
        # correctness:
        pquery.validate(verbose=True)

        # Get the machine learning task type corresponding to this predictive
        # query (e.g. binary classification, regression, link prediction, etc.)
        print(pquery.get_task_type())

        # Suggest a training table generation plan and use it to generate a
        # training table from this query, to be used in `Trainer.fit`:
        training_table_plan = pquery.suggest_training_table_plan()
        training_table = pquery.generate_training_table(training_table_plan)

        # Suggest a prediction table generation plan and use it to generate a
        # prediction table from this query, to be used in `Trainer.predict`:
        pred_table_plan = pquery.suggest_prediction_table_plan()
        pred_table = pquery.generate_prediction_table(pred_table_plan)

    Args:
        graph: The :class:`~kumoai.graph.Graph` object which the predictive
            query is defined over.
        query: A string representation of the predictive query.
    """
    def __init__(
        self,
        graph: Graph,
        query: str,
    ) -> None:
        self.graph = graph
        self.query = query

        # A predictive query owns a trainer object, which is used internally
        # to support `fit` and `predict` directly on this object. A user can
        # also inspect the training table, prediction table, and trainer
        # objects, but cannot set them; any advanced configuration must be
        # done directly via `Trainer`:
        self._train_table: Optional[Union[TrainingTable,
                                          TrainingTableJob]] = None
        self._prediction_table: Optional[Union[PredictionTable,
                                               PredictionTableJob]] = None

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

    @property
    def id(self) -> str:
        r"""Returns the unique ID for this predictive query, determined from
        its schema and the schema of its associated graph. Two queries that
        differ either in their syntax or in their graph will have different
        ids.
        """
        return self.save()

    @property
    def train_table(self) -> Union[TrainingTable, TrainingTableJob]:
        r"""Returns the training table that was last generated by this
        predictive query. If the predictive query has not yet generated a
        training table, raises a :class:`ValueError`.

        Note that the training table may be of type
        :class:`~kumoai.pquery.TrainingTable` or
        :class:`~kumoai.pquery.TrainingTableJob`, depending on whether the
        training table was generated with or without waiting for its
        completion, respectively.
        """
        if not self._train_table:
            raise ValueError(
                "This predictive query has not yet generated a training "
                "table. Please call `generate_training_table` to generate "
                "a training table before proceeding.")

        return self._train_table

    @property
    def prediction_table(self) -> Union[PredictionTable, PredictionTableJob]:
        r"""Returns the prediction table that was last generated by this
        predictive query. If the predictive query has not yet generated a
        prediction table, raises a :class:`ValueError`.

        Note that the prediction table may be of type
        :class:`~kumoai.pquery.PredictionTable` or
        :class:`~kumoai.pquery.PredictionTableJob`, depending on whether the
        prediction table was generated with or without waiting for its
        completion, respectively.
        """
        if not self._prediction_table:
            raise ValueError(
                "This predictive query has not yet generated a prediction "
                "table. Please call `generate_prediction_table` to generate a "
                "prediction table before proceeding.")

        return self._prediction_table

    def get_task_type(self) -> TaskType:
        r"""Returns the task type of this predictive query. The task type of
        the query corresponds to the machine learning problem that this query
        translates to in the Kumo platform; for more information about possible
        task types, please visit the Kumo `documentation
        <https://docs.kumo.ai/docs/task-types/>`__.
        """
        try:
            self.validate(verbose=False)
        except ValueError as e:
            raise ValueError(
                f"Predictive query {self.query} is improperly configured, so "
                f"a task type cannot be obtained. Please ensure your query "
                f"has a valid configuration before proceeding. You can use "
                f"the `validate` method to verify the validity of your query."
            ) from e
        task_type, _ = global_state.client.pquery_api.infer_task_type(
            pquery_string=self.query, graph_id=self.graph.id)
        return task_type

    def validate(self, verbose: bool = True) -> Self:
        r"""Validates the syntax of this predictive query, ensuring that
        the query is formulated correctly in Kumo's Predictive Query Language
        and that the query makes semantic sense (defines a suitable predictive
        problem) on this :class:`~kumoai.graph.Graph`.

        Args:
            verbose: Whether to log non-error output of this validation.

        Raises:
            ValueError:
                if validation fails.

        Example:
            >>> import kumoai
            >>> query = kumoai.PredictiveQuery(...)  # doctest: +SKIP
            >>> query.validate()  # doctest: +SKIP
            ValidationResponse(warnings=[], errors=[])
        """
        self.graph.save()  # Need a valid graph ID; also validates graph.

        resp = global_state.client.pquery_api.validate(
            self._to_api_pquery_resource())

        if not resp.ok:
            raise ValueError(resp.error_message())
        if verbose:
            if resp.empty():
                logger.info("Query %s is configured correctly.", self.query)
            else:
                logger.warning(resp.message())

        return self

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

    def _to_api_pquery_resource(
        self,
        name: Optional[str] = None,
    ) -> PQueryResource:
        return PQueryResource(
            name=name,
            query_string=self.query,
            graph=self.graph._to_api_graph_definition(),
            desc="",
        )

    def save(self, name: Optional[str] = None) -> PredictiveQueryID:
        r"""Saves a predictive query to Kumo, returning a unique ID for this
        query. If a name is provided, saves it as a named, re-usable template.

        Args:
            name: Optional name for the template. If provided, saves the
                query as a named template. If the name already exists,
                that template will be overwritten.

        Example:
            >>> import kumoai
            >>> query = kumoai.PredictiveQuery(...)  # doctest: +SKIP
            >>> query.save()  # doctest: +SKIP
            pquery-xxx
            >>> query.save("my_template")  # doctest: +SKIP
            my_template
        """
        try:
            self.validate(verbose=False)
        except ValueError as e:
            raise ValueError(
                f"Predictive query {self.query} is improperly configured, so "
                f"it cannot be saved. Please ensure your query "
                f"has a valid configuration before proceeding. You can use "
                f"the `validate` method to verify the validity of your query."
            ) from e

        if name is not None:
            template_resource = global_state.client.pquery_api.get_if_exists(
                name)
            if template_resource is not None:
                template_string = template_resource.query_string
                logger.warning(
                    ("Predictive query template %s already exists, with "
                     "query string %s. This template will be overridden with "
                     "configuration %s."), name, template_string, self.query)

        self.graph.save()
        return global_state.client.pquery_api.create(
            pquery=self._to_api_pquery_resource(name))

    @classmethod
    def load(cls, pq_id_or_template: str) -> 'PredictiveQuery':
        r"""Loads a predictive query from either a predictive query ID or a
        named template. Returns a :class:`~kumoai.pquery.PredictiveQuery`
        object that contains the loaded query along with its associated graph,
        tables, etc.
        """
        api = global_state.client.pquery_api
        res = api.get_if_exists(pq_id_or_template)
        if not res:
            raise ValueError(
                f"Predictive query {pq_id_or_template} was not found.")
        return cls(
            graph=Graph._from_api_graph_definition(res.graph),
            query=res.query_string,
        )

    @classmethod
    def load_from_training_job(cls, training_job_id: str) -> 'PredictiveQuery':
        r"""Loads a predictive query from a training job, regardless of the
        training job's status. Returns a
        :class:`~kumoai.pquery.PredictiveQuery` object that contains the loaded
        query along with its associated graph, tables, etc.
        """
        train_api: TrainingJobAPI = global_state.client.training_job_api
        job = train_api.get(training_job_id)
        id_or_name = job.config.pquery_id
        return PredictiveQuery.load(pq_id_or_template=id_or_name)

    # Training & Prediction Table Generation ##################################

    @overload
    def generate_training_table(
        self,
        plan: Optional[TrainingTableGenerationPlan] = None,
    ) -> TrainingTable:
        pass

    @overload
    def generate_training_table(
        self,
        plan: Optional[TrainingTableGenerationPlan] = None,
        *,
        non_blocking: Literal[False],
    ) -> TrainingTable:
        pass

    @overload
    def generate_training_table(
        self,
        plan: Optional[TrainingTableGenerationPlan] = None,
        *,
        non_blocking: Literal[True],
    ) -> TrainingTableJob:
        pass

    @overload
    def generate_training_table(
        self,
        plan: Optional[TrainingTableGenerationPlan] = None,
        *,
        non_blocking: bool,
    ) -> Union[TrainingTable, TrainingTableJob]:
        pass

    def generate_training_table(
        self,
        plan: Optional[TrainingTableGenerationPlan] = None,
        *,
        non_blocking: bool = False,
        custom_tags: Mapping[str, str] = {},
    ) -> Union[TrainingTable, TrainingTableJob]:
        r"""Generates a training table from the specified :attr:`query`
        string.

        Args:
            plan: A specification of the parameters for training table
                generation. If not provided, will use an intelligently
                generated default plan based on the query and graph. This plan
                is equivalent to the plan inferred with
                ``suggest_training_table_plan(run_mode=RunMode.NORMAL)``.
            non_blocking: Whether this operation should return immediately
                after launching the training table generation job, or await
                completion of the generated training table.
            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[TrainingTable, TrainingTableJob]:
                If ``non_blocking=False``, returns a training table object. If
                ``non_blocking=True``, returns a training table future object.
        """
        pq_id = self.save()

        # TODO(manan): improve this...
        if not plan:
            plan = self.suggest_training_table_plan()

        train_table_job_api = global_state.client.generate_train_table_job_api
        job_id: GenerateTrainTableJobID = train_table_job_api.create(
            GenerateTrainTableRequest(
                dict(custom_tags), pq_id, plan,
                graph_snapshot_id=self.graph.snapshot(
                    non_blocking=non_blocking)))

        self._train_table = TrainingTableJob(job_id=job_id)
        if non_blocking:
            return self._train_table
        self._train_table = self._train_table.attach()
        return self._train_table

    @overload
    def generate_prediction_table(
        self,
        plan: Optional[PredictionTableGenerationPlan] = None,
    ) -> PredictionTable:
        pass

    @overload
    def generate_prediction_table(
        self,
        plan: Optional[PredictionTableGenerationPlan] = None,
        *,
        non_blocking: Literal[False],
    ) -> PredictionTable:
        pass

    @overload
    def generate_prediction_table(
        self,
        plan: Optional[PredictionTableGenerationPlan] = None,
        *,
        non_blocking: Literal[True],
    ) -> PredictionTableJob:
        pass

    @overload
    def generate_prediction_table(
        self,
        plan: Optional[PredictionTableGenerationPlan] = None,
        *,
        non_blocking: bool,
    ) -> Union[PredictionTable, PredictionTableJob]:
        pass

    def generate_prediction_table(
        self,
        plan: Optional[PredictionTableGenerationPlan] = None,
        *,
        non_blocking: bool = False,
        custom_tags: Mapping[str, str] = {},
    ) -> Union[PredictionTable, PredictionTableJob]:
        r"""Generates a prediction table from the predictive query
        :attr:`query` string.

        Args:
            plan: A specification of the parameters for prediction table
                generation. If not provided, will use an intelligently
                generated default plan based on the query and graph. This plan
                is equivalent to the plan inferred with
                ``suggest_prediction_table_plan(run_mode=RunMode.NORMAL)``.
            non_blocking: Whether this operation should return immediately
                after launching the prediction table generation job, or await
                completion of the generated prediction table.
            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[PredictionTable, PredictionTableJob]:
                If ``non_blocking=False``, returns a prediction table object.
                If ``non_blocking=True``, returns a prediction table future
                object.
        """
        pq_id = self.save()

        if not plan:
            plan = self.suggest_prediction_table_plan()

        bp_table_api = global_state.client.generate_prediction_table_job_api
        job_id: GeneratePredictionTableJobID = bp_table_api.create(
            GeneratePredictionTableRequest(
                dict(custom_tags), pq_id, plan,
                graph_snapshot_id=self.graph.snapshot(
                    non_blocking=non_blocking)))

        self._prediction_table = PredictionTableJob(job_id=job_id)
        if non_blocking:
            return self._prediction_table
        self._prediction_table = self._prediction_table.result()
        return self._prediction_table

    # Training & Prediction ###################################################

    def suggest_training_table_plan(
        self,
        run_mode: RunMode = RunMode.FAST,
    ) -> TrainingTableGenerationPlan:
        r"""Suggests a training table generation plan given the predictive
        query and graph. This training table generation plan can be used to
        alter the approach Kumo uses to generate the training table for your
        predictive query.

        Args:
            run_mode: A representation of how quickly you would like your
                predictive query to complete. Faster run modes correspond to
                lower training times, at the cost of potentially lower
                performance.
        """
        self.graph.save()
        req = SuggestModelPlanRequest(
            query_string=self.query,
            graph_id=self.graph.id,
            run_mode=run_mode,
        )
        return global_state.client.pquery_api.suggest_training_table_plan(req)

    def suggest_prediction_table_plan(self, ) -> PredictionTableGenerationPlan:
        r"""Suggests a prediction table generation plan given the predictive
        query and graph. This prediction table generation plan can be used to
        alter the approach Kumo uses to generate the prediction table for your
        predictive query.
        """
        return PredictionTableGenerationPlan(anchor_time=InferredType.VALUE)

    def suggest_model_plan(
        self,
        run_mode: RunMode = RunMode.FAST,
        train_table_spec: Optional[TrainingTableSpec] = None,
    ) -> ModelPlan:
        r"""Suggests a modeling plan given the predictive query and graph. This
        model plan can be used to alter the approach Kumo uses to train your
        machine learning model.

        Args:
            run_mode: A representation of how quickly you would like your
                predictive query to complete. Faster run modes correspond to
                lower training times, at the cost of potentially lower
                performance.
            train_table_spec: Needed if the original train table has been
                modified by adding a weight column.
        """
        self.graph.save()
        req = SuggestModelPlanRequest(
            query_string=self.query,
            graph_id=self.graph.id,
            run_mode=run_mode,
            train_table_spec=train_table_spec,
        )
        return global_state.client.pquery_api.suggest_model_plan(req)

    @overload
    def fit(
        self,
        training_table_plan: Optional[TrainingTableGenerationPlan] = None,
        model_plan: Optional[ModelPlan] = None,
    ) -> Tuple[Trainer, TrainingJobResult]:
        pass

    @overload
    def fit(
        self,
        training_table_plan: Optional[TrainingTableGenerationPlan] = None,
        model_plan: Optional[ModelPlan] = None,
        *,
        non_blocking: Literal[False],
    ) -> Tuple[Trainer, TrainingJobResult]:
        pass

    @overload
    def fit(
        self,
        training_table_plan: Optional[TrainingTableGenerationPlan] = None,
        model_plan: Optional[ModelPlan] = None,
        *,
        non_blocking: Literal[True],
    ) -> Tuple[Trainer, TrainingJob]:
        pass

    @overload
    def fit(
        self,
        training_table_plan: Optional[TrainingTableGenerationPlan] = None,
        model_plan: Optional[ModelPlan] = None,
        *,
        non_blocking: bool,
    ) -> Tuple[Trainer, Union[TrainingJobResult, TrainingJob]]:
        pass

    def fit(
        self,
        training_table_plan: Optional[TrainingTableGenerationPlan] = None,
        model_plan: Optional[ModelPlan] = None,
        *,
        non_blocking: bool = False,
    ) -> Tuple[Trainer, Union[TrainingJobResult, TrainingJob]]:
        r"""Trains a Kumo model on this predictive query, given optional
        additional specifications of the training table generation plan and
        the model plan.

        Args:
            training_table_plan: A specification of the parameters for training
                table generation. If not provided, will use an intelligently
                generated default plan based on the query and graph. This plan
                is equivalent to the plan inferred with
                ``suggest_training_table_plan(run_mode=RunMode.NORMAL)``.
            model_plan: A specification of the parameters for model training.
                If not provided, will use an intelligently generated default
                plan based on the query and graph. This plan
                is equivalent to the plan inferred with
                ``suggest_model_plan(run_mode=RunMode.NORMAL)``.
            non_blocking: Whether this operation should return immediately
                after launching the training job, or await completion of the
                training job.

        Returns:
            Tuple[Trainer, Union[TrainingJobResult, TrainingJob]]:
                A tuple with two elements. The first element is the trainer
                object used to launch the training job. The second element
                is either a training job object (if ``non_blocking=True``)
                or a training job future object (if ``non_blocking=False``).
        """
        # If we have already generated the training table, use it with Trainer:
        if self._train_table is None:
            # Nonblocking generate:
            self._train_table = self.generate_training_table(
                training_table_plan, non_blocking=True)

        # TODO(manan): what if `self._train_table` represents a failed job?
        model_plan = model_plan or self.suggest_model_plan()
        trainer = Trainer(model_plan)
        return (trainer,
                trainer.fit(self.graph, self.train_table,
                            non_blocking=non_blocking))

    def generate_baseline(
        self,
        metrics: List[str],
        train_table: Union[TrainingTable, TrainingTableJob],
        *,
        non_blocking: bool = False,
    ) -> Union[BaselineJob, BaselineJobResult]:
        r"""Runs a baseline model on this predictive query, given metrics and
        optional additional specifications of the training table generation
        plan.

        Args:
            metrics (List[str]): A list to metrics that baseline model will be
                evaluated on.
            train_table (Union[TrainingTable, TrainingTableJob]): 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 (bool): Whether this operation should
                return immediately after launching the baseline job, or await
                completion of the baseline job. Defaults to False.

        Returns:
            Union[BaselineJob, BaselineJobResult]:  either a baseline job
                object (if ``non_blocking=True``) or a baseline job future
                object (if ``non_blocking=False``).
        """ # noqa
        baseline_trainer = BaselineTrainer(metrics)
        return baseline_trainer.run(self.graph, train_table,
                                    non_blocking=non_blocking)
