from kedro.framework.project import pipelines
from kedro.io import DataCatalog
from kedro.runner import SequentialRunner
from celery import shared_task, Task
from .backends import init_backend
from fastapi.encoders import jsonable_encoder

class KedroGraphqlTask(Task):
    _db = None

    @property
    def db(self):
        if self._db is None:
            self._db = init_backend()
        return self._db

    def before_start(self, task_id, args, kwargs):
        """Handler called before the task starts.

        .. versionadded:: 5.2

        Arguments:
            task_id (str): Unique id of the task to execute.
            args (Tuple): Original arguments for the task to execute.
            kwargs (Dict): Original keyword arguments for the task to execute.

        Returns:
            None: The return value of this handler is ignored.
        """
        self.db.update(task_id = task_id, values = {"status": "STARTED"})

    def on_success(self, retval, task_id, args, kwargs):
        """Success handler.

        Run by the worker if the task executes successfully.

        Arguments:
            retval (Any): The return value of the task.
            task_id (str): Unique id of the executed task.
            args (Tuple): Original arguments for the executed task.
            kwargs (Dict): Original keyword arguments for the executed task.

        Returns:
            None: The return value of this handler is ignored.
        """
        self.db.update(task_id = task_id, values = {"status": "SUCCESS"})


    def on_retry(self, exc, task_id, args, kwargs, einfo):
        """Retry handler.

        This is run by the worker when the task is to be retried.

        Arguments:
            exc (Exception): The exception sent to :meth:`retry`.
            task_id (str): Unique id of the retried task.
            args (Tuple): Original arguments for the retried task.
            kwargs (Dict): Original keyword arguments for the retried task.
            einfo (~billiard.einfo.ExceptionInfo): Exception information.

        Returns:
            None: The return value of this handler is ignored.
        """
        self.db.update(task_id = task_id, values = {"status": "RETRY", "task_exception": str(exc), "task_einfo": str(einfo)})
   
    def on_failure(self, exc, task_id, args, kwargs, einfo):
        """Error handler.

        This is run by the worker when the task fails.

        Arguments:
            exc (Exception): The exception raised by the task.
            task_id (str): Unique id of the failed task.
            args (Tuple): Original arguments for the task that failed.
            kwargs (Dict): Original keyword arguments for the task that failed.
            einfo (~billiard.einfo.ExceptionInfo): Exception information.

        Returns:
            None: The return value of this handler is ignored.
        """
        self.db.update(task_id = task_id, values = {"status": "FAILURE", "task_exception": str(exc), "task_einfo": str(einfo)})

    def after_return(self, status, retval, task_id, args, kwargs, einfo):
        """Handler called after the task returns.

        Arguments:
            status (str): Current task state.
            retval (Any): Task return value/exception.
            task_id (str): Unique id of the task.
            args (Tuple): Original arguments for the task.
            kwargs (Dict): Original keyword arguments for the task.
            einfo (~billiard.einfo.ExceptionInfo): Exception information.

        Returns:
            None: The return value of this handler is ignored.
        """
        self.db.update(task_id = task_id, values = {"status": status, "task_einfo": str(einfo)})


@shared_task(bind = True, base = KedroGraphqlTask)
def run_pipeline(self, name: str, inputs: dict, outputs: dict, parameters: dict):
    catalog = {**inputs, **outputs}
    io = DataCatalog().from_config(catalog = catalog)

    ## add parameters to DataCatalog e.g. {"params:myparam":"value"}
    params = {"params:"+k:v for k,v in parameters.items()}
    params["parameters"] = parameters
    io.add_feed_dict(params)

    SequentialRunner().run(pipelines[name], catalog = io)
    
    return "success"