"""DFCX End to End Conversation Functions"""

# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import time
import traceback
import uuid

from typing import Dict
from operator import attrgetter
from threading import Thread

import pandas as pd

from google.cloud.dialogflowcx_v3beta1 import services
from google.cloud.dialogflowcx_v3beta1 import types
from google.api_core import exceptions as core_exceptions
from proto.marshal.collections import repeated
from proto.marshal.collections import maps

from dfcx_scrapi.core import scrapi_base
from dfcx_scrapi.core import flows
from dfcx_scrapi.core import pages

logging.basicConfig(
    format="[dfcx] %(levelname)s:%(message)s", level=logging.INFO
)

MAX_RETRIES = 3


class DialogflowConversation(scrapi_base.ScrapiBase):
    """Class that wraps the SessionsClient to hold end to end conversations
    and maintain internal session state
    """

    def __init__(
        self,
        config=None,
        creds_path: str = None,
        creds_dict: Dict = None,
        creds=None,
        agent_id: str = None,
        language_code: str = "en",
    ):

        super().__init__(
            creds_path=creds_path,
            creds_dict=creds_dict,
            creds=creds,
            agent_id=agent_id,
        )

        logging.info(
            f"create conversation with creds_path: \
                {creds_path} | agent_id: {agent_id}",
        )

        if agent_id or config["agent_path"]:
            self.agent_id = agent_id or config["agent_path"]

        self.language_code = language_code or config["language_code"]

        self.start_time = None
        self.query_result = None
        self.session_id = None
        self.turn_count = None
        self.agent_env = {}  # empty
        self.restart()
        self.flows = flows.Flows(creds=self.creds)
        self.pages = pages.Pages(creds=self.creds)

    @staticmethod
    def _get_match_type_from_map(match_type: int):
        """Translates the match_type enum int value into a more descriptive
        string.
        """
        match_type_map = {
            0: "MATCH_TYPE_UNSPECIFIED",
            1: "INTENT",
            2: "DIRECT_INTENT",
            3: "PARAMETER_FILLING",
            4: "NO_MATCH",
            5: "NO_INPUT",
            6: "EVENT",
        }

        return match_type_map[match_type]

    @staticmethod
    def _validate_test_set_input(test_set: pd.DataFrame):
        """Validates that all pages referenced in the test set exist in the
        agent.
        """
        mask = test_set.page_id.isna().to_list()
        invalid_pages = set(test_set.page_display_name[mask].to_list())

        if invalid_pages:
            raise Exception(
                "The following Pages are invalid and missing Page "
                f"IDs: \n{invalid_pages}\n\nPlease ensure that your Page "
                "Display Names do not contain typos.\nFor Default Start Page "
                "use the special page display name START_PAGE."
            )

    @staticmethod
    def progress_bar(current, total, bar_length=50, type_="Progress"):
        """Display progress bar for processing."""
        percent = float(current) * 100 / total
        arrow = "-" * int(percent / 100 * bar_length - 1) + ">"
        spaces = " " * (bar_length - len(arrow))
        print(
            f"{type_}({current}/{total})" + f"[{arrow}{spaces}] {percent:.2f}%",
            end="\r",
        )

    @staticmethod
    def _build_query_params_object(parameters, current_page, disable_webhook):
        if parameters:
            query_params = types.session.QueryParameters(
                disable_webhook=disable_webhook,
                parameters=parameters,
                current_page=current_page,
            )
        else:
            query_params = types.session.QueryParameters(
                disable_webhook=disable_webhook, current_page=current_page
            )

        return query_params

    @staticmethod
    def _build_query_input_object(input_obj, language_code):
        if "dtmf" in input_obj:
            digits = str(input_obj["dtmf"])

            finish_digit = None
            if "finish_digit" in input_obj:
                finish_digit = str(input_obj["finish_digit"])

            dtmf_input = types.session.DtmfInput(
                digits=digits, finish_digit=finish_digit
            )
            query_input = types.session.QueryInput(
                dtmf=dtmf_input,
                language_code=language_code,
            )

        elif "intent" in input_obj:
            intent_input = types.session.IntentInput(intent=input_obj["intent"])
            query_input = types.session.QueryInput(
                intent=intent_input, language_code=language_code
            )

        elif "event" in input_obj:
            event_input = types.session.EventInput(event=input_obj["event"])
            query_input = types.session.QueryInput(
                event=event_input, language_code=language_code
            )

        elif "text" in input_obj:
            text = input_obj["text"]
            logging.debug("Input text: %s", text)
            text_input = types.session.TextInput(text=text)
            query_input = types.session.QueryInput(
                text=text_input,
                language_code=language_code,
            )

        return query_input

    @staticmethod
    def _gather_text_responses(text_message):

        flat_texts = "\n".join(text_message.text)

        return flat_texts

    def _gather_response_messages(self, response_messages):
        rm_gathered = []
        for msg in response_messages:
            if msg.payload:
                msg = {
                    "payload": self.recurse_proto_marshal_to_dict(msg.payload)
                }

            elif msg.play_audio:
                msg = {"play_audio": {"audio_uri": msg.play_audio.audio_uri}}

            elif msg.live_agent_handoff:
                msg = {
                    "live_agent_handoff": self.recurse_proto_marshal_to_dict(
                        msg.live_agent_handoff.metadata
                    )
                }

            elif msg.conversation_success:
                msg = {
                    "conversation_success": self.recurse_proto_marshal_to_dict(
                        msg.conversation_success.metadata
                    )
                }

            elif msg.output_audio_text:
                msg = {"output_audio_text": msg.output_audio_text.text}

            elif msg.text:
                msg = {"text": self._gather_text_responses(msg.text)}

            rm_gathered.append(msg)

        return rm_gathered

    def _gather_query_result_parameters(self, input_parameters):
        output_parameters = {}
        for param in input_parameters:
            val = input_parameters[param]

            # If we find a RepeatedComposite (i.e. List) we will recurse
            # down and convert to lists/dics/str as needed.
            if isinstance(val, repeated.RepeatedComposite):
                val = self.recurse_proto_repeated_composite(val)

            elif isinstance(val, maps.MapComposite):
                val = self.recurse_proto_marshal_to_dict(val)

            output_parameters[param] = val

        return output_parameters

    def _page_id_mapper(self):
        """Initializes the agent_pages_map dataframe.

        This dataframe contains the flow_display_name, page_display_name,
        and page_id for each page in the agent.
        """
        agent_pages_map = pd.DataFrame()
        flow_map = self.flows.get_flows_map(agent_id=self.agent_id)
        for flow_id in flow_map.keys():

            page_map = self.pages.get_pages_map(flow_id=flow_id)

            flow_mapped = pd.DataFrame.from_dict(page_map, orient="index")
            flow_mapped["page_id"] = flow_mapped.index

            flow_mapped = flow_mapped.rename(columns={0: "page_display_name"})

            # add start page
            start_page_id = flow_id + "/pages/START_PAGE"
            flow_mapped = pd.concat(
                [
                    flow_mapped,
                    pd.DataFrame(
                        columns=["page_display_name", "page_id"],
                        data=[["START_PAGE", start_page_id]],
                    ),
                ]
            )

            flow_mapped.insert(0, "flow_display_name", flow_map[flow_id])
            agent_pages_map = pd.concat([agent_pages_map, flow_mapped])

        self.agent_pages_map = agent_pages_map.reset_index(drop=True)

    def _get_reply_results(self, utterance, page_id, results, i):
        """Get results of single text utterance to CX Agent.

        Args:
          utterance: Text to send to the bot for testing.
          page_id: Specified CX Page to send the utterance request to
          results: Pandas Dataframe to capture and store the results
          i: Internal tracking for Python Threading
        """
        response = self.reply(
            send_obj={"text": utterance}, current_page=page_id, restart=True
        )

        target_page = response["page_name"]

        results["target_page"][i] = target_page
        results["match"][i] = response["match"]

    def _get_intent_detection(self, test_set: pd.DataFrame):
        """Gets the results of a subset of Intent Detection tests.

        NOTE - This is an internal method used by run_intent_detection to
        manage parallel intent detection requests and should not be used as a
        standalone function.
        """

        self._page_id_mapper()
        test_set_mapped = pd.merge(
            test_set,
            self.agent_pages_map,
            on=["flow_display_name", "page_display_name"],
            how="left",
        )
        utterances = list(test_set_mapped["utterance"])
        page_ids = list(test_set_mapped["page_id"])

        self._validate_test_set_input(test_set_mapped)

        threads = [None] * len(utterances)
        results = {
            "target_page": [None] * len(utterances),
            "match":[None] * len(utterances),
        }
        for i, (utterance, page_id) in enumerate(zip(utterances, page_ids)):
            threads[i] = Thread(
                target=self._get_reply_results,
                args=(utterance, page_id, results, i),
            )
            threads[i].start()

        for idx, _ in enumerate(threads):
            threads[idx].join()

        test_set_mapped["target_page"] = results["target_page"]
        test_set_mapped["match"] = results["match"]
        test_set_mapped = test_set_mapped.drop(columns=["page_id"])
        intent_detection = test_set_mapped.copy()

        return intent_detection


    def restart(self):
        """Starts a new session/conversation for this agent"""
        self.session_id = uuid.uuid4()
        self.turn_count = 0

    def set_agent_env(self, param, value):
        """Setting changes related to the environment"""
        logging.info("setting agent_env param:[%s] = value:[%s]", param, value)
        self.agent_env[param] = value

    def checkpoint(self, msg=None, start=False):
        """Print a checkpoint to time progress and debug bottleneck"""
        if start:
            start_time = time.perf_counter()
            self.start_time = start_time
        else:
            start_time = self.start_time
        duration = round((time.perf_counter() - start_time), 2)
        if duration > 2:
            if msg:
                print(f"{duration:0.2f}s {msg}")

    def reply(
        self,
        send_obj: Dict[str, str],
        restart: bool = False,
        retries: int = 0,
        current_page: str = None,
        checkpoints: bool = False,
    ):
        """Runs intent detection on one utterance and gets the agent reply.

        Args:
          send_obj: Dictionary with the following structure:
            {'text': str,
            'params': Dict[str,str],
            'dtmf': str}
          restart: Boolean flag that determines whether to use the existing
            session ID or start a new conversation with a new session ID.
            Passing True will create a new session ID on subsequent calls.
            Defaults to False.
          retries: used for recurse calling this func if API fails
          current_page: Specify the page id to start the conversation from
          checkpoints: Boolean flag to enable/disable Checkpoint timer
            debugging. Defaults to False.

        Returns:
          A dictionary for the agent reply to to the submitted text.
            Includes keys response_messages, confidence, page_name,
            intent_name, match_type, match, other_intents, and params.
        """
        text = send_obj.get("text")
        send_params = send_obj.get("params")

        if not text:
            logging.warning(f"Input Text is empty. {send_obj}")

        if text and len(text) > 256:
            logging.warning(
                "Text input is too long. Truncating to 256 characters."
            )
            text = text[0:256]
            logging.warning(f"TRUNCATED TEXT: {text}")

        custom_environment = self.agent_env.get("environment")
        disable_webhook = self.agent_env.get("disable_webhook") or False

        if checkpoints:
            self.checkpoint(start=True)

        if restart:
            self.restart()

        client_options = self._set_region(self.agent_id)
        session_client = services.sessions.SessionsClient(
            credentials=self.creds, client_options=client_options
        )
        session_path = f"{self.agent_id}/sessions/{self.session_id}"

        if custom_environment:
            logging.info("req using env: %s", custom_environment)
            session_path = (
                f"{self.agent_id}/environments/"
                f"{custom_environment}/sessions/{self.session_id}"
            )

        # Build Query Params object
        query_params = self._build_query_params_object(
            send_params, current_page, disable_webhook
        )

        # Build Query Input object
        query_input = self._build_query_input_object(
            send_obj, self.language_code
        )

        request = types.session.DetectIntentRequest(
            session=session_path,
            query_input=query_input,
            query_params=query_params,
        )

        logging.debug("query_params: %s", query_params)
        logging.debug("request %s", request)

        response = None
        try:
            response = session_client.detect_intent(request=request)

        except core_exceptions.InternalServerError as err:
            logging.error(
                "---- ERROR --- InternalServerError caught on CX.detect %s", err
            )
            logging.error("text: %s", text)
            logging.error("query_params: %s", query_params)
            logging.error("query_input: %s", query_input)
            return {}

        except core_exceptions.ClientError as err:
            logging.error(
                "---- ERROR ---- ClientError caught on CX.detect %s", err
            )
            template = "An exception of type {0} occurred. \nArguments:\n{1!r}"
            message = template.format(type(err).__name__, err.args)
            logging.error("err name %s", message)

            logging.error("text %s", text)
            logging.error("query_params %s", query_params)
            logging.error("query_input %s", query_input)
            logging.error(traceback.print_exc())
            retries += 1
            if retries < MAX_RETRIES:
                logging.error("retrying")
                return self.reply(send_obj, restart=restart, retries=retries)
            else:
                logging.error("MAX_RETRIES exceeded")
                return {}

        if checkpoints:
            self.checkpoint("<< got response")
        query_result = response.query_result
        logging.debug("dfcx>qr %s", query_result)
        self.query_result = query_result
        reply = {}

        # Gather Response Messages into List of Dicts
        if query_result.response_messages:
            response_messages = self._gather_response_messages(
                query_result.response_messages
            )
        else:
            response_messages = None

        # Convert params structures from Proto to standard python data types
        if query_result.parameters:
            params = self._gather_query_result_parameters(
                query_result.parameters
            )
        else:
            params = None

        reply["response_messages"] = response_messages
        reply["confidence"] = query_result.intent_detection_confidence
        reply["page_name"] = query_result.current_page.display_name
        reply["intent_name"] = query_result.intent.display_name
        reply["match_type"] = self._get_match_type_from_map(
            query_result.match.match_type
        )
        reply["match"] = query_result.match
        reply["other_intents"] = self.format_other_intents(query_result)
        reply["params"] = params

        logging.debug("reply %s", reply)

        return reply

    def format_other_intents(self, query_result):
        """Unwind protobufs into more friendly dict"""
        other_intents = query_result.diagnostic_info.get(
            "Alternative Matched Intents"
        )
        items = []
        rank = 0
        for alt in other_intents:
            items.append(
                {
                    "name": alt.get("DisplayName"),
                    "score": alt.get("Score"),
                    "rank": rank,
                }
            )
            rank += 1

        if self:
            return items

        return None

    def getpath(self, obj, xpath, default=None):
        """Get data at a pathed location out of object internals"""
        elem = obj
        try:
            for xpitem in xpath.strip("/").split("/"):
                try:
                    xpitem = int(xpitem)
                    elem = elem[xpitem]  # dict
                except ValueError:
                    elem = elem.get(xpitem)  # array
        except KeyError:
            logging.warning("failed to getpath: %s ", xpath)
            return default

        logging.info("OK getpath: %s", xpath)
        if self:
            return elem

        return None

    def run_intent_detection(
        self,
        test_set: pd.DataFrame,
        chunk_size: int = 300,
        rate_limit: float = 20,
    ):
        """Tests a set of utterances for intent detection against a CX Agent.

        This function uses Python Threading to run tests in parallel to
        expedite intent detection testing for Dialogflow CX agents. The default
        quota for Text requests/min is 1200. Ref:
          https://cloud.google.com/dialogflow/quotas#table

        Args:
          test_set: A Pandas DataFrame with the following schema.
            flow_display_name: str
            page_display_name: str
              - NOTE, when using the Default Start Page of a Flow you must
                define it as the special display name START_PAGE
            utterance: str
          chunk_size: Determines the number of text requests to send in
            parallel. This should be adjusted based on your test_set size and
            the Quota limits set for your GCP project. Default is 300.
          rate_limit: Number of seconds to wait between running test set chunks

        Returns:
          A Pandas DataFrame consisting of the original
            DataFrame plus an additional column for the detected intent with
            the following schema.
              flow_display_name: str
              page_display_name: str
              utterance: str
              detected_intent: str
              confidence: float
              target_page: str
        """

        result = pd.DataFrame()
        for start in range(0, test_set.shape[0], chunk_size):
            test_set_chunk = test_set.iloc[start : start + chunk_size]
            result_chunk = self._get_intent_detection(test_set=test_set_chunk)
            result = pd.concat([result, result_chunk])
            self.progress_bar(start, test_set.shape[0])
            time.sleep(rate_limit)
        self.progress_bar(test_set.shape[0], test_set.shape[0])

        result = self._unpack_match(result)
        return result

    def _unpack_match(self, df: pd.DataFrame):
        """ Unpacks a 'match' column into four component columns.

        Args:
          df: dataframe containing a column named match of types.Match

        Returns:
          A copy of df with columns match_type, confidence, parameters_set,
            and detected_intent instead of match.
        """
        df = (
            df
            .copy()
            .assign(
                match_type = lambda df: df.match.apply(
                    attrgetter("match_type._name_")),
                confidence = lambda df: df.match.apply(
                    attrgetter("confidence")),
                parameters_set = lambda df: df.match.apply(
                    attrgetter("parameters")),
                detected_intent = lambda df: df.match.apply(
                    attrgetter("intent.display_name"))
            )
            .assign(
                parameters_set = lambda df: df.parameters_set.apply(
                    lambda p: self.recurse_proto_marshal_to_dict(
                        p) if p else "")
            )
            .drop(columns="match")
        )
        return df
