import logging
import os
import time
from typing import Optional, Union
from uuid import UUID

from acelerai_pipelines.http_aceler_pipelines import AcelerAIHttpPipelinesClient
from acelerai_pipelines.models.run import Run, RunStatus
from acelerai_pipelines.models.pagination_data import PaginationData
from acelerai_pipelines.models.process_status import ProcessStatusType
from acelerai_pipelines.utils import validate_process_status_params

logger = logging.getLogger("PipelinesClient")


class PipelinesClient:
    """
    Synchronous client to interact with ACELER.AI Pipelines.
    """

    def __init__(self, token: str):
        """
        Constructor for PipelinesClient.
        :param token: Token to authenticate with ACELER.AI
        """
        self.acelerai_client = AcelerAIHttpPipelinesClient(token)
        self.__mode = os.environ.get("EXEC_LOCATION", "LOCAL")

    def close(self):
        """Close underlying client connections."""
        self.acelerai_client.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def execute_run_ondemand(self, deploy_key: str, payload: Optional[dict] = None) -> Run:
        """
        Execute a run on demand synchronously.
        :param deploy_key: The deploy key/id to be executed
        :param payload: The payload to be used in the run
        :return: The run object
        """
        logger.debug(f"Executing run on demand - Deploy ID: {deploy_key}")

        deploy = self.acelerai_client.get_deploy_by_fkey(deploy_key)
        if deploy is None:
            raise Exception(
                f"Deploy {deploy_key} not found in your account, please check the deploy key or your credentials"
            )

        if deploy.DeployInfo.Actived == False:
            raise Exception(
                f"Deploy {deploy_key} is not active, please activate the deploy in pipelines dashboard"
            )

        run = self.acelerai_client.execute_run_ondemand(deploy_key, payload)
        if run is None:
            raise Exception(f"Error executing run on demand - Deploy ID: {deploy_key}")
        return run

    def get_run(self, run_id: Union[UUID, str], wait_result: bool = True) -> Run:
        """
        Get the result of a run synchronously.
        :param run_id: The run id
        :param wait_result: If True, blocks and polls until the run terminates
        :return: The run object
        """
        run_id = str(run_id)
        logger.debug(f"Getting run result - Run ID: {run_id}")

        while True:
            run = self.acelerai_client.get_run(run_id)
            if run is None:
                raise Exception(f"Error fetching run - Run ID: {run_id}")

            if not wait_result:
                return run

            if run.RunStatus == RunStatus.Error:
                logger.error(f"Error in run - Run ID: {run_id}")
                return run

            if run.RunStatus == RunStatus.Completed:
                logger.debug(f"Run completed - Run ID: {run_id}")
                return run

            if run.RunStatus == RunStatus.Cancelled:
                logger.error(f"Run was cancelled - Run ID: {run_id}")
                return run

            if run.RunStatus == RunStatus.InternalError:
                logger.error(
                    f"Internal error in run, please communicate with support - Run ID: {run_id}"
                )
                return run

            if run.RunStatus == RunStatus.Lost:
                logger.error(f"Run was lost - Run ID: {run_id}")
                return run

            if run.RunStatus == RunStatus.NoResources:
                logger.error(f"Run failed (No Resources) - Run ID: {run_id}")
                return run

            if run.RunStatus == RunStatus.NotLaunched:
                logger.error(f"Run failed (Not Launched) - Run ID: {run_id}")
                return run

            if run.RunStatus == RunStatus.Pending:
                logger.debug(f"Run is pending - Run ID: {run_id}")

            if run.RunStatus == RunStatus.Running:
                logger.debug(f"Run is running - Run ID: {run_id}")

            time.sleep(3)

    def get_deploy_runs(
        self, deploy_key: str, page: int = 1, page_size: int = 10
    ) -> PaginationData[Run]:
        """
        Get runs for a deploy synchronously.
        :param deploy_key: The deploy key
        :param page: The page number
        :param page_size: The page size
        :return: The paginated list of runs
        """
        logger.debug(f"Getting runs - Deploy Key: {deploy_key}")
        data_page = self.acelerai_client.get_deploy_runs(deploy_key, page, page_size)
        if data_page is None:
            raise Exception(f"Error getting runs for deploy: {deploy_key}")
        return data_page

    def send_process_status(
        self,
        status: Union[ProcessStatusType, str],
        process_name: str,
        title: str,
        content: Optional[str] = None,
    ):
        """
        Send a process status for a run synchronously.
        :param status: ProcessStatusType enum or string ("Ok", "Error", or "Warning")
        :param process_name: Name of the process (mandatory, max 200 chars)
        :param title: Title of the sub-process status (mandatory, max 200 chars)
        :param content: Content / details (optional)
        """
        status_val, process_name, title, run_id_str = validate_process_status_params(
            status, process_name, title
        )
        logger.debug(
            f"Sending ProcessStatus for Run ID {run_id_str}: [{status_val}] {title} (Process: {process_name})"
        )
        self.acelerai_client.send_process_status(
            run_id_str, status_val, process_name, title, content
        )