import os
import streamlit as st
from PIL import Image
import pytesseract
import openai
from openai import OpenAI
import subprocess

# ========== ✅ 设置 GPT Key ==========
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))


# ========== ✅ 页面设置 ==========
st.set_page_config(page_title="AI Cold Email Generator", layout="centered")
st.title("📬 AI Cold Email Generator")
st.write("上传两位 LinkedIn 截图（每人可多张），系统将识别并生成 Coffee Chat 邀请邮件")

# ========== ✅ 函数区 ==========

def extract_text_from_images(uploaded_files):
    """从多张图片中提取文本"""
    all_text = ""
    for uploaded_file in uploaded_files:
        image = Image.open(uploaded_file).convert("RGB")
        text = pytesseract.image_to_string(image)
        all_text += text + "\n"
    return all_text.strip()

def generate_email(profile1_text, profile2_text):
    """调用 GPT 生成邮件"""
    prompt = f"""
    根据以下两人的 LinkedIn Profile，提取共同点，并生成一封自然、有温度的 Coffee Chat 邀请邮件（避免模板化）：

    📌 Profile 1:
    {profile1_text}

    📌 Profile 2:
    {profile2_text}

    ✉️ 邮件内容如下：
    """
    response = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": "你是一位 AI Networking 邮件助手，擅长生成 Coffee Chat 邮件"},
            {"role": "user", "content": prompt}
        ],
        temperature=0.7,
        max_tokens=600
    )
    return response.choices[0].message.content.strip()

def check_tesseract():
    """调试用：查看 OCR 是否安装成功"""
    try:
        version = subprocess.check_output(["tesseract", "--version"])
        return version.decode()
    except Exception as e:
        return str(e)

# ========== ✅ 页面上传区 ==========

col1, col2 = st.columns(2)
with col1:
    files1 = st.file_uploader("上传你的 LinkedIn 截图（支持多张）", type=["png", "jpg", "jpeg"], accept_multiple_files=True, key="profile1")
with col2:
    files2 = st.file_uploader("上传目标联系人的截图（支持多张）", type=["png", "jpg", "jpeg"], accept_multiple_files=True, key="profile2")

if st.button("🚀 生成邮件"):
    if files1 and files2:
        with st.spinner("正在识别图像并生成邮件..."):
            profile1_text = extract_text_from_images(files1)
            profile2_text = extract_text_from_images(files2)
            email = generate_email(profile1_text, profile2_text)
            st.success("✅ 邮件生成成功")
            st.text_area("📨 Coffee Chat 邮件内容", email, height=300)
    else:
        st.warning("请上传双方的截图")

# ========== ✅ OCR 安装测试区 ==========
with st.expander("🛠️ 查看 Tesseract 安装状态（调试用）"):
    result = check_tesseract()
    if "tesseract" in result.lower():
        st.success("✅ Tesseract 安装成功")
        st.code(result, language="bash")
    else:
        st.error("❌ Tesseract 未安装")
        st.text(result)