import gradio as gr
import asyncio
import nest_asyncio
import re
import io
import sys
import os

from protonmail import ProtonMail

nest_asyncio.apply()

username = "laygyi5566@proton.me"
password = "htz175039"

# Session setup
proton = ProtonMail()
if not os.path.exists("session.pickle"):
    proton.login(username, password)
    proton.save_session("session.pickle")

proton.load_session("session.pickle", auto_save=True)


async def get_verification_code_async():
    # Capture stdout for logging
    old_stdout = sys.stdout
    sys.stdout = mystdout = io.StringIO()

    try:
        messages = proton.get_messages()
        if not messages:
            return "No messages found.", "", "No verification code found.", mystdout.getvalue()

        message = proton.read_message(messages[0])
        sender = message.sender.address
        subject = message.subject
        body = message.body

        # Match 6-digit or 8-char hex (Runway)
        code_match_6 = re.search(r'\b\d{6}\b', body)
        code_match_runway = re.search(r'\b[A-F0-9]{8}\b', body)

        if code_match_6:
            verification_code = code_match_6.group()
        elif code_match_runway:
            verification_code = code_match_runway.group()
        else:
            verification_code = "No verification code found."

        # Save email HTML
        with open("message.html", "w", encoding="utf-8") as f:
            f.write(body)

        logs = mystdout.getvalue()
        return sender, subject, verification_code, logs

    finally:
        sys.stdout = old_stdout


def get_verification_code_sync():
    return asyncio.get_event_loop().run_until_complete(get_verification_code_async())


def display_verification():
    sender, subject, code, logs = get_verification_code_sync()

    if code != "No verification code found.":
        html_code = f"""
        <div style="font-size: 20px; margin-top: 10px;">
            <b>Verification Code:</b> <span id="vcode" style="color: green;">{code}</span>
            <button onclick="navigator.clipboard.writeText(document.getElementById('vcode').innerText)" 
                    style="margin-left: 10px; padding: 4px 10px;">Copy</button>
        </div>
        """
    else:
        html_code = f"<div style='color: red; font-size: 18px;'>{code}</div>"

    info_text = f"<b>Sender:</b> {sender}<br><b>Subject:</b> {subject}"
    log_html = f"<pre style='background:black; color:white; padding:10px; max-height:300px; overflow:auto;'>{logs}</pre>"

    return html_code, info_text, log_html


with gr.Blocks() as demo:
    gr.Markdown("# 📨 ProtonMail Verification Code Fetcher")

    with gr.Row():
        btn = gr.Button("🔍 Get Latest Verification Code")

    code_area = gr.HTML()
    info_area = gr.HTML()
    log_area = gr.HTML()

    btn.click(fn=display_verification, outputs=[code_area, info_area, log_area])

demo.launch()