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
# 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}'
    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
import traceback
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:
        try:
            payload = message['payload']
            headers = payload['headers']
            date = next(header['value'] for header in headers if header['name'] == 'Date')
            from_email = next(header['value'] for header in headers if header['name'] == 'From')
            subject = next(header['value'] for header in headers if header['name'] == 'Subject')
            try:
                body = base64.urlsafe_b64decode(payload['parts'][0]['body']['data']).decode('utf-8')
                emails.append({'date': date, 'from': from_email, 'subject': subject, 'body': body})
            except:
                print(traceback.format_exc())
                #print(headers)
                #print(payload)
            
        except:
            print(traceback.format_exc())
    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='Sheet1!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)

            if should_send_followup(last_three_emails):
                followup_email = generate_followup_email(last_three_emails)
                recipient_email = last_three_emails[-1]['from']
                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()