# Standard library imports
import json
import urllib.parse
import webbrowser
from time import sleep
from typing import cast
from colorama import Fore
from datetime import datetime
from tqdm import tqdm

# Local/application imports
from rapidata.api_client.exceptions import ApiException
from rapidata.api_client.models.campaign_artifact_model import CampaignArtifactModel
from rapidata.api_client.models.order_state import OrderState
from rapidata.api_client.models.preview_order_model import PreviewOrderModel
from rapidata.api_client.models.submit_order_model import SubmitOrderModel
from rapidata.api_client.models.preliminary_download_model import (
    PreliminaryDownloadModel,
)
from rapidata.api_client.models.workflow_artifact_model import WorkflowArtifactModel
from rapidata.rapidata_client.order.rapidata_results import RapidataResults
from rapidata.service.openapi_service import OpenAPIService
from rapidata.rapidata_client.config import (
    logger,
    managed_print,
    rapidata_config,
    tracer,
)
from rapidata.rapidata_client.api.rapidata_api_client import (
    suppress_rapidata_error_logging,
)


class RapidataOrder:
    """
    An instance of a Rapidata order.

    Used to interact with a specific order in the Rapidata system, such as starting,
    pausing, and retrieving results.

    Args:
        name: The name of the order.
        order_id: The ID of the order.
        openapi_service: The OpenAPIService instance for API interaction.
    """

    def __init__(
        self,
        name: str,
        order_id: str,
        openapi_service: OpenAPIService,
    ):
        self.id = order_id
        self.name = name
        self.__created_at: datetime | None = None
        self._openapi_service = openapi_service
        self.__workflow_id: str = ""
        self.__campaign_id: str = ""
        self.__pipeline_id: str = ""
        self._max_retries = 10
        self._retry_delay = 2
        self.order_details_page = (
            f"https://app.{self._openapi_service.environment}/order/detail/{self.id}"
        )
        logger.debug("RapidataOrder initialized")

    @property
    def created_at(self) -> datetime:
        """Returns the creation date of the order."""
        if not self.__created_at:
            self.__created_at = self._openapi_service.order_api.order_order_id_get(
                self.id
            ).order_date
        return self.__created_at

    def run(self) -> "RapidataOrder":
        """Runs the order to start collecting responses."""
        with tracer.start_as_current_span("RapidataOrder.run"):
            logger.info("Starting order '%s'", self)
            self._openapi_service.order_api.order_order_id_submit_post(
                self.id, SubmitOrderModel(ignoreFailedDatapoints=True)
            )
            logger.debug("Order '%s' has been started.", self)
            managed_print(
                f"Order '{self.name}' is now viewable under: {self.order_details_page}"
            )
            return self

    def pause(self) -> None:
        """Pauses the order."""
        with tracer.start_as_current_span("RapidataOrder.pause"):
            logger.info("Pausing order '%s'", self)
            self._openapi_service.order_api.order_order_id_pause_post(self.id)
            logger.debug("Order '%s' has been paused.", self)
            managed_print(f"Order '{self}' has been paused.")

    def unpause(self) -> None:
        """Unpauses/resumes the order."""
        with tracer.start_as_current_span("RapidataOrder.unpause"):
            logger.info("Unpausing order '%s'", self)
            self._openapi_service.order_api.order_order_id_resume_post(self.id)
            logger.debug("Order '%s' has been unpaused.", self)
            managed_print(f"Order '{self}' has been unpaused.")

    def delete(self) -> None:
        """Deletes the order."""
        with tracer.start_as_current_span("RapidataOrder.delete"):
            logger.info("Deleting order '%s'", self)
            self._openapi_service.order_api.order_order_id_delete(self.id)
            logger.debug("Order '%s' has been deleted.", self)
            managed_print(f"Order '{self}' has been deleted.")

    def get_status(self) -> str:
        """
        Gets the status of the order.

        States:
            Created: The order has been created but not started yet.\n
            Preview: The order has been set up and ready but not collecting responses yet.\n
            Submitted: The order has been submitted and is being reviewed.\n
            ManualReview: The order is in manual review - something went wrong with the automatic approval.\n
            Processing: The order is actively being processed.\n
            Paused: The order has been paused.\n
            Completed: The order has been completed.\n
            Failed: The order has failed.
        """
        with tracer.start_as_current_span("RapidataOrder.get_status"):
            return self._openapi_service.order_api.order_order_id_get(self.id).state

    def display_progress_bar(self, refresh_rate: int = 5) -> None:
        """
        Displays a progress bar for the order processing using tqdm.

        Args:
            refresh_rate: How often to refresh the progress bar, in seconds.
        """
        if refresh_rate < 1:
            raise ValueError("refresh_rate must be at least 1")

        if self.get_status() == OrderState.CREATED:
            raise Exception("Order has not been started yet. Please start it first.")

        while self.get_status() == OrderState.SUBMITTED:
            managed_print(
                f"Order '{self}' is submitted and being reviewed. Standby...", end="\r"
            )
            sleep(1)

        if self.get_status() == OrderState.MANUALREVIEW:
            raise Exception(
                f"Order '{self}' is in manual review. It might take some time to start. To speed up the process, contact support (info@rapidata.ai).\nOnce started, run this method again to display the progress bar."
            )

        with tqdm(
            total=100,
            desc="Processing order",
            unit="%",
            bar_format="{desc}: {percentage:3.0f}%|{bar}| completed [{elapsed}<{remaining}, {rate_fmt}]",
            disable=rapidata_config.logging.silent_mode,
        ) as pbar:
            last_percentage = 0
            while True:
                current_percentage = self._workflow_progress.completion_percentage
                if current_percentage > last_percentage:
                    pbar.update(current_percentage - last_percentage)
                    last_percentage = current_percentage

                if current_percentage >= 100:
                    break

                sleep(refresh_rate)

    @property
    def _workflow_progress(self):
        """
        Gets the workflow progress.
        """
        progress = None
        for _ in range(self._max_retries // 2):
            try:
                with suppress_rapidata_error_logging():
                    workflow_id = self.__get_workflow_id()
                    progress = self._openapi_service.workflow_api.workflow_workflow_id_progress_get(
                        workflow_id
                    )
                break
            except Exception:
                sleep(self._retry_delay * 2)
        if not progress:
            raise Exception("Failed to get progress. Please try again later.")
        return progress

    def get_results(self, preliminary_results: bool = False) -> RapidataResults:
        """
        Gets the results of the order.
        If the order is still processing, this method will block until the order is completed and then return the results.

        Args:
            preliminary_results: If True, returns the preliminary results of the order. Defaults to False.
                Note that preliminary results are not final and may not contain all the datapoints & responses. Only the onese that are already available.
        """
        with tracer.start_as_current_span("RapidataOrder.get_results"):
            logger.info("Getting results for order '%s'...", self)
            if preliminary_results and self.get_status() not in [OrderState.COMPLETED]:
                return self.__get_preliminary_results()

            elif preliminary_results and self.get_status() in [OrderState.COMPLETED]:
                managed_print("Order is already completed. Returning final results.")

            while (state := self.get_status()) not in [
                OrderState.COMPLETED,
                OrderState.PAUSED,
                OrderState.MANUALREVIEW,
                OrderState.FAILED,
            ]:
                sleep(5)
                logger.debug(
                    "Order '%s' is in state %s not yet completed. Waiting...",
                    self,
                    state,
                )

            try:
                return RapidataResults(
                    json.loads(
                        self._openapi_service.order_api.order_order_id_download_results_get(
                            order_id=self.id
                        )
                    )
                )
            except (ApiException, json.JSONDecodeError) as e:
                raise Exception(f"Failed to get order results: {str(e)}") from e

    def view(self) -> None:
        """
        Opens the order details page in the browser.

        Raises:
            Exception: If the order is not in processing state.
        """
        logger.info("Opening order details page in browser...")
        could_open_browser = webbrowser.open(self.order_details_page)
        if not could_open_browser:
            encoded_url = urllib.parse.quote(
                self.order_details_page, safe="%/:=&?~#+!$,;'@()*[]"
            )
            managed_print(
                Fore.RED
                + f"Please open this URL in your browser: '{encoded_url}'"
                + Fore.RESET
            )

    def preview(self) -> None:
        """
        Opens a preview of the order in the browser.

        Raises:
            Exception: If the order is not in processing state.
        """
        logger.info("Opening order preview in browser...")
        if self.get_status() == OrderState.CREATED:
            logger.info("Order is still in state created. Setting it to preview.")
            self._openapi_service.order_api.order_order_id_preview_post(
                self.id, PreviewOrderModel(ignoreFailedDatapoints=True)
            )
            logger.info("Order is now in preview state.")

        campaign_id = self.__get_campaign_id()
        auth_url = f"https://app.{self._openapi_service.environment}/order/detail/{self.id}/preview?campaignId={campaign_id}"
        could_open_browser = webbrowser.open(auth_url)
        if not could_open_browser:
            encoded_url = urllib.parse.quote(auth_url, safe="%/:=&?~#+!$,;'@()*[]")
            managed_print(
                Fore.RED
                + f"Please open this URL in your browser: '{encoded_url}'"
                + Fore.RESET
            )

    def __get_pipeline_id(self) -> str:
        """Internal method to fetch and cache the pipeline ID."""
        if not self.__pipeline_id:
            for _ in range(self._max_retries):
                try:
                    self.__pipeline_id = (
                        self._openapi_service.order_api.order_order_id_get(
                            self.id
                        ).pipeline_id
                    )
                    break
                except Exception:
                    sleep(self._retry_delay)
            else:
                raise Exception("Failed to fetch pipeline ID after retries.")
        return self.__pipeline_id

    def __get_workflow_id(self) -> str:
        """Internal method to fetch and cache the workflow ID."""
        if not self.__workflow_id:
            self.__fetch_workflow_and_campaign_ids()
        return self.__workflow_id

    def __get_campaign_id(self) -> str:
        """Internal method to fetch and cache the campaign ID."""
        if not self.__campaign_id:
            self.__fetch_workflow_and_campaign_ids()
        return self.__campaign_id

    def __fetch_workflow_and_campaign_ids(self) -> None:
        """Fetches workflow and campaign IDs from the pipeline."""
        if self.__workflow_id and self.__campaign_id:
            return
        pipeline_id = self.__get_pipeline_id()
        for _ in range(self._max_retries):
            try:
                pipeline = self._openapi_service.pipeline_api.pipeline_pipeline_id_get(
                    pipeline_id
                )
                self.__workflow_id = cast(
                    WorkflowArtifactModel,
                    pipeline.artifacts["workflow-artifact"].actual_instance,
                ).workflow_id
                self.__campaign_id = cast(
                    CampaignArtifactModel,
                    pipeline.artifacts["campaign-artifact"].actual_instance,
                ).campaign_id
                return
            except Exception:
                sleep(self._retry_delay)
        raise Exception("Failed to fetch workflow and campaign IDs after retries.")

    def __get_preliminary_results(self) -> RapidataResults:
        """Internal method to fetch preliminary results."""
        try:
            pipeline_id = self.__get_pipeline_id()
            download_id = self._openapi_service.pipeline_api.pipeline_pipeline_id_preliminary_download_post(
                pipeline_id, PreliminaryDownloadModel(sendEmail=False)
            ).download_id

            elapsed = 0
            timeout = 60
            while elapsed < timeout:
                preliminary_results = self._openapi_service.pipeline_api.pipeline_preliminary_download_preliminary_download_id_get(
                    preliminary_download_id=download_id
                )
                if preliminary_results:
                    return RapidataResults(json.loads(preliminary_results.decode()))
                sleep(1)
                elapsed += 1
            raise Exception("Timeout waiting for preliminary results.")
        except (ApiException, json.JSONDecodeError) as e:
            raise Exception(f"Failed to get preliminary results: {str(e)}") from e

    def __str__(self) -> str:
        return f"RapidataOrder(name='{self.name}', order_id='{self.id}')"

    def __repr__(self) -> str:
        return f"RapidataOrder(name='{self.name}', order_id='{self.id}')"
