import os
import seaborn as sns
from openai import OpenAI
import gradio as gr
import re
import matplotlib.pyplot as plt
from app.utils import (
    format_prediction_ouptut,
    create_input_instruction,
    sentiment_flow_plot,
    display_sentiment_score_table,
    EXAMPLE_CONVERSATIONS,
)

openai_args = {"api_key": ""}

SAMPLE_INPUT = """
Visitor: Heyyy
Visitor: How are you this evening
Agent: better now ;) call me
Visitor: I am at work for now, be off around 10pm
Visitor: Need some company
Visitor: Are you independent honey
Agent: well since you arent available at the moment ill just come out and say-these sites are bad news. \
    did you know that most of the girls on here are here against their will? \
    Most of them got dragged into this lifestyle by an abuser, \
    oftentimes before they were of legal consenting age. isnt that sad?
Agent: we are with some guys who are trying to spread awareness of the realities of this "industry".
Agent: https://exoduscry.com/choice/
Visitor: Thanks
Agent: i encourage you to watch this video. it is jarring to think about how bad someone else's options must be to choose to be on these sites
Visitor: Ooohhh
Agent: selling their body to make ends meet or appease a pimp
Visitor: That's really awful
Agent: it is. you seem like the kind of guy who wouldnt wont to proliferate that kind of harmful lifestyle. am i right in thinking that?
Visitor: Well iam just looking for attention
Visitor: My marriage is not going well lol
Agent: i know that it is hard to find ourselves lonely and without much alternative to meet that perceived need but \
    its humbling to think that our needs can force someone else into such a dark place
Agent: hey, thanks for sharing that my man. i know it can be hard
Agent: marraige is the most humbling of relationships, isnt it?
Visitor: She leaves with her friends n no time for me
Agent: ive been there my guy. i know that it is alot easier to numb that loneliness for sure
Visitor: I want to be faithful
Agent: does your wife know how you feel when she chooses her friends instead of you?
Visitor: I been drinking lately
Visitor: Yes, she takes pills
Agent: if so, i hope you are praying for her to realize the hurt she is causing and to seek change
Visitor: She had surgery 4 yes ago n it's been hard for her n her addiction on pills
Visitor: Yes for now i am looking for a female friend to talk n see what can we do for each other
Agent: that is hard my man. physical pain is a huge obstacle in life for sure so i hear you
Visitor: Well chat later. thanks
Agent: have you considered pursuing other men who can encourage you instead of looking for the easy way out?
Agent: what is your name my friend? i will be praying for you by name if you wouldnt mind sharing it
Agent: well, i gotta run. watch that video i sent and i will definitely be praying for you. \
    I hope you pray for yourself and for your wife - God can definitely intervene and cause complete change in the situation if He wills it. \
    He is good and He hears you. You are loved by Him, brother. Good night
"""

