import gradio as gr
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException
import time
import random
import string
import logging
import re

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def setup_driver():
    """设置 Chrome 驱动"""
    chrome_options = Options()
    chrome_options.add_argument('--no-sandbox')
    chrome_options.add_argument('--disable-dev-shm-usage')
    chrome_options.add_argument('--disable-gpu')
    chrome_options.add_argument('--disable-software-rasterizer')
    chrome_options.add_argument('--headless')
    chrome_options.add_argument('--disable-blink-features=AutomationControlled')
    chrome_options.add_experimental_option('excludeSwitches', ['enable-automation'])
    chrome_options.add_experimental_option('useAutomationExtension', False)
    
    driver = webdriver.Chrome(options=chrome_options)
    return driver

def get_verification_code(driver, email, max_retries=10, delay=2):
    """获取验证码"""
    email_url = f'https://linshiyou.com/#/{email}'
    
    try:
        driver.get(email_url)
        
        for attempt in range(max_retries):
            try:
                driver.refresh()
                time.sleep(delay)
                
                body_content = driver.find_element(By.CLASS_NAME, "body").get_attribute('outerHTML')
                logger.info(f"获取到邮件内容: {body_content}")
                
                match = re.search(r'\b\d{6}\b', body_content)
                if match:
                    verification_code = match.group(0)
                    logger.info(f"获取到验证码: {verification_code}")
                    return verification_code
                    
            except Exception as e:
                logger.error(f"第 {attempt + 1} 次获取验证码失败: {str(e)}")
                
            logger.info(f"等待验证码，重试第 {attempt + 1} 次")
            time.sleep(delay)
            
    except Exception as e:
        logger.error(f"打开邮箱页面失败: {str(e)}")
        
    return None

def generate_random_email():
    """生成随机邮箱地址"""
    try:
        random_prefix = ''.join(random.choices(string.ascii_letters + string.digits, k=8))
        email = f"{random_prefix}@youxiang.dev"
        logger.info(f"生成随机邮箱: {email}")
        return {"status": "success", "email": email}
    except Exception as e:
        logger.error(f"生成邮箱失败: {str(e)}")
        return {"status": "error", "message": str(e)}

def get_email_code(email: str):
    """获取指定邮箱的验证码"""
    if not email:
        return {"status": "error", "message": "邮箱地址不能为空"}
        
    logger.info(f"开始获取邮箱验证码: {email}")
    driver = None
    
    try:
        driver = setup_driver()
        code = get_verification_code(driver, email)
        
        if code:
            logger.info(f"成功获取验证码: {code}")
            return {"status": "success", "code": code}
        else:
            logger.error("未能获取验证码")
            return {"status": "error", "message": "未能获取验证码"}
            
    except Exception as e:
        logger.error(f"获取验证码过程出错: {str(e)}")
        return {"status": "error", "message": str(e)}
    finally:
        if driver:
            try:
                driver.quit()
            except:
                pass

# 创建 Gradio 界面
with gr.Blocks(title="邮箱验证码服务") as demo:
    gr.Markdown("""
    # 邮箱验证码服务
    
    提供两个主要功能：
    1. 生成随机临时邮箱
    2. 获取邮箱验证码
    """)
    
    with gr.Tab("生成随机邮箱"):
        with gr.Row():
            email_btn = gr.Button("生成邮箱", variant="primary")
            email_output = gr.JSON(label="结果")
    
    with gr.Tab("获取验证码"):
        with gr.Row():
            email_input = gr.Textbox(
                label="邮箱地址",
                placeholder="请输入要获取验证码的邮箱地址",
                info="支持 @youxiang.dev 域名的邮箱"
            )
        with gr.Row():
            code_btn = gr.Button("获取验证码", variant="primary")
            code_output = gr.JSON(label="结果")
            
    # 绑定事件处理
    email_btn.click(
        fn=generate_random_email,
        inputs=[],
        outputs=email_output,
        api_name="generate_email"
    )
    
    code_btn.click(
        fn=get_email_code,
        inputs=email_input,
        outputs=code_output,
        api_name="get_code"
    )

# 启用队列处理请求
demo.queue()

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)