import openai
import os
from dotenv import load_dotenv
import streamlit as st

load_dotenv()
openai.api_key = os.getenv("OPENAI_API_KEY")


function_descriptions = [
    {
        "name": "extract_info_from_email",
        "description": "categorise & extract key info from an email, such as use case, company name, contact details, etc.",
        "parameters": {
            "type": "object",
            "properties": {
                "companyName": {
                    "type": "string",
                    "description": "the name of the company that sent the email"
                },                                        
                "product": {
                    "type": "string",
                    "description": "Try to identify which product the client is interested in selling or buying, if any"
                },
                "type": {
                    "type": "string",
                    "description": "Try to identify if the email sender is trying to sell a product or wants to buy a product, if any"
                },
                "amount":{
                    "type": "string",
                    "description": "Try to identify the amount of products the client wants to purchase or sell, if any"
                },
                "category": {
                    "type": "string",
                    "description": "Try to categorise this email into categories like those: 1. Sales; 2. customer support; 3. consulting; 4. partnership; 5. Marketing; etc."
                },
                "nextStep":{
                    "type": "string",
                    "description": "What is the suggested next step to move this forward?"
                },
                "priority": {
                    "type": "string",
                    "description": "Try to give a priority score to this email based on how likely this email will leads to a good business opportunity, from 0 to 10; 10 most important"
                },
            },
            "required": ["companyName", "amount", "product", "type","priority", "category", "nextStep"]
        }
    }
]

test_email = """
Dear Jason 
I hope this message finds you well. I'm Shirley from Gucci;

I'm looking to purchase some company T-shirt for my team, we are a team of 100k people, and we want to get 2 t-shirt per person.

Please let me know the price and timeline you can work with;

Looking forward

Shirley Lou
"""

email = st.text_area("Email:",test_email, height=350)

def extract_data(email,function_descriptions):
    prompt = f"Please extract key information from this email: {email} "
    message = [{"role": "user", "content": prompt}]
    
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo-0613",
        #model="gpt-4",
        messages=message,
        functions = function_descriptions, 
        function_call="auto"
    )

    print(response)
    return response.choices[0]["message"]["function_call"]["arguments"]

if st.button("Extract Data"):
    arguments = extract_data(email,function_descriptions)
    st.json(arguments)