import aepp
from aepp import connector
from copy import deepcopy
import requests
from typing import IO, Union
import logging


class DataIngestion:
    """
    Class that manages sending data via authenticated methods.
    For Batch and Streaming messages.
    """

    loggingEnabled = False
    logger = None

    def __init__(
        self,
        config: dict = aepp.config.config_object,
        header=aepp.config.header,
        loggingObject: dict = None,
        **kwargs,
    ):
        """
        Instantiate the DataAccess class.
        Arguments:
            config : OPTIONAL : config object in the config module.
            header : OPTIONAL : header object  in the config module.
        Additional kwargs will update the header.
        """
        if loggingObject is not None and sorted(
            ["level", "stream", "format", "filename", "file"]
        ) == sorted(list(loggingObject.keys())):
            self.loggingEnabled = True
            self.logger = logging.getLogger(f"{__name__}")
            self.logger.setLevel(loggingObject["level"])
            formatter = logging.Formatter(loggingObject["format"])
            if loggingObject["file"]:
                fileHandler = logging.FileHandler(loggingObject["filename"])
                fileHandler.setFormatter(formatter)
                self.logger.addHandler(fileHandler)
            if loggingObject["stream"]:
                streamHandler = logging.StreamHandler()
                streamHandler.setFormatter(formatter)
                self.logger.addHandler(streamHandler)
        self.connector = connector.AdobeRequest(config_object=config, header=header)
        self.header = self.connector.header
        self.header.update(**kwargs)
        self.sandbox = self.connector.config["sandbox"]
        self.endpoint = (
            aepp.config.endpoints["global"] + aepp.config.endpoints["ingestion"]
        )
        self.endpoint_streaming = aepp.config.endpoints["streaming"]["collection"]
        self.STREAMING_REFERENCE = {
            "header": {
                "schemaRef": {
                    "id": "https://ns.adobe.com/{TENANT_ID}/schemas/{SCHEMA_ID}",
                    "contentType": "application/vnd.adobe.xed-full+json;version={SCHEMA_VERSION}",
                },
                "imsOrgId": "{IMS_ORG_ID}",
                "datasetId": "{DATASET_ID}",
                "createdAt": "1526283801869",
                "source": {"name": "{SOURCE_NAME}"},
            },
            "body": {
                "xdmMeta": {
                    "schemaRef": {
                        "id": "https://ns.adobe.com/{TENANT_ID}/schemas/{SCHEMA_ID}",
                        "contentType": "application/vnd.adobe.xed-full+json;version={SCHEMA_VERSION}",
                    }
                },
                "xdmEntity": {
                    "person": {
                        "name": {
                            "firstName": "Jane",
                            "middleName": "F",
                            "lastName": "Doe",
                        },
                        "birthDate": "1969-03-14",
                        "gender": "female",
                    },
                    "workEmail": {
                        "primary": True,
                        "address": "janedoe@example.com",
                        "type": "work",
                        "status": "active",
                    },
                },
            },
        }

    def createBatch(
        self,
        datasetId: str = None,
        format: str = "json",
        multiline: bool = False,
        enableDiagnostic: bool = False,
        partialIngestionPercentage: int = 0,
    ) -> dict:
        """
        Create a new batch in Catalog Service.
        Arguments:
            datasetId : REQUIRED : The Dataset ID for the batch to upload data to.
            format : REQUIRED : the format of the data send.(default json)
            multiline : OPTIONAL : If you wish to upload multi-line JSON.
        """
        if datasetId is None:
            raise ValueError("Require a dataSetId")
        if self.loggingEnabled:
            self.logger.debug(f"Using createBatch with following format ({format})")
        obj = {
            "datasetId": datasetId,
            "inputFormat": {"format": format, "isMultiLineJson": False},
        }
        if multiline is True:
            obj["inputFormat"]["isMultiLineJson"] = True
        if enableDiagnostic != False:
            obj["enableErrorDiagnostics"] = True
        if partialIngestionPercentage > 0:
            obj["partialIngestionPercentage"] = partialIngestionPercentage
        path = "/batches"
        res = self.connector.postData(self.endpoint + path, data=obj)
        return res

    def deleteBatch(self, batchId: str = None) -> str:
        """
        Delete a batch by applying the revert action on it.
        Argument:
            batchId : REQUIRED : Batch ID to be deleted
        """
        if batchId is None:
            raise ValueError("Require a batchId argument")
        if self.loggingEnabled:
            self.logger.debug(f"Starting deleteBatch for ID: ({batchId})")
        path = f"/batches/{batchId}"
        params = {"action": "REVERT"}
        res = self.connector.postData(self.endpoint + path, params=params)
        return res

    def replayBatch(self, datasetId: str = None, batchIds: list = None) -> dict:
        """
        You can replay a batch that has already been ingested. You need to provide the datasetId and the list of batch to be replay.
        Once specify through that action, you will need to re-upload batch information via uploadSmallFile method with JSON format and then specify the completion.
        You will need to re-use the batchId provided for the re-upload.
        Arguments:
            dataSetId : REQUIRED : The dataset ID attached to the batch
            batchIds : REQUIRED : The list of batchID to replay.
        """
        if datasetId is None:
            raise ValueError("Require a dataset ID")
        if batchIds is None or type(batchIds) != list:
            raise ValueError("Require a list of batch ID")
        if self.loggingEnabled:
            self.logger.debug(f"Starting replayBatch for dataset ID: ({datasetId})")
        path = "/batches"
        predecessors = [f"${batchId}" for batchId in batchIds]
        data = {
            "datasetId": datasetId,
            "inputFormat": {"format": "json"},
            "replay": {"predecessors": predecessors, "reason": "replace"},
        }
        res = self.connector.patchData(self.endpoint + path, data=data)
        return res

    def uploadSmallFile(
        self,
        batchId: str = None,
        datasetId: str = None,
        filePath: str = None,
        data: Union[list, dict] = None,
        verbose: bool = False,
    ) -> dict:
        """
        Upload a small file (<256 MB) to the filePath location in the dataset.
        Arguments:
            batchId : REQUIRED : The batchId referencing the batch processed created beforehand.
            datasetId : REQUIRED : The dataSetId related to where the data are ingested to.
            filePath : REQUIRED : the filePath that will store the value.
            data : REQUIRED : The data to be uploaded (following the type provided). List or Dictionary, depending if multiline is enabled.
            verbose: OPTIONAL : if you wish to see comments around the
        """
        if batchId is None:
            raise Exception("require a batchId")
        if datasetId is None:
            raise Exception("require a dataSetId")
        if filePath is None:
            raise Exception("require a filePath value")
        if data is None:
            raise Exception("require data to be passed")
        if verbose:
            print(f"Your data is in {type(data)} format")
        if self.loggingEnabled:
            self.logger.debug(f"uploadSmallFile as format: ({type(data)})")
        privateHeader = deepcopy(self.header)
        privateHeader["Content-Type"] = "application/octet-stream"
        path = f"/batches/{batchId}/datasets/{datasetId}/files/{filePath}"
        res = self.connector.putData(
            self.endpoint + path, data=data, headers=privateHeader
        )
        return res

    def uploadSmallFileFinish(
        self, batchId: str = None, action: str = "COMPLETE", verbose: bool = False
    ) -> dict:
        """
        Send an action to signify that the import is done.
        Arguments:
            batchId : REQUIRED : The batchId referencing the batch processed created beforehand.
            action : REQUIRED : either one of these actions:
                COMPLETE (default value)
                ABORT
                FAIL
                REVERT
        """
        if batchId is None:
            raise Exception("require a batchId")
        if action is None or action not in ["COMPLETE", "ABORT", "FAIL", "REVERT"]:
            raise Exception("Not a valid action has been passed")
        path = f"/batches/{batchId}"
        if self.loggingEnabled:
            self.logger.debug(f"Finishing upload for batch ID: ({batchId})")
        params = {"action": action}
        res = self.connector.postData(
            self.endpoint + path, params=params, verbose=verbose
        )
        return res

    def uploadLargeFileStartEnd(
        self,
        batchId: str = None,
        datasetId: str = None,
        filePath: str = None,
        action: str = "INITIALIZE",
    ) -> dict:
        """
        Start / End the upload of a large file with a POST method defining the action (see parameter)
        Arguments:
            batchId : REQUIRED : The batchId referencing the batch processed created beforehand.
            datasetId : REQUIRED : The dataSetId related to where the data are ingested to.
            filePath : REQUIRED : the filePath that will store the value.
            action : REQUIRED : Action to either INITIALIZE or COMPLETE the upload.
        """
        if batchId is None:
            raise Exception("require a batchId")
        if datasetId is None:
            raise Exception("require a dataSetId")
        if filePath is None:
            raise Exception("require a filePath value")
        params = {"action": action}
        if self.loggingEnabled:
            self.logger.debug(
                f"Starting or Ending large upload for batch ID: ({batchId})"
            )
        path = f"/batches/{batchId}/datasets/{datasetId}/files/{filePath}"
        res = self.connector.postData(self.endpoint + path, params=params)
        return res

    def uploadLargeFilePart(
        self,
        batchId: str = None,
        datasetId: str = None,
        filePath: str = None,
        data: bytes = None,
        contentRange: str = None,
    ) -> dict:
        """
        Continue the upload of a large file with a PATCH method.
        Arguments:
            batchId : REQUIRED : The batchId referencing the batch processed created beforehand.
            datasetId : REQUIRED : The dataSetId related to where the data are ingested to.
            filePath : REQUIRED : the filePath that will store the value.
            data : REQUIRED : The data to be uploaded (in bytes)
            contentRange : REQUIRED : The range of bytes of the file being uploaded with this request.
        """
        if batchId is None:
            raise Exception("require a batchId")
        if datasetId is None:
            raise Exception("require a dataSetId")
        if filePath is None:
            raise Exception("require a filePath value")
        if data is None:
            raise Exception("require data to be passed")
        if contentRange is None:
            raise Exception("require the content range to be passed")
        privateHeader = deepcopy(self.header)
        privateHeader["Content-Type"] = "application/octet-stream"
        privateHeader["Content-Range"] = contentRange
        if self.loggingEnabled:
            self.logger.debug(f"Uploading large part for batch ID: ({batchId})")
        path = f"/batches/{batchId}/datasets/{datasetId}/files/{filePath}"
        res = requests.patch(self.endpoint + path, data=data, headers=privateHeader)
        res_json = res.json()
        return res_json

    def headFileStatus(
        self, batchId: str = None, datasetId: str = None, filePath: str = None
    ) -> dict:
        """
        Check the status of a large file upload.
        Arguments:
            batchId : REQUIRED : The batchId referencing the batch processed created beforehand.
            datasetId : REQUIRED : The dataSetId related to where the data are ingested to.
            filePath : REQUIRED : the filePath that reference the file.
        """
        if batchId is None:
            raise Exception("require a batchId")
        if datasetId is None:
            raise Exception("require a dataSetId")
        if filePath is None:
            raise Exception("require a filePath value")
        if self.loggingEnabled:
            self.logger.debug(f"Head File Status batch ID: ({batchId})")
        path = f"/batches/{batchId}/datasets/{datasetId}/files/{filePath}"
        res = self.connector.headData(self.endpoint + path)
        return res

    def getPreviewBatchDataset(
        self,
        batchId: str = None,
        datasetId: str = None,
        format: str = "json",
        delimiter: str = ",",
        quote: str = '"',
        escape: str = "\\",
        charset: str = "utf-8",
        header: bool = True,
        nrow: int = 5,
    ) -> dict:
        """
        Generates a data preview for the files uploaded to the batch so far. The preview can be generated for all the batch datasets collectively or for the selected datasets.
        Arguments:
            batchId : REQUIRED : The batchId referencing the batch processed created beforehand.
            datasetId : REQUIRED : The dataSetId related to where the data are ingested to.
            format : REQUIRED : Format of the file ('json' default)
            delimiter : OPTIONAL : The delimiter to use for parsing column values.
            quote : OPTIONAL : The quote value to use while parsing data.
            escape : OPTIONAL : The escape character to use while parsing data.
            charset : OPTIONAL : The encoding to be used (default utf-8)
            header : OPTIONAL : The flag to indicate if the header is supplied in the dataset files.
            nrow : OPTIONAL : The number of rows to parse. (default 5) - cannot be 10 or greater
        """
        if batchId is None:
            raise Exception("require a batchId")
        if datasetId is None:
            raise Exception("require a dataSetId")
        if format is None:
            raise Exception("require a format type")
        params = {
            "delimiter": delimiter,
            "quote": quote,
            "escape": escape,
            "charset": charset,
            "header": header,
            "nrow": nrow,
        }
        if self.loggingEnabled:
            self.logger.debug(f"getPreviewBatchDataset for dataset ID: ({datasetId})")
        path = f"/batches/{batchId}/datasets/{datasetId}/preview"
        res = self.connector.getData(self.endpoint + path, params=params)
        return res

    def streamMessage(
        self,
        inletId: str = None,
        data: dict = None,
        synchronousValidation: bool = False,
    ) -> dict:
        """
        Send a dictionary to the connection for streaming ingestion.
        Arguments:
            inletId : REQUIRED : the connection ID to be used for ingestion
            data : REQUIRED : The data that you want to ingest to Platform.
            synchronousValidation : OPTIONAL : An optional query parameter, intended for development purposes.
                If set to true, it can be used for immediate feedback to determine if the request was successfully sent.
        """
        if inletId is None:
            raise Exception("Require a connectionId to be present")
        if data is None and type(data) != dict:
            raise Exception("Require a dictionary to be send for ingestion")
        if self.loggingEnabled:
            self.logger.debug(f"Starting Streaming single message")
        params = {"synchronousValidation": synchronousValidation}
        path = f"/collection/{inletId}"
        res = self.connector.postData(
            self.endpoint_streaming + path, data=data, params=params
        )
        return res

    def streamMessages(
        self,
        inletId: str = None,
        data: list = None,
        synchronousValidation: bool = False,
    ) -> dict:
        """
        Send a dictionary to the connection for streaming ingestion.
        Arguments:
            inletId : REQUIRED : the connection ID to be used for ingestion
            data : REQUIRED : The list of data that you want to ingest to Platform.
            synchronousValidation : OPTIONAL : An optional query parameter, intended for development purposes.
                If set to true, it can be used for immediate feedback to determine if the request was successfully sent.
        """
        if inletId is None:
            raise Exception("Require a connectionId to be present")
        if data is None and type(data) != list:
            raise Exception("Require a list of dictionary to be send for ingestion")
        if self.loggingEnabled:
            self.logger.debug(f"Starting Streaming multiple messages")
        params = {"synchronousValidation": synchronousValidation}
        data = {"messages": data}
        path = f"/collection/batch/{inletId}"
        res = self.connector.postData(
            self.endpoint_streaming + path, data=data, params=params
        )
        return res
