import awswrangler
import boto3
import botocore
import logging
import re
import time

from psi_pinpoint.errors import EndpointNotFound
from psi_pinpoint import get_logger
from typeguard import typechecked
from typing import Any, Dict, List, Union


@typechecked
class PinpointClient:
    """Wrapper class around the Boto3 Pinpoint client to make interacting with
    the service a smidge easier"""

    def __init__(
        self,
        aws_session: Union[boto3.Session, Any] = None,
        backoff_multiplier: float = 1.0,
        max_retries: int = 30,
        wait_seconds: float = 2.0,
        logger: logging.Logger = None,
    ) -> None:
        self.__aws_session = aws_session or boto3.Session()
        self.__client = self.aws_session.client("pinpoint")
        self.backoff_multiplier = backoff_multiplier
        self.max_retries = max_retries
        self.wait_seconds = wait_seconds
        self.__logger = logger or get_logger()

    @property
    def aws_session(self) -> Union[boto3.Session, Any]:
        return self.__aws_session

    @property
    def backoff_multiplier(self) -> float:
        return self.__backoff_multiplier

    @backoff_multiplier.setter
    def backoff_multiplier(self, backoff_multiplier: float):
        assert backoff_multiplier >= 1, "Backoff multiplier must be greater than or equal to 1"
        self.__backoff_multiplier = backoff_multiplier

    @property
    def client(self) -> Any:
        return self.__client

    @property
    def logger(self) -> logging.Logger:
        return self.__logger

    @property
    def max_retries(self) -> int:
        return self.__max_retries

    @max_retries.setter
    def max_retries(self, max_retries: int):
        assert max_retries >= 0, "Max retries must be greater than or equal to 0"
        self.__max_retries = max_retries

    @property
    def wait_seconds(self) -> float:
        return self.__wait_seconds

    @wait_seconds.setter
    def wait_seconds(self, wait_seconds: float):
        assert wait_seconds > 0, "Wait seconds must be greater than 0"
        self.__wait_seconds = wait_seconds

    def create_application(self, app_name: str, tags: dict) -> dict:
        """Create a new Pinpoint Application (Projecet)

        Parameters
        ----------
        app_name : str
            Friendly name of an Application/Project

        tags : dict
            tag key/tag values to apply to the app e.g. {'t_AppID': 'SVC01234'}

        Returns
        -------
        dict : {
            'Arn': 'arn:aws:mobiletargeting:us-east-1:490126044637:apps/fc3e1678715c489fb6ac6e31306c627f',
            'Id': 'fc3e1678715c489fb6ac6e31306c627f',
            'Name': 'pp-utils-test',
            'tags': {
                'owner': 'jonathan.morris@pearson.com',
                't_AppID': 'SVC01234'
            },
            'CreationDate': '2022-06-06T19:42:15.824Z'
        }
        """
        return self._make_api_call(
            "create_app",
            **{"CreateApplicationRequest": {"Name": app_name, "tags": tags}},
        )["ApplicationResponse"]

    def get_application(self, app_id: str = None, app_name: str = None) -> dict:
        """Return a Pinpoint Application (Projecet) by ID or name

        Parameters
        ----------
        app_id : str
            An Application ID (Project ID in the console) for direct lookup

        app_name : str
            Friendly name of an Application/Project. Gets a list of all apps
            in an account and tries to match by name.

        Returns
        -------
        dict: {
            'Arn': 'arn:aws:mobiletargeting:us-east-1:490126044637:apps/fc3e1678715c489fb6ac6e31306c627f',
            'Id': 'fc3e1678715c489fb6ac6e31306c627f',
            'Name': 'pp-utils-test',
            'tags': {
                'owner': 'jonathan.morris@pearson.com',
                't_AppID': 'SVC01234'
            },
            'CreationDate': '2022-06-06T19:42:15.824Z'
        }
        """
        if app_id:
            self.logger.info(f"Fetching app {app_id}")
            result = self._make_api_call(ApplicationId=app_id, api="get_app")["ApplicationResponse"]

        elif app_name:
            self.logger.info(f"Fetching app {app_name}")
            result = self._get_from_list(
                filter_key="Name",
                filter_value=app_name,
                results_selector="ApplicationsResponse",
                api="get_apps",
            )

        else:
            raise TypeError(
                f"Must provide one of {list(locals().keys())[1:]}"
            )  # Kinda hacky. Should probably use the inspect module or just write out each param
        return result

    def list_applications(self) -> list:
        """Return a list of all applications

        Returns
        -------
        list : [
            {
                'Arn': 'arn:aws:mobiletargeting:us-east-1:490126044637:apps/fc3e1678715c489fb6ac6e31306c627f',
                'Id': 'fc3e1678715c489fb6ac6e31306c627f',
                'Name': 'pp-utils-test',
                'tags': {
                    'owner': 'jonathan.morris@pearson.com',
                    't_AppID': 'SVC01234'
                },
                'CreationDate': '2022-06-06T19:42:15.824Z'
            },
            ...
        ]
        """
        return self._paginate("get_apps", "ApplicationsResponse")

    def get_campaign(
        self,
        app_id: str = None,
        app_name: str = None,
        campaign_id: str = None,
        campaign_name: str = None,
    ) -> dict:
        app_id = app_id or self.get_application(app_name=app_name)["Id"]

        if campaign_id:
            self.logger.info(f"Fetching campaign {campaign_id}")
            result = self._make_api_call(
                ApplicationId=app_id,
                CampaignId=campaign_id,
                api="get_campaign",
            )["CampaignResponse"]

        elif campaign_name:
            self.logger.info(f"Fetching campaign {campaign_name}")
            result = self._get_from_list(
                filter_key="Name",
                filter_value=campaign_name,
                results_selector="CampaignsResponse",
                api="get_campaigns",
                ApplicationId=app_id,
            )

        else:
            raise TypeError("Must provide one of campaign_id, campaign_name")
        return result

    def list_campaigns(self, app_id: str = None, app_name: str = None) -> list:
        app_id = app_id or self.get_application(app_name=app_name)["Id"]
        return self._paginate("get_campaigns", "CampaignsResponse", ApplicationId=app_id)

    def create_export_job(
        self,
        application_id: str,
        role_arn: str,
        s3_url_prefix: str,
        segment_id: str,
        segment_version: int = None,
    ) -> Dict[str, Any]:
        """Creates an export job for a given application and segment.

        Parameters
        ----------
        application_id : str
            The ID of the Pinpoint application to export from.
        role_arn : str
            The IAM role used to perform the export.
            Must have a trust policy allowing the AWS service "pinpoint.amazonaws.com" to assume the role.
        s3_url_prefix : str
            An S3 path like s3://bucket/example/export/path/ where the export results are stored.
        segment_id : str
            The ID of the Pinpoint segment to export from.
        segment_version : int
            (optional) The version of the Pinpoint segment to export from.

        Returns
        -------
        Dict[str, Any]
            The export job response dictionary with the structure:

            {
                'ApplicationId': 'string',
                'CompletedPieces': 123,
                'CompletionDate': 'string',
                'CreationDate': 'string',
                'Definition': {
                'RoleArn': 'string',
                'S3UrlPrefix': 'string',
                'SegmentId': 'string',
                'SegmentVersion': 123
                },
                'FailedPieces': 123,
                'Failures': [
                'string',
                ],
                'Id': 'string',
                'JobStatus': 'CREATED'|'PREPARING_FOR_INITIALIZATION'|'INITIALIZING'|'PROCESSING'|'PENDING_JOB'|'COMPLETING'|'COMPLETED'|'FAILING'|'FAILED',
                'TotalFailures': 123,
                'TotalPieces': 123,
                'TotalProcessed': 123,
                'Type': 'string'
            }
        """
        self.logger.info(f"Creating export job with S3 url prefix: {s3_url_prefix}")
        assert re.match(
            "^s3://([^/]+)/(.*?([^/]+)?)/$", s3_url_prefix
        ), "Invalid S3 url prefix format, requires a format like s3://bucket/example/export/path/"
        export_job_request = {
            "RoleArn": role_arn,
            "S3UrlPrefix": s3_url_prefix,
            "SegmentId": segment_id,
        }
        if segment_version:
            export_job_request["SegmentVersion"] = segment_version
        try:
            resp = self.client.create_export_job(
                ApplicationId=application_id,
                ExportJobRequest=export_job_request,
            )
            export_job_resp = resp.get("ExportJobResponse") or {}
            assert isinstance(
                export_job_resp, dict
            ), f"Expected export job response to be a dict, got: {export_job_resp}"
            self.logger.info(f"Created export job with id: {export_job_resp.get('Id')}")
            return export_job_resp
        except Exception as e:
            self.logger.error(f"Failed to create export job with error: {e}")
            raise e

    def get_export_job(
        self,
        application_id: str,
        job_id: str,
        wait_for_completion: bool = False,
    ) -> Dict[str, Any]:
        """Get an export job response. Optionally wait for the job to complete or fail before returning.

        Wait behavior is controlled by this class's backoff_multiplier, max_retries, and wait_seconds attributes.

        Parameters
        ----------
        application_id : str
            The ID of the Pinpoint application for which the export was performed.
        job_id : str
            The ID of an existing Pinpoint export job.
        wait_for_completion : bool
            (optional) Whether or not to wait for the export job to complete before returning.
            Wait behavior is controlled by this class's backoff_multiplier, max_retries, and wait_seconds attributes.

        Returns
        -------
        Dict[str, Any]
            The export job response dictionary with the structure:

            {
                'ApplicationId': 'string',
                'CompletedPieces': 123,
                'CompletionDate': 'string',
                'CreationDate': 'string',
                'Definition': {
                'RoleArn': 'string',
                'S3UrlPrefix': 'string',
                'SegmentId': 'string',
                'SegmentVersion': 123
                },
                'FailedPieces': 123,
                'Failures': [
                'string',
                ],
                'Id': 'string',
                'JobStatus': 'CREATED'|'PREPARING_FOR_INITIALIZATION'|'INITIALIZING'|'PROCESSING'|'PENDING_JOB'|'COMPLETING'|'COMPLETED'|'FAILING'|'FAILED',
                'TotalFailures': 123,
                'TotalPieces': 123,
                'TotalProcessed': 123,
                'Type': 'string'
            }
        """
        self.logger.info(f"Getting export job with id: {job_id}")
        export_job_status = "PENDING_JOB"
        wait_seconds = self.wait_seconds
        retry_count = 0
        try:
            export_job_resp = {}
            while retry_count <= self.max_retries:
                resp = self.client.get_export_job(ApplicationId=application_id, JobId=job_id)
                export_job_resp = resp.get("ExportJobResponse") or {}
                assert isinstance(
                    export_job_resp, dict
                ), f"Expected export job response to be a dict, got: {export_job_resp}"
                export_job_status = export_job_resp.get("JobStatus")
                if not wait_for_completion or export_job_status in ["COMPLETED", "FAILED"]:
                    break
                self.logger.debug(f"Got export job with status: {export_job_status}")
                time.sleep(wait_seconds)
                wait_seconds *= self.backoff_multiplier
                retry_count += 1
            self.logger.info(f"Got export job with status: {export_job_status}")
            return export_job_resp
        except Exception as e:
            self.logger.error(f"Failed to get export job with exception: {e}")
            raise e

    def get_export_job_results(
        self,
        application_id: str = None,
        job_id: str = None,
        s3_url_prefix: str = None,
        wait_for_completion: bool = False,
    ) -> List[dict]:
        """Gets the results of an export job and returns them as a list of dictionaries.
        Accepts either an application ID and job ID OR an S3 url prefix, but not both.

        To retrieve the export job results, the export job must be complete. You may use the wait
        parameter with the application ID and job ID to wait for the export job to complete.

        Parameters
        ----------
        application_id : str
            (conditional) The ID of the Pinpoint application for which the export was performed.
            If set, job_id must also be set.
            If set, s3_url_prefix cannot be set.
        job_id : str
            (conditional) The ID of an existing Pinpoint export job.
            If set, application_id must also be set.
            If set, s3_url_prefix cannot be set.
        s3_url_prefix : str
            (conditional) An S3 path like s3://bucket/example/export/path/ where the export results are stored.
            If set, application_id and job_id cannot be set.
        wait_for_completion : bool
            (optional) Whether or not to wait for the export job to complete before returning.
            This function only waits if using parameters application_id and job_id.
            Wait behavior is controlled by this class's backoff_multiplier, max_retries, and wait_seconds attributes.

        Returns
        -------
        List[dict]
            The export results as a list of dictionaries.
        """
        if application_id or job_id:
            assert (
                application_id and job_id
            ), "Cannot get export job results: you must specify both application_id and job_id together"
            assert (
                not s3_url_prefix
            ), "Cannot get export job results: cannot specify an s3_url_prefix if application_id or job_id are specified"
            export_job_resp = self.get_export_job(application_id, job_id, wait_for_completion)
            export_job_status = export_job_resp.get("JobStatus")
            assert (
                export_job_status == "COMPLETED"
            ), f"Cannot get export job results: export job has status {export_job_status}, requires status COMPLETED"
            s3_url_prefix = export_job_resp.get("Definition", {}).get("S3UrlPrefix")
        else:
            assert (
                s3_url_prefix
            ), "Cannot get export job results: must specify s3_url_prefix or application_id+job_id"

        self.logger.info(f"Getting export job results from S3 url prefix: {s3_url_prefix}")
        assert isinstance(
            s3_url_prefix, str
        ), f"S3 url prefix must be a string, got: {s3_url_prefix}"
        assert re.match(
            "^s3://([^/]+)/(.*?([^/]+)?)/$", s3_url_prefix
        ), "Invalid S3 url prefix format, requires a format like s3://bucket/example/export/path/"
        df = awswrangler.s3.read_json(
            path=s3_url_prefix,
            path_suffix=".gz",
            ignore_empty=True,
            use_threads=False,
            lines=True,
        )
        return df.to_dict(orient="records")

    def send_messages(
        self,
        app_id: str = None,
        app_name: str = None,
        message_request: dict = None,
    ) -> dict:
        """Send a message to a list of endpoints
        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/pinpoint.html#Pinpoint.Client.send_messages

        Parameters
        ----------
        app_id : str
            An Application ID (Project ID in the console) for direct lookup

        app_name : str
            Friendly name of an Application/Project. Gets a list of all apps
            in an account and tries to match by name.

        message_request: dict
            Message configuration dict. See the doc linked above for details.
            Example Email message config: {
                "MessageConfiguration": {
                    "EmailMessage": {
                        "SimpleEmail": {
                            "Subject": {
                                "Data": "Email Subject"
                            },
                            "TextPart": {
                                "Data": "This is an example plain text email body."
                            }
                        }
                    }
                }
                "Endpoints": {
                    "john.doe@pearson.com": {},
                    "jane.doe@pearson.com": {}
                }
            }

        Returns
        -------
        dict : {
            "ApplicationId": "0670d277a2284ab89136e31a53716b38",
            "RequestId": "81ef986a-61ce-44d1-bd9c-5f41f112bc6b",
            "Result": {
                "john.doe@pearson.com": {
                    "Address": "john.doe@pearson.com",
                    "DeliveryStatus": "SUCCESSFUL",
                    "MessageId": "ggoj0rsveqptvo59idadahfq59ejvh8mnjsmq5g0",
                    "StatusCode": 200,
                    "StatusMessage": "02000000ggoj0rsv-eqptvo59-idad-ahfq-59ej-vh8mnjsmq5g0-000000"
                },
                "jane.doe@pearson.com": {
                    "Address": "jane.doe@pearson.com",
                    "DeliveryStatus": "SUCCESSFUL",
                    "MessageId": "arrptporpcmj9pt7us4lrn33mf2acnsc4phveh80",
                    "StatusCode": 200,
                    "StatusMessage": "02000000arrptpor-pcmj9pt7-us4l-rn33-mf2a-cnsc4phveh80-000000"
                }
            }
        }
        """
        app_id = app_id or self.get_application(app_name=app_name)["Id"]
        try:
            return self._make_api_call(
                "send_messages",
                ApplicationId=app_id,
                MessageRequest=message_request,
            )["MessageResponse"]
        except botocore.exceptions.ClientError as ex:
            if "endpoint did not exist" in ex.response["Message"].lower():
                raise EndpointNotFound(ex.response["Message"])
            raise ex

    def send_users_messages(
        self,
        app_id: str = None,
        app_name: str = None,
        send_users_message_request: dict = None,
    ) -> dict:
        """Send a message to a list of users
        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/pinpoint.html#Pinpoint.Client.send_users_messages

        Parameters
        ----------
        app_id : str
            An Application ID (Project ID in the console) for direct lookup

        app_name : str
            Friendly name of an Application/Project. Gets a list of all apps
            in an account and tries to match by name.

        send_users_message_request: dict
            Message configuration dict. See the doc linked above for details.
            Example Email message config: {
                "MessageConfiguration": {
                    "EmailMessage": {
                        "SimpleEmail": {
                            "Subject": {
                                "Data": "Email Subject"
                            },
                            "TextPart": {
                                "Data": "This is an example plain text email body."
                            }
                        }
                    }
                }
                "Users": {
                    "john.doe@pearson.com": {},
                    "jane.doe@pearson.com": {}
                }
            }

        Returns
        -------
        dict : {
            "ApplicationId": "0670d277a2284ab89136e31a53716b38",
            "RequestId": "81ef986a-61ce-44d1-bd9c-5f41f112bc6b",
            "Result": {
                "john.doe@pearson.com": {
                    "john.doe@pearson.com": {
                        "Address": "john.doe@pearson.com",
                        "DeliveryStatus": "SUCCESSFUL",
                        "MessageId": "ggoj0rsveqptvo59idadahfq59ejvh8mnjsmq5g0",
                        "StatusCode": 200,
                        "StatusMessage": "02000000ggoj0rsv-eqptvo59-idad-ahfq-59ej-vh8mnjsmq5g0-000000"
                    }
                },
                "jane.doe@pearson.com": {
                    "jane.doe@pearson.com": {
                        "Address": "jane.doe@pearson.com",
                        "DeliveryStatus": "SUCCESSFUL",
                        "MessageId": "arrptporpcmj9pt7us4lrn33mf2acnsc4phveh80",
                        "StatusCode": 200,
                        "StatusMessage": "02000000arrptpor-pcmj9pt7-us4l-rn33-mf2a-cnsc4phveh80-000000"
                    }
                }
            }
        }
        """
        app_id = app_id or self.get_application(app_name=app_name)["Id"]
        return self._make_api_call(
            "send_users_messages",
            ApplicationId=app_id,
            SendUsersMessageRequest=send_users_message_request,
        )["SendUsersMessageResponse"]

    def get_user_endpoints(
        self, app_id: str = None, app_name: str = None, user_id: str = None
    ) -> list:
        """Lookup all endpoints associated with a given user_id
        https: //boto3.amazonaws.com/v1/documentation/api/latest/reference/services/pinpoint.html#Pinpoint.Client.get_user_endpoints

        Parameters
        ----------
        app_id : str
            An Application ID (Project ID in the console) for direct lookup

        app_name : str
            Friendly name of an Application/Project. Gets a list of all apps
            in an account and tries to match by name.

        user_id : str
            A user's unique ID. For Courier, this is just their email address.

        Returns
        -------
        list : [
            {
                'Address': 'jonathan.morris@pearson.com',
                'ApplicationId': '0670d277a2284ab89136e31a53716b38',
                'ChannelType': 'EMAIL',
                'CohortId': '45',
                'CreationDate': '2022-08-08T16: 26: 11.156Z',
                'EffectiveDate': '2022-08-08T16: 26: 11.156Z',
                'EndpointStatus': 'ACTIVE',
                'Id': 'jonathan.morris@pearson.com',
                'OptOut': 'NONE',
                'RequestId': 'df17fe05-9c6a-414e-ac60-be97f668577a',
                'User': {'UserId': 'jonathan.morris@pearson.com'}
            },
            ...
        ]
        """
        app_id = app_id or self.get_application(app_name=app_name)["Id"]
        return self._make_api_call("get_user_endpoints", ApplicationId=app_id, UserId=user_id,)[
            "EndpointsResponse"
        ]["Item"]

    def _make_api_call(self, api: str, **kwargs) -> dict:
        """Return the result of an API call

        Parameters
        ----------
        api : str
            the Pinpoint client method to invoke, such as get_app or create_app

        kwargs : Any
            Any arguments needed by the given `api` parameter
        """
        return getattr(self.client, api)(**kwargs)

    def _get_from_list(
        self,
        filter_key: str,
        filter_value: Any,
        results_selector: str,
        api: str = None,
        **kwargs,
    ) -> dict:
        """Look up a resource by name or any other attribute and return the
        first match

        Parameters
        ----------
        filter_key : str
            The resource attribute to match on e.g. "Name"

        filter_value : Any
            The attribute value to match e.g. "psi-general-coms"

        results_selector : str
            Top-level key used to identify the collection of resources returned
            by a LIST call, e.g. `JourneysResponse` for the `list_journeys` API

        api : str
            the API to call e.g. `list_journeys`

        kwargs: Any
            Any additional arguments to pass to the API
        """
        results = self._paginate(api, results_selector, **kwargs)
        for result in results:
            if result.get(filter_key) == filter_value:
                return result
        return {}

    def _paginate(self, api: str, results_selector: str, **kwargs) -> list:
        """Paginate a list of results with a native Pinpoint paginator if
        available or manually with results that return `NextToken`

        Parameters
        ----------
        results_selector : str
            Top-level key used to identify the collection of resources returned
            by a LIST call, e.g. `JourneysResponse` for the `list_journeys` API

        api : str
            the API to call e.g. `list_journeys`

        kwargs: Any
            Any additional arguments to pass to the API"""
        results = []
        if self.client.can_paginate(api):
            paginator = self.client.get_paginator(api)
            iterator = paginator.paginate()
            results.extend(result.get(results_selector, []) for result in iterator)
        else:
            self.logger.warning(f"Unable to get paginator for {api}. Trying manually...")
            manual = getattr(self.client, api)(**kwargs)
            results.extend(manual[results_selector].get("Item", []))
            while manual.get("NextToken"):
                manual = getattr(self.client, api)(manual["NextToken"], **kwargs)
                results.extend(manual[results_selector].get("Item", []))
        return results


if __name__ == "__main__":
    pinpoint_client = PinpointClient()
    message_response = pinpoint_client.send_messages(
        app_id="0670d277a2284ab89136e31a53716b38",
        message_request={
            "MessageConfiguration": {
                "EmailMessage": {
                    "SimpleEmail": {
                        "Subject": {"Data": "Email Subject"},
                        "TextPart": {"Data": "This is an example plain text email body."},
                    }
                }
            },
            "Endpoints": {"john.doe@pearson.com": {}, "jane.doe@pearson.com": {}},
        },
    )
    print(message_response)