SAMPLE_OUTPUT = """
Visitor: Heyyy
[Greeting]
Visitor: How are you this evening
[Greeting]
Agent: better now ;) call me
[Openness]
Visitor: I am at work for now, be off around 10pm
[Interest]
Visitor: Need some company
[Interest]
Visitor: Are you independent honey
[Interest]
Agent: well since you arent available at the moment ill just come out and say-these sites are bad news. \
    did you know that most of the girls on here are here against their will? \
    Most of them got dragged into this lifestyle by an abuser, \
    oftentimes before they were of legal consenting age. isnt that sad?
[Informative]
Agent: we are with some guys who are trying to spread awareness of the realities of this "industry".
[Informative]
Agent: https://exoduscry.com/choice/
[Informative]
Visitor: Thanks
[Acceptance]
Agent: i encourage you to watch this video. it is jarring to think about how bad someone else's options must be to choose to be on these sites
[Informative]
Visitor: Ooohhh
[Interest]
Agent: selling their body to make ends meet or appease a pimp
[Informative]
Visitor: That's really awful
[Remorse]
Agent: it is. you seem like the kind of guy who wouldnt wont to proliferate that kind of harmful lifestyle. am i right in thinking that?
[Accusatory]
Visitor: Well iam just looking for attention
[Anxious]
Visitor: My marriage is not going well lol
[Anxious]
Agent: i know that it is hard to find ourselves lonely and without much alternative to meet that perceived need but \
    its humbling to think that our needs can force someone else into such a dark place
[Informative]
Agent: hey, thanks for sharing that my man. i know it can be hard
[Acceptance]
Agent: marraige is the most humbling of relationships, isnt it?
[Openness]
Visitor: She leaves with her friends n no time for me
[Annoyed]
Agent: ive been there my guy. i know that it is alot easier to numb that loneliness for sure
[Acceptance]
Visitor: I want to be faithful
[Acceptance]
Agent: does your wife know how you feel when she chooses her friends instead of you?
[Curiosity]
Visitor: I been drinking lately
[Anxious]
Visitor: Yes, she takes pills
[Anxious]
Agent: if so, i hope you are praying for her to realize the hurt she is causing and to seek change
[Interest]
Visitor: She had surgery 4 yes ago n it's been hard for her n her addiction on pills
[Anxious]
Visitor: Yes for now i am looking for a female friend to talk n see what can we do for each other
[Informative]
Agent: that is hard my man. physical pain is a huge obstacle in life for sure so i hear you
[Acceptance]
Visitor: Well chat later. thanks
[Openness]
Agent: have you considered pursuing other men who can encourage you instead of looking for the easy way out?
[Informative]
Agent: what is your name my friend? i will be praying for you by name if you wouldnt mind sharing it
[Openness]
Agent: well, i gotta run. watch that video i sent and i will definitely be praying for you. \
    I hope you pray for yourself and for your wife - God can definitely intervene and cause complete change in the situation if He wills it. \
    He is good and He hears you. You are loved by Him, brother. Good night
[Openness]

# Conversation Sentiment Analysis Report

The Visitor begins the conversation with a friendly and casual tone, expressing a desire for company and showing interest in the Agent. \
However, as the Agent provides information about the harsh realities of the commercial sex industry, the Visitor's sentiment shifts to acceptance of the information \
and a sense of confusion and remorse about the situation.

The Visitor then reveals personal issues, indicating anxiety and seeking attention due to marital problems. \
The sentiment continues to be anxious as the Visitor discusses personal struggles with alcohol and his wife's pill addiction, \
showing a need for companionship and support.

Despite the heavy topics, the Visitor expresses a desire to remain faithful and shows interest in finding a friend, albeit with a hint of desperation. \
The Visitor openly takes the Agent's information and the conversation flows smoothly as both the Visitor and the Agent \
show openness toward each other.
"""


def gpt_process_response(resp) -> list:
    """Preprocess the response from GPT into lists of speakers, messages, labels, and
    the summary of the conversation.

    Args:
        resp (str): response from gpt

    Raises:
        gr.Error: if GPT produces response with invalid format.

    Returns:
        list: list of messages with label
    """
    response_pattern = "```(.*)```(.*)"

    try:
        labeled_conv, summary = re.search(response_pattern, resp, re.DOTALL).groups()
        labeled_conv = labeled_conv.strip()
        summary = summary.strip()

        msg_pattern = r"(Agent|Visitor): (.*)\n\[(.*)\]"
        labeled_messages = re.findall(msg_pattern, labeled_conv)

        speakers, messages, labels = [], [], []
        for speaker, message, label in labeled_messages:
            speakers.append(speaker)
            messages.append(message)
            labels.append(label)

        return speakers, messages, labels, summary
    except:
        raise gr.Error("GPT produced output in wrong format!")


