from flask import Flask, request, jsonify
from flask_cors import CORS
import imaplib
import email
from email.header import decode_header
import joblib

app = Flask(__name__)
CORS(app)

model = joblib.load('spam_classifier_model.pkl')
vectorizer = joblib.load('vectorizer.pkl')


@app.route('/fetch-emails', methods=['POST'])
def fetch_emails():
    data = request.json
    username = data['username']
    password = data['password']

    try:
        mail = imaplib.IMAP4_SSL("outlook.office365.com")
        mail.login(username, password)
        mail.select("inbox")
    except imaplib.IMAP4.error:
        return jsonify({"error": "Login failed. Check your email and password."}), 401

    status, messages = mail.search(None, "ALL")
    email_ids = messages[0].split()

    emails = []

    for email_id in email_ids:
        res, msg = mail.fetch(email_id, "(RFC822)")
        for response_part in msg:
            if isinstance(response_part, tuple):
                msg = email.message_from_bytes(response_part[1])
                subject, encoding = decode_header(msg["Subject"])[0]
                if isinstance(subject, bytes):
                    subject = subject.decode(encoding if encoding else "utf-8")
                from_ = msg.get("From")
                name, email_address = email.utils.parseaddr(from_)
                body = ""
                if msg.is_multipart():
                    for part in msg.walk():
                        if part.get_content_type() == "text/plain" and part.get("Content-Disposition") is None:
                            body += part.get_payload(decode=True).decode(part.get_content_charset() or "utf-8")
                else:
                    body = msg.get_payload(decode=True).decode(msg.get_content_charset() or "utf-8")

                emails.append({"id": email_id.decode(), "from": from_, "name": name, "email_address": email_address,
                               "subject": subject, "body": body})

    return jsonify(emails)


@app.route('/classify-email', methods=['POST'])
def classify_email():
    data = request.json
    email_body = data['body']
    email_vectorized = vectorizer.transform([email_body])
    prediction = model.predict(email_vectorized)
    result = "Suspicious" if prediction == 1 else "Not suspicious"
    return jsonify({"result": result})


@app.route('/mark-safe', methods=['POST'])
def mark_safe():
    data = request.json
    email_id = data['email_id']
    # Logic to mark email as safe
    return jsonify({"message": f"Email {email_id} marked as safe"})


@app.route('/delete-email', methods=['POST'])
def delete_email():
    data = request.json
    email_id = data['email_id']

    # Connect to the mail server and delete the email
    username = data['username']
    password = data['password']
    try:
        mail = imaplib.IMAP4_SSL("outlook.office365.com")
        mail.login(username, password)
        mail.select("inbox")
        mail.store(email_id, '+FLAGS', '\\Deleted')
        mail.expunge()
        return jsonify({"message": f"Email {email_id} deleted"})
    except imaplib.IMAP4.error:
        return jsonify({"error": "Failed to delete email"}), 500


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