import os
import gradio as gr
import spaces
from openai import OpenAI
from PIL import Image
import torch
import base64
from io import BytesIO

# 初始化GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
zero = torch.Tensor([0]).to(device)
print(f"Using device: {zero.device}")

# 初始化OpenAI客户端
try:
    api_key = os.environ.get('OPENAI_API_KEY')
    if not api_key:
        raise ValueError("No OpenAI API key found in environment variables")
        
    client = OpenAI(
        api_key=api_key,
        # 可以添加额外的配置，但不是必需的
        timeout=60.0,  # 设置超时时间
        max_retries=2  # 设置重试次数
    )
    print("Successfully initialized OpenAI client")
    
    # 验证API密钥是否有效
    test_response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": "test"}],
        max_tokens=5
    )
    print("API key verified successfully")
    
except ValueError as ve:
    print(f"Error: {str(ve)}")
    client = None
except Exception as e:
    print(f"Error initializing OpenAI client: {str(e)}")
    if "api_key" in str(e).lower():
        print("API key validation failed")
    client = None

def encode_image_to_base64(image):
    """将图片转换为base64编码"""
    if isinstance(image, str):  # 如果是文件路径
        with Image.open(image) as img:
            buffered = BytesIO()
            img.save(buffered, format="PNG")
            img_str = base64.b64encode(buffered.getvalue()).decode()
    else:  # 如果是PIL Image对象
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
    return img_str

@spaces.GPU
def analyze_slide(image, course_title):
    """分析幻灯片内容并生成讲解"""
    print(f"Current device: {zero.device}")  # 验证GPU使用情况
    
    if not client:
        return f"错误：OpenAI API密钥未正确配置或验证失败。\nAPI密钥: {'已设置' if api_key else '未设置'}\n请在Hugging Face Space的Secrets中设置正确的OPENAI_API_KEY。"
    
    if not image:
        return "请上传幻灯片图片"
    
    # 将图片转换为base64
    base64_image = encode_image_to_base64(image)
    
    try:
        # 调用GPT-4 Vision API
        response = client.chat.completions.create(
            model="gpt-4-vision-preview",  # 使用支持图像的模型
            messages=[
                {
                    "role": "system",
                    "content": "你是一位专业的教育讲师，请基于上传的课程幻灯片图片内容，生成详细的讲解。讲解应该清晰、专业且易于理解。"
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": f"这是一节'{course_title}'课程的幻灯片。请详细讲解这张幻灯片的内容，包括主要概念、重点和难点。"
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{base64_image}"
                            }
                        }
                    ]
                }
            ],
            max_tokens=2000
        )
        
        return response.choices[0].message.content
    
    except Exception as e:
        error_message = str(e)
        if "api_key" in error_message.lower():
            return "错误：OpenAI API密钥无效或未正确配置。请检查API密钥设置。"
        return f"处理过程中出现错误: {error_message}"

# 创建Gradio界面
demo = gr.Interface(
    fn=analyze_slide,
    inputs=[
        gr.Image(type="pil", label="上传幻灯片图片"),
        gr.Textbox(label="课程标题", placeholder="请输入课程标题")
    ],
    outputs=gr.Textbox(label="讲解内容"),
    title="课程幻灯片智能讲解系统",
    description="上传课程幻灯片图片并输入课程标题，系统将生成详细的讲解内容。\n注意：请确保在Hugging Face Space的Secrets中设置了有效的OPENAI_API_KEY。",
    examples=[],
    cache_examples=False
)

# 启动应用
if __name__ == "__main__":
    demo.launch()