def get_completion(conversation, model="gpt-4-1106-preview"):
    prompt = f"""
    The EPIK Project is about mobilizing male allies \
    to disrupt the commercial sex market, \
    equipping them to combat the roots of exploitation \
    and encouraging them to collaborate effectively \
    with the wider anti-trafficking movement. \
    You are an adept expert conversation sentiment analyzer. \
    Your job is to analyze the conversation and provide a report \
    based on the sentiment flow of the conversation on the visitor's \
    perspective. Visitor indicates the potential buyer, and Agent indicates the volunteer from EPIK. \
    The conversation is going to be given in the format:
    
    Visitor: <Visitor's message here>
    Agent: <Agent's message here>
    
    The actual conversation is delimited by triple backticks
    ```{conversation}```
    
    Here is the list of sentiment labels you should use delimited by square brackets. \
    ["Openness", "Anxious", "Confused", "Disapproval", "Remorse", "Accusatory", \
    "Denial", "Obscene", "Uninterested", "Annoyed", "Informative", "Greeting", \
    "Interest", "Curiosity", "Acceptance"]
    
    Your output should look like:
    ```
    Speaker: <Speaker's message here>
    [sentiment label]
    ...
    Speaker: <Speaker's message here>
    [sentiment label]
    ```
    
    where Speaker can either be Visitor or Agent. Then, you should write your report on the sentiment flow \
    on the Visitor's side below.

    Here is a sample input delimited by triple backticks

    ```{SAMPLE_INPUT}```

    Here is a same output that you should try to aim for delimited by sqaure brackets

    [{SAMPLE_OUTPUT}]
    """
    try:
        client = OpenAI(api_key=openai_args["api_key"])
        messages = [{"role": "user", "content": prompt}]
        response = client.chat.completions.create(
            model=model,
            messages=messages,
            temperature=0,  # this is the degree of randomness of the model's output
        )

        speakers, sentences, labels, summary = gpt_process_response(
            response.choices[0].message.content
        )
        return format_prediction_ouptut(speakers, sentences, labels), summary
    except BaseException:
        raise gr.Error("Can't connect to GPT! Please input a valid OpenAI's API key.")


def set_key(key):
    openai_args["api_key"] = key
    return


def gpt_ui():
    with gr.Blocks() as gpt_model:
        gr.Markdown("# GPT 4.0")
        with gr.Row():
            with gr.Column(scale=3):
                gr.Markdown(
                    """Another approach to analyze and predict the sentiment labels for each
                    message in a conversation is to utilize the power of GPT as a Large Language
                    Model. As of December 2023, GPT 4.0 is the latest version of GPT under the
                    name `gpt-4-1106-preview`, and it is connected with the app through Python's
                    `openai` module.
                    
                    One advantage of GPT over COSMIC is that it can also generate a
                    summary for the conversation beside the sentiment labels. However,
                    the trade-off is that GPT is not free to use. With
                    `gpt-4-1106-preview` model, the pricing is $0.01/1,000 input tokens
                    and $0.03/1,000 output tokens.  
                    
                    Before running the model, please log in with your OpenAI's API key first.
                    The key can be created with OpenAI
                    [here](https://platform.openai.com/api-keys). Make sure to save your API key
                    in a secured file for future use. 
                    """
                )
            with gr.Column(scale=1):
                api_key = gr.Textbox(
                    label="OpenAI's API key",
                    lines=1,
                    type="password",
                    placeholder="Please enter you API key here",
                )
                btn_key = gr.Button(value="Submit Key")
                btn_key.click(set_key, inputs=api_key)

        create_input_instruction()
        with gr.Row():
            with gr.Column():
                example_dropdown = gr.Dropdown(
                    choices=["-- Not Selected --"] + list(EXAMPLE_CONVERSATIONS.keys()),
                    value="-- Not Selected --",
                    label="Select an example",
                )

                gr.Markdown('<p style="text-align: center;color: gray;">--- OR ---</p>')

                conversation_input = gr.Textbox(
                    label="Input your conversation",
                    placeholder="Plese input your conversation here",
                    lines=15,
                    max_lines=15,
                )

                def on_example_change(input):
                    if input in EXAMPLE_CONVERSATIONS:
                        return EXAMPLE_CONVERSATIONS[input]

                    return ""

                example_dropdown.input(
                    on_example_change,
                    inputs=example_dropdown,
                    outputs=conversation_input,
                )

            with gr.Column():
                output_box = gr.Textbox(
                    value="",
                    label="Predicted Sentiment Labels",
                    lines=22,
                    max_lines=22,
                    interactive=False,
                )

        btn = gr.Button(value="Submit")
        report_md = gr.Markdown(value="")

        btn.click(
            get_completion,
            inputs=conversation_input,
            outputs=[output_box, report_md],
        )

        gr.Markdown("# Sentiment Flow Plot")
        with gr.Row():
            with gr.Column(scale=1):
                display_sentiment_score_table()
            with gr.Column(scale=2):
                plot_box = gr.Plot(label="Analysis Plot")

        plot_btn = gr.Button(value="Plot Sentiment Flow")

        plot_btn.click(sentiment_flow_plot, inputs=[output_box], outputs=[plot_box])

        # reset all outputs whenever a change in the input is detected
        conversation_input.change(
            lambda x: ("", "", None),
            conversation_input,
            outputs=[output_box, report_md, plot_box],
        )

    return gpt_model