#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
简化版DeepSeek OCR测试
"""

import base64
from pathlib import Path
from openai import OpenAI

# API配置
CLARIFAI_PAT = "941fba50c8c04be590a9b2d21b7d8347"

# 初始化OpenAI客户端
client = OpenAI(
    base_url="https://api.clarifai.com/v2/ext/openai/v1",
    api_key=CLARIFAI_PAT
)

def test_basic_connection():
    print("🔍 测试API连接...")
    try:
        response = client.chat.completions.create(
            model="https://clarifai.com/deepseek-ai/deepseek-ocr/models/DeepSeek-OCR",
            messages=[{"role": "user", "content": "Hello, 你好"}],
            max_tokens=50
        )
        print("✅ 连接成功！")
        print(f"响应: {response.choices[0].message.content}")
        return True
    except Exception as e:
        print(f"❌ 连接失败: {e}")
        return False

def ocr_image(image_path):
    print(f"📸 处理图片: {image_path}")
    
    if not Path(image_path).exists():
        print(f"❌ 文件不存在: {image_path}")
        return None
    
    try:
        print("🔄 读取图片...")
        image_base64 = base64.b64encode(Path(image_path).read_bytes()).decode()
        print(f"📊 图片编码完成，长度: {len(image_base64)}")
        
        print("🤖 调用DeepSeek-OCR...")
        response = client.chat.completions.create(
            model="https://clarifai.com/deepseek-ai/deepseek-ocr/models/DeepSeek-OCR",
            messages=[{
                "role": "user",
                "content": [
                    {"type": "text", "text": "请提取图片中的所有中文和英文文字内容"},
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/png;base64,{image_base64}"}
                    }
                ]
            }],
            temperature=0.0,
            max_tokens=1024
        )
        
        result = response.choices[0].message.content
        print("✅ OCR完成！")
        return result
        
    except Exception as e:
        print(f"❌ OCR失败: {e}")
        return None

def main():
    print("🚀 DeepSeek OCR 简化测试")
    print("=" * 40)
    
    if not test_basic_connection():
        return
    
    print("\n" + "=" * 40)
    
    image_path = "../user_upload/截屏2025-10-24 12.34.04.png"
    result = ocr_image(image_path)
    
    if result:
        print("\n🎯 OCR结果:")
        print("-" * 30)
        print(result)
        print("-" * 30)
        
        with open("ocr_result.txt", "w", encoding="utf-8") as f:
            f.write(result)
        print("💾 结果已保存到 ocr_result.txt")
    
    print("\n✅ 测试完成！")

if __name__ == "__main__":
    main()