from typing import (
    Any,
    List,
    Dict,
    Union,
    Optional,
    Set,
    Generator
)

import requests
import json
import pandas as pd
import base64
import logging

logger: logging.Logger = logging.getLogger(__name__)


class ServerClient():
    """ServerClient to make calls to a ai server instance

    Example:

    ```python
    >>> import ai_server

    # define access keys
    >>> loginKeys = {"secretKey":"<your_secret_key>","accessKey":"<your_access_key>"}

    # create connection object by passing in the secret key, access key and base url for the api
    >>> server_connection = ai_server.ServerClient(access_key=loginKeys['accessKey'], secret_key=loginKeys['secretKey'], base='<Your deployed server Monolith URL>')
    ```
    """

    # Class attribute to hold a singleton instance
    da_server = None

    def __init__(
        self,
        access_key: str,
        secret_key: str,
        base: str
    ) -> None:
        """
        Args:
            access_key (`str`):
                A user's access key is a unique identifier used for authentication and authorization. It will allow users or applications to access specific resources or perform designated actions within an ai-server instance.
            secret_key (`str`):
                A user's confidential and securely stored piece of information that complements an access key, used in cryptographic processes to validate identity and ensure secure communication
            base (`str`):
                main url to access the server api
        """
        self.access_key: str = access_key                             # set the access key as an instance attribute
        # set the secret key as an instance attribute
        self.secret_key: str = secret_key
        # set the base url as an instance attribute
        self.main_url: str = base

        if self.main_url.endswith('\/'):
            self.main_url = self.main_url[:-1]

        # TODO provide definitons for all of these attributes
        # used to keep track of the authorization header after the user has authenticated
        self.auth_headers: Dict = {}
        self.login()
        # boolean flag to make sure the users have authenticated and the cookies are set
        self.connected: bool = True
        # keep track of the the insights created from this connection
        self.open_insights: Set = set()
        self.cur_insight: str = self.make_new_insight()

        # map to keep track of requests and responses outside of this class or in separate threads
        self.monitors = {}

        # set the instance of this class as the class attribute da_server
        ServerClient.da_server = self

    def login(
        self
    ):
        """
        Register / validate the users secret & access key combination. 

        Store the cookies to be used for other api calls after authentication
        """
        combined = self.access_key + ":" + self.secret_key
        combined_enc = base64.b64encode(combined.encode('ascii'))
        headers = {
            'Authorization': f"Basic {combined_enc.decode('ascii')}"
        }
        self.auth_headers: Dict = headers.copy()
        logger.info(self.auth_headers)
        headers.update(
            {'disableRedirect': 'true'}
        )

        # authentication url
        api_url = "/auth/whoami"
        url = self.main_url + api_url

        self.r = requests.get(
            url,
            headers=headers
        )

        self.r.raise_for_status()

        # display the login response
        json_response = self.r.json()

        logger.info(json_response)

        if 'errorMessage' in json_response:
            if json_response['errorMessage'] == "null principal":
                raise AuthenticationError(
                    url + f" USERID = INVALID could not login using user access key '{self.access_key}' with invalid secret key")

            raise AuthenticationError(json_response['errorMessage'])

        self.cookies = self.r.cookies

        # display the cookies
        logger.info(self.cookies)

    # TODO add dec
    def get_auth_headers(
        self
    ) -> Dict:
        '''Get the autheroization headers used to authenticate the user'''
        return self.auth_headers

    def get_openai_endpoint(
        self
    ):
        '''Get the corrensponding openai endpoint for the CFG AI server'''
        return self.main_url + "/model/openai"

    def make_new_insight(
        self
    ) -> str:
        """
        Create a new insight (temporal space) to operate within the ai-server at set it as the current insight.
        """
        if not self.connected:
            return "Please login"

        response = requests.post(
            self.main_url + "/engine/runPixel",
            cookies=self.cookies,
            data={
                'expression': 'META | true',
                'insightId': 'new'
            }
        )

        # raise HTTP error if one occurs
        response.raise_for_status()

        json_output = response.json()
        logger.info(json_output)
        self.cur_insight = json_output['insightID']
        self.open_insights.add(self.cur_insight)

        logger.info(f"Current open insights are -- {self.open_insights}")

        return self.cur_insight

    def run_pixel(
        self,
        payload: str,
        insight_id: Optional[str] = None,
        full_response: Optional[bool] = False
    ):
        '''
        Send a pixel payload to the platforms /api/engine/runPixel endpoint.

        /api/engine/runPixel is the AI server's primary endpoint that consumes a flexible payload. The payload must contain two parameters, namely:
            1.) expression - The @payload passed is placed here and must comply and be defined in the Server's DSL (Pixel) which dictates what action should be taken

            2.) insightId - the temporal workspace identifier which isoloates the actions being performed

        Args:
            payload (`str`): DSL (Pixel) instruction on what specific action should be performed 
            insight_id (`str`): Unique identifier for the temporal worksapce where actions are being isolated. Options are:
                - insight_id = '' -> Creates a temporary one time insight created but is not stored /     cannot be referenced for future state
                - insight_id = 'new' -> Creates a new insight which can be referenced in the same sesssion
                - insight_id = '<uuid/guid string>' -> Uses an existing insight id that exists in the user session
            full_response (`bool`): Indicate whether to return the full json response or only the actions output

        Returns:
            `Union[Any, Dict]`: The output object from the runPixel response or the entire runPixel response.
        '''
        # going to create an insight if insight not available
        if not self.connected:
            return "Please login"

        if insight_id is None:
            insight_id = self.cur_insight
            if insight_id is None:
                # the insight_id is still null
                logger.info(
                    "insight_id and self.cur_insight are both undefined. Creating new insight")
                self.cur_insight = self.make_new_insight()
                insight_id = self.cur_insight

        logger.info(f"Current insight_id is set to {insight_id}")

        pixel_payload = {
            'expression': payload,
            'insightId': insight_id
        }

        api_url = "/engine/runPixel"
        response = requests.post(
            self.main_url + api_url,
            cookies=self.cookies,
            data=pixel_payload
        )

        response_dict = response.json()
        logger.info(response_dict)
        if ("ERROR" in response_dict['pixelReturn'][0]['operationType']):
            raise Exception(response_dict['pixelReturn'][0]['output'])

        if full_response:
            return response_dict
        else:
            return self.get_pixel_output(response_dict)

    def get_partial_responses(self, job_id: str) -> Generator:
        '''
        Make continuois post requests to the partial endpoint until we see a change in the message state

        Args:
            job_id (`str`): The identifier of the job being run

        Returns:
            `Generator`: A generator that can be iterated through to get partials
        '''

        assert job_id is not None, "Please provide a valid job ID."

        partial_endpoint = self.main_url + '/engine/partial'
        cookies = self.cookies
        payload = {
            'jobId': job_id,
        }

        # keep track of when the streams start coming since it could be delayed
        started_streaming = False
        while True:
            response = requests.post(
                partial_endpoint, cookies=cookies, data=payload).json()
            msg = response.get('message')
            msg_length = len(msg)

            if msg_length != 0 and not started_streaming:
                started_streaming = True
                yield msg['new']
            elif msg_length != 0:
                yield msg['new']
            elif msg_length == 0 and started_streaming:
                break

    def get_pixel_output(
        self,
        response: Dict
    ) -> Union[Any, List]:
        '''
        Utility method to grab the output of runPixel call.

        Args:
            response (`Dict`): The runPixel response from the Tomcat Server

        Returns:
            `Union[Any, List]`: The output object or a list of objects 
        '''
        main_output = response['pixelReturn'][0]['output']

        if isinstance(main_output, list):
            output = main_output[0]
        else:
            output = main_output

        return output

    def logout(
        self
    ) -> None:
        '''Closes the connection to the server.'''
        requests.get(self.main_url + "/logout/all", cookies=self.cookies)

        self.cookies = None         # reset the connection attributes for the class
        self.connected = False      # reset the connection attributes for the class

    def send_request(
        self,
        payload_struct: Dict
    ) -> None:
        '''
        Constructs the pixel payload from a PayloadStruct for various server resources such as ModelEngine, StorageEngine and DatabaseEngine

        Args:
            payload_struct (`dict`): the actual payload being sent to the AI Server

        Returns:
            `None`
        '''

        # this is a synchronous call
        # but I dont want to bother the server proxy and leave it as is
        epoc = payload_struct['epoc']

        input_payload_message = json.dumps(payload_struct)

        logger.info("Sending a PayloadStruct " + input_payload_message)

        # RemoteEngineRun is responsible for handling ModelEngine, StorageEngine and DatabaseEngine via RESTServer
        output_payload_message = self.run_pixel(
            payload="RemoteEngineRun(payload=\"<e>" +
            input_payload_message + "</e>\");",
            insight_id=payload_struct['insightId']
        )

        if epoc in self.monitors:
            self.monitors[epoc] = output_payload_message

    def get_open_insights(
        self
    ) -> List[str]:
        '''
        Retrieves a list of insight IDs created after establishing the connection.

        Args:

        Returns:
            `List[str]`: List of insight IDs
        '''
        # MyOpenInsights is pre0defined server level reactor to handle the action of getting the insightIds
        open_insights = self.run_pixel(payload='MyOpenInsights();')

        # keep track of the open insights within the python object itself
        self.open_insights = set(open_insights)

        # return a list of insight ids
        return open_insights

    def drop_insights(
        self,
        insight_ids: Union[str, List[str]]
    ) -> None:
        '''
        Utility method to close insight(s). By default, if the the current/working insight is closed, then we will attempt to set to the first insight in open_insights as the current insight

        Args:
            insight_ids (`str` | `list`): a single insight or a list of insights the user wishes to close

        Returns:
            `None`
        '''
        if isinstance(insight_ids, list):
            for id in insight_ids:
                self.run_pixel(insight_id=id, payload='DropInsight();')
                self.open_insights.discard(id)
        else:
            self.run_pixel(insight_id=insight_ids, payload='DropInsight();')
            self.open_insights.discard(insight_ids)

        if self.cur_insight not in self.open_insights:
            if (len(self.open_insights) > 0):
                self.cur_insight = self.open_insights[0]
            else:
                self.cur_insight = None

    def import_data_product(
        self,
        project_id: str,
        insight_id: str,
        sql: str
    ) -> Union[pd.DataFrame, None]:
        '''
        Accesses data within a saved insight (data product) via a REST API.

        Args:
        project_id (`str`):
            Given project/app unique identifier 
        insight_id (`str`):
            Shared insight indentifier.
        sql (`bool`):
            SQL statement that is executable on the frames within the data product

        Returns:
          `pd.DataFrame` Pandas Dataframe based on the SQL statement passed in
        '''
        base_url = self.main_url + '/project-' + project_id + \
            '/jdbc_json?insightId=' + insight_id + '&open=true&sql='

        dataProductUrl = base_url + sql
        response = requests.get(
            dataProductUrl, cookies=self.cookies
        ).json()

        try:
            return pd.DataFrame(response['dataArray'], columns=response['columns'])
        except:
            try:
                return pd.DataFrame(response['data'], columns=response['columns'])
            except:
                return None

    def upload_files(
        self,
        files: List[str],
        project_id: Optional[str] = None,
        insight_id: Optional[str] = None,
        path: Optional[str] = None
    ) -> List[str]:
        '''
        Uploads files from the local device to the server.

        Args:
            files (`List[str]`):
                List of file paths to upload from local device
            insight_id (`str`):
                Unique identifier for the temporal worksapce where actions are being isolated
            project_id (`bool`):
                Given project/app unique identifier
            path (`str`):
                Specific upload path

        Returns (`List[str]`):
            List of file names that have been successfully uploaded

        '''
        # .../Monolith/api/uploadFile/baseUpload?insightId=de43ce0d-db2e-4ab9-a807-336bb86c4ea0&projectId=4c14bc58-973f-4293-87ed-a5d32c24f418&path=version/assets/

        if isinstance(files, str):
            files = [files]

        param = ''
        path = path or ''

        if insight_id or project_id or path:
            if insight_id == None:
                insight_id = self.cur_insight

            param += f'insightId={insight_id}'

            if project_id:
                if param:
                    param += '&'
                param += f'projectId={project_id}'

            if path:
                if param:
                    param += '&'
                param += f'path={path}'
        else:
            param += f'insightId={self.cur_insight}'

        param = f'?{param}'

        url = f'{self.main_url}/uploadFile/baseUpload{param}'

        logger.info("The upload url is " + url)

        insight_file_paths = []
        for filepath in files:
            with open(filepath, 'rb') as fobj:
                response = requests.post(
                    url,
                    cookies=self.cookies,
                    files={'file': fobj}
                )
                insight_file_paths.append(response.json()[0]['fileName'])

        return insight_file_paths


class AuthenticationError(Exception):
    pass
