import os
import openai
from google.oauth2.credentials import Credentials
from google.auth.transport.requests import Request
#from google.auth.transport.urllib3 import Request
from google_auth_oauthlib.flow import InstalledAppFlow
from googleapiclient.discovery import build
#import gspread
from datetime import datetime, timedelta
import base64
import json
from dateutil.parser import parse as parse_date
import traceback
# Set up the OpenAI API key
openai_api_key = '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'
from openai import OpenAI

client = OpenAI(
  api_key=openai_api_key,
)
# Set up the scopes for Gmail and Sheets APIs
SCOPES = ['https://www.googleapis.com/auth/gmail.readonly',
          'https://www.googleapis.com/auth/gmail.send',
          'https://www.googleapis.com/auth/spreadsheets']

# Authenticate and initialize Gmail and Sheets APIs
def authenticate_gmail_sheets():
    creds = None
    if os.path.exists('token.json'):
        creds = Credentials.from_authorized_user_file('token.json', SCOPES)
    if not creds or not creds.valid:
        if creds and creds.expired and creds.refresh_token:
            creds.refresh(Request())
        else:
            flow = InstalledAppFlow.from_client_secrets_file('credentials.json', SCOPES)
            creds = flow.run_local_server(port=0)
        with open('token.json', 'w') as token:
            token.write(creds.to_json())
    gmail_service = build('gmail', 'v1', credentials=creds)
    sheets_service = build('sheets', 'v4', credentials=creds)
    return gmail_service, sheets_service

# Get the emails sent by me where the response from the recipient was more than 5 days ago or no response
def get_old_emails(gmail_service):
    #query = f'from:{my_email}'
    query = f'from:me after:5d'
    # days = 7
    # today = datetime.utcnow()
    # start_date = today - timedelta(days=days+2)  # Considering 2 weekend days
    # query = f'from:me after:{start_date.strftime("%Y/%m/%d")}'
    result = gmail_service.users().messages().list(userId='me', q=query).execute()
    messages = result.get('messages', [])
    
    old_emails = []
    for msg in messages:
        message = gmail_service.users().messages().get(userId='me', id=msg['id']).execute()
        headers = message['payload']['headers']
        thread_id = message['threadId']
        date_str = next(header['value'] for header in headers if header['name'] == 'Date')
        #date = datetime.strptime(date_str, '%a, %d %b %Y %H:%M:%S %z')
        try:
            date = parse_date(date_str)
        except ValueError:
            print(f"Error parsing date: {date_str}")
            continue

        if (datetime.now(date.tzinfo) - date).days > 5:
            thread = gmail_service.users().threads().get(userId='me', id=thread_id).execute()
            messages_in_thread = thread['messages']
            if len(messages_in_thread) == 1 or all(msg['labelIds'] == ['SENT'] for msg in messages_in_thread[1:]):
                old_emails.append(thread_id)
    return old_emails

# Get the last 3 emails from a thread
# Get the last 3 emails from a thread
def get_last_three_emails(gmail_service, thread_id):
    thread = gmail_service.users().threads().get(userId='me', id=thread_id).execute()
    messages = thread['messages'][-3:]
    emails = []
    for message in messages:
        payload = message['payload']
        headers = payload['headers']
        date_str = next(header['value'] for header in headers if header['name'] == 'Date')
        
        # Initialize body content as empty string
        body = ""
        
        # Check if the message has parts
        if 'parts' in payload:
            for part in payload['parts']:
                # Check if the part has a body with data
                if 'body' in part and 'data' in part['body']:
                    # Decode and append the part's body data
                    body += base64.urlsafe_b64decode(part['body']['data']).decode('utf-8', errors='ignore') + "\n"

        # If no parts or body data found, try decoding the entire payload body
        elif 'body' in payload:
            body = base64.urlsafe_b64decode(payload['body']['data']).decode('utf-8', errors='ignore') + "\n"
        
        # Append each email's details to the list
        emails.append({'date': date_str, 'from': message['sender'], 'subject': payload['subject'], 'body': body})
    
    return emails


# Generate follow-up email using GPT-4
def generate_followup_email(last_three_emails):
    prompt = "Based on the following email thread, generate a professional follow-up email:\n\n"
    for email in last_three_emails:
        prompt += f"From: {email['from']}\nDate: {email['date']}\nSubject: {email['subject']}\n\n{email['body']}\n\n"
    prompt += "Follow-up email:"
    response = client.chat.completions.create(
        messages=[
            {
                "role": "system",
                "content": ""
            },
        {
                "role": "user",
                "content": prompt
  # Assuming query_text is defined earlier
            }
        ],
        model="gpt-4o",
        max_tokens=500
        )
    #response = openai.Completion.create(engine="gpt-4", prompt=prompt, max_tokens=500)
    return response.choices[0].message.content

