app.py

import os
import json
import openai
import threading
from flask import Flask, redirect, url_for, session, request, jsonify, render_template
from google_auth_oauthlib.flow import Flow
from google.auth.transport.requests import Request
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
import pickle
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import MultinomialNB
from transformers import pipeline

os.environ['OAUTHLIB_INSECURE_TRANSPORT'] = '1'

# Flask App Configuration
app = Flask(__name__)
app.secret_key = "sm-project"

# Google OAuth2 Configuration
CLIENT_SECRETS_FILE = "credentials.json"
SCOPES = ["https://www.googleapis.com/auth/gmail.readonly"]
REDIRECT_URI = "http://127.0.0.1:5000/oauth2callback"

# OAuth Flow Setup
flow = Flow.from_client_secrets_file(
    CLIENT_SECRETS_FILE,
    scopes=SCOPES,
    redirect_uri=REDIRECT_URI
)

# Initialize the zero-shot classification pipeline from Hugging Face
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

def classify_email(email, labels, user_feedback, important_emails, normal_emails):
    """Classify a single email asynchronously."""
    if email['sender'] in user_feedback.get('not_important', []):
        normal_emails.append(email)
        return
    result = classifier(email['body'], candidate_labels=labels)
    if result['labels'][0] == 'Important':
        important_emails.append(email)
    else:
        normal_emails.append(email)

def classify_emails(emails, user_feedback):
    """Classify emails using multiple threads."""
    important_emails = []
    normal_emails = []
    labels = ["Important", "Not Important"]
    threads = []
    
    for email in emails:
        thread = threading.Thread(target=classify_email, args=(email, labels, user_feedback, important_emails, normal_emails))
        threads.append(thread)
        thread.start()
    
    for thread in threads:
        thread.join()
    
    return important_emails, normal_emails

# Function to Fetch Gmail Service
def get_gmail_service():
    creds = Credentials.from_authorized_user_info(session["credentials"])
    return build("gmail", "v1", credentials=creds)

# Gmail API - Fetch Emails with Pagination
def fetch_emails(page_token=None):
    service = get_gmail_service()
    query_params = {"userId": "me", "maxResults": 50}  # Increased batch size
    if page_token:
        query_params["pageToken"] = page_token
    results = service.users().messages().list(**query_params).execute()
    messages = results.get("messages", [])

    emails = []
    for msg in messages:
        msg_data = service.users().messages().get(userId="me", id=msg["id"]).execute()
        headers = msg_data.get("payload", {}).get("headers", [])

        subject = next((h["value"] for h in headers if h["name"] == "Subject"), "No Subject")
        sender = next((h["value"] for h in headers if h["name"] == "From"), "Unknown Sender")

        body = extract_email_body(msg_data)

        emails.append({"subject": subject, "sender": sender, "body": body})
    
    next_page_token = results.get("nextPageToken", None)
    return emails, next_page_token

# Optimized Email Body Extraction
def extract_email_body(msg_data):
    """Extract plain text body from an email message."""
    parts = msg_data.get("payload", {}).get("parts", [])
    if parts:
        for part in parts:
            if part.get("mimeType") == "text/plain":
                return part.get("body", {}).get("data", "No body available")
    return "No body available"

@app.route("/dashboard")
def dashboard():
    if "credentials" not in session:
        return redirect(url_for("login"))
    
    page_token = request.args.get("page_token")
    emails, next_page_token = fetch_emails(page_token)
    user_feedback = session.get("user_feedback", {})
    important_emails, normal_emails = classify_emails(emails, user_feedback)
    
    return render_template("dashboard.html", 
                           important_emails=important_emails, 
                           normal_emails=normal_emails, 
                           next_page_token=next_page_token)

@app.route("/feedback", methods=["POST"])
def feedback():
    data = request.json
    sender = data.get("sender")
    
    if sender:
        user_feedback = session.get("user_feedback", {})
        user_feedback.setdefault("not_important", []).append(sender)
        session["user_feedback"] = user_feedback
    
    return jsonify({"message": "Feedback received."})

@app.route("/logout")
def logout():
    session.pop("credentials", None)
    return redirect(url_for("home"))

@app.route("/")
def home():
    if "credentials" in session:
        return redirect(url_for("dashboard"))
    return render_template("index.html")

@app.route("/login")
def login():
    auth_url, _ = flow.authorization_url(prompt="consent")
    return redirect(auth_url)

@app.route("/oauth2callback")
def oauth2callback():
    flow.fetch_token(authorization_response=request.url)
    credentials = flow.credentials
    session["credentials"] = credentials_to_dict(credentials)
    return redirect(url_for("dashboard"))

def credentials_to_dict(credentials):
    return {
        "token": credentials.token,
        "refresh_token": credentials.refresh_token,
        "token_uri": credentials.token_uri,
        "client_id": credentials.client_id,
        "client_secret": credentials.client_secret,
        "scopes": credentials.scopes
    }

if __name__ == "__main__":
    app.run(debug=True)