import json
import time
import warnings
from distutils.version import StrictVersion
from typing import Union

import pandas as pd
from requests import HTTPError

from .cas_management import CASManagement
from ..core import current_session
from .score_definitions import ScoreDefinitions
from .service import Service


class ScoreExecution(Service):
    """
    The Score Execution API is used to produce a score by
    executing the mapped code generated by score objects using the score definition.

    See Also
    --------
    `REST Documentation <https://developers.sas.com/rest-apis/scoreExecution-v2>`

    """

    _SERVICE_ROOT = "/scoreExecution"
    _cas_management = CASManagement()
    _score_definitions = ScoreDefinitions()
    _services = Service()

    (
        list_executions,
        get_execution,
        update_execution,
        delete_execution,
    ) = Service._crud_funcs("/executions", "execution")

    @classmethod
    def create_score_execution(
        cls,
        score_definition_id: str,
        description: str = "",
        output_server_name: str = "cas-shared-default",
        output_library_name: str = "Public",
        output_table_name: str = "",
    ):
        """Creates the score definition service.

        Parameters
        --------
        score_definition_id: str
            A score definition id representing score definition existing on the server that needs to be executed.
        description: str, optional
            Description of score execution. Defaults to an empty string.
        output_server_name: str, optional
            The name of the output server the output table and output library is stored in. Defaults to "cas-shared-default".
        output_library_name: str, optional
            The name of the output library the output table is stored in. Defaults to "Public".
        output_table_name: str, optional
            The name of the output table the score execution or analysis output will be stored in. Defaults to an empty string.

        Returns
        -------
        RestObj

        """

        # Gets information about the scoring object from the score definition
        score_definition = cls._score_definitions.get_definition(score_definition_id)
        if not score_definition:
            raise HTTPError
        score_exec_name = score_definition.get("name")
        model_uuid = score_definition.get("objectDescriptor").get("uri").split("/")[-1]
        model_uri = f"/modelManagement/models/{model_uuid}"
        model_name = score_definition.get("objectDescriptor").get("name")
        model_input_library = score_definition.get("inputData").get("libraryName")
        model_table_name = score_definition.get("inputData").get("tableName")

        # Defining a default output table name if none is provided
        if not output_table_name:
            output_table_name = f"{model_name}_{score_definition_id}"

        headers_score_exec = {"Content-Type": "application/json"}

        create_score_exec = {
            "name": f"{score_exec_name}",
            "description": description,
            "hints": {
                "objectURI": f"{model_uri}",
                "inputTableName": f"{model_table_name}",
                "inputLibraryName": f"{model_input_library}",
            },
            "scoreDefinitionId": f"{score_definition_id}",
            "outputTable": {
                "tableName": output_table_name,
                "libraryName": output_library_name,
                "serverName": output_server_name,
            },
        }

        # Creating the score execution
        score_execution = cls.post(
            "executions",
            data=json.dumps(create_score_exec),
            headers=headers_score_exec,
        )

        return score_execution

    @classmethod
    def poll_score_execution_state(
        cls, score_execution: Union[dict, str], timeout: int = 300
    ):
        """Checks the state of the score execution.

        Parameters
        --------
        score_execution: str or dict
            A running score_execution.
        timeout: int
            Time limit for checking the score_execution state.

        Returns
        -------
        String

        """
        if type(score_execution) is str:
            exec_id = score_execution
        else:
            exec_id = score_execution.get("id")

        start_poll = time.time()
        while time.time() - start_poll < timeout:
            score_execution_state = cls.get(f"executions/{exec_id}/state")
            if score_execution_state == "completed":
                print("Score execution state is 'completed'")
                return "completed"
            elif score_execution_state == "failed":
                # TODO: Grab score execution logs and return those
                print("The score execution state is failed.")
                return "failed"
            elif time.time() - start_poll > timeout:
                print("The score execution is still running, but polling time ran out.")
                return "timeout"

    @classmethod
    def get_score_execution_results(
        cls, score_execution: Union[dict, str], use_cas_gateway: False
    ):
        """Generates an output table for the score_execution results.

        Parameters
        --------
        score_execution: str or dict
            A running score_execution.

        Returns
        -------
        Table reference

        """
        try:
            import swat
        except ImportError:
            swat = None

        if type(score_execution) is str:
            score_execution = cls.get_execution(score_execution)

        server_name = score_execution.get("outputTable").get("serverName")
        library_name = score_execution.get("outputTable").get("libraryName")
        table_name = score_execution.get("outputTable").get("tableName")

        # If swat is not available, then
        if not swat:
            output_table = cls._no_gateway_get_results(
                server_name, library_name, table_name
            )
            return output_table
        else:
            session = current_session()
            cas = session.as_swat()
            if not use_cas_gateway:
                output_table = cls._no_gateway_get_results(
                    server_name, library_name, table_name
                )
                return output_table
            else:
                cas.loadActionSet("gateway")
                gateway_code = f"""
import pandas as pd
import numpy as np
            
table = gateway.read_table({{"caslib": "{library_name}", "name": "{table_name}"}})
            
gateway.return_table("Execution Results", df = table, label = "label", title = "title")"""

                output_table = cas.gateway.runlang(
                    code=gateway_code, single=True, timeout_millis=10000
                )
                output_table = pd.DataFrame(output_table["Execution Results"])
                return output_table

    @classmethod
    def _no_gateway_get_results(cls, server_name, library_name, table_name):
        """Helper method that builds the output table.

        Parameters
        --------
        server_name: str
            CAS server where original table is stored.
        library_name: CAS library where original table is stored.
        table_name: Table that contains row and columns information to build the output table

        Returns
        -------
        Pandas Dataframe

        """
        if pd.__version__ >= StrictVersion("1.0.3"):
            from pandas import json_normalize
        else:
            from pandas.io.json import json_normalize

        warnings.warn(
            "Without swat installed, the amount of rows from the output table that "
            "can be collected are memory limited by the CAS worker."
        )

        output_columns = cls._cas_management.get(
            f"servers/{server_name}/"
            f"caslibs/{library_name}/"
            f"tables/{table_name}/columns?limit=10000"
        )
        columns = json_normalize(output_columns)
        column_names = columns["name"].to_list()

        session = current_session()

        output_rows = session.get(
            f"casRowSets/servers/{server_name}/"
            f"caslibs/{library_name}/"
            f"tables/{table_name}/rows?limit=10000"
        )
        output_table = pd.DataFrame(
            json_normalize(output_rows.json()["items"])["cells"].to_list(),
            columns=column_names,
        )
        return output_table