# Send follow-up email using Gmail API
def send_followup_email(gmail_service, thread_id, to_email, followup_email):
    message = gmail_service.users().messages().send(userId='me', body={
        'raw': base64.urlsafe_b64encode(f'MIME-Version: 1.0\nContent-Type: text/plain; charset="UTF-8"\n\n{followup_email}'.encode('utf-8')).decode('utf-8'),
        'threadId': thread_id
    }).execute()
    return message

# Detect the context type of the follow-up email
def detect_context_type(followup_email):
    context_types = ["Subscription", "Support", "Inquiry", "Feedback", "Updates", "Billing", "Promotion", "Newsletter", "Webinar", "Partnership", "Demo", "Training", "Renewal", "Survey", "Announcement", "Invitation", "Reminder", "Confirmation", "Cancellation", "Report", "Alert", "Onboarding", "Discount", "Referral", "Technical Issue"]
    
    response = client.chat.completions.create(
        messages=[
            {
                "role": "system",
                "content": ""
            },
        {
                "role": "user",
                "content": f"Classify the following email into one of the following categories: {', '.join(context_types)}:\n\n{followup_email}\n\nCategory:"
  # Assuming query_text is defined earlier
            }
        ],
        model="gpt-4o",
        max_tokens=10
        )

    response_text = response.choices[0].message.content
    
    #prompt = f"Classify the following email into one of the following categories: {', '.join(context_types)}:\n\n{followup_email}\n\nCategory:"
    #response = openai.Completion.create(engine="gpt-4", prompt=prompt, max_tokens=10)
    return response_text

# Update Google Spreadsheet
def update_google_sheet(sheets_service, spreadsheet_id, sheet_name, data):
    sheet = sheets_service.spreadsheets()
    sheet.values().append(spreadsheetId=spreadsheet_id, range=(sheet_name + '!A1'), valueInputOption='RAW', insertDataOption='INSERT_ROWS', body={'values': data}).execute()

def should_send_followup(last_three_emails):
    # Create the prompt for GPT-4
    prompt = "Based on the following email thread, determine if a follow-up email is required. Respond with 'Yes' or 'No'.\n\n"
    for email in last_three_emails:
        prompt += f"From: {email['from']}\nDate: {email['date']}\nSubject: {email['subject']}\n\n{email['body']}\n\n"
    prompt += "Is a follow-up email required? Answer yes or no"

    response = client.chat.completions.create(
            messages=[
                {
                    "role": "system",
                    "content": ""
                },
            {
                    "role": "user",
                    "content": prompt
    # Assuming query_text is defined earlier
                }
            ],
            model="gpt-4o",
            max_tokens=10,
            temperature=0
            )

    followup_needed = response.choices[0].message.content.lower()
    print(followup_needed)

    # Extract the response text and strip any leading/trailing whitespace
    #followup_needed = response.choices[0].text.strip()

    # Return True if follow-up is needed, False otherwise
    return followup_needed.lower() == 'yes'

# Main function to automate follow-up emails
def automate_followup_emails():
    spreadsheet_id = '1AziMFQdivL-na1G-DtUbARofrdemBgf-rjlPXRhrg1g'

    gmail_service, sheets_service = authenticate_gmail_sheets()
    old_emails = get_old_emails(gmail_service)
    
    for thread_id in old_emails:
        try:
            last_three_emails = get_last_three_emails(gmail_service, thread_id)
            recipient_email = last_three_emails[-1]['from']
            if should_send_followup(last_three_emails):
                followup_email = generate_followup_email(last_three_emails)
                context_type = detect_context_type(followup_email)
                #print(recipient_email, my_email, context_type,followup_email)
                
                #send_followup_email(gmail_service, thread_id, recipient_email, followup_email)
                
                data = [
                    [recipient_email, context_type, last_three_emails[-1]['date'], str(datetime.now()), last_three_emails[-1]['body'], followup_email]
                ]
                #print(data)
                update_google_sheet(sheets_service, spreadsheet_id, 'Sheet1', data)
            else:
                data = [
                    [recipient_email, context_type, last_three_emails[-1]['date'], str(datetime.now()), last_three_emails[-1]['body'], 'No follow-up email needed']
                ]
                update_google_sheet(sheets_service, spreadsheet_id, 'Sheet2', data)
        except:
            print(traceback.format_exc())


automate_followup_emails()