import os
import requests
import base64
import json
import time
from pathlib import Path

def fix_old_photo(image_path, output_path="fixed_photo.png"):
    """
    使用Gemini API修复老照片
    
    Args:
        image_path: 输入图片路径
        output_path: 输出图片路径
    """
    # 从环境变量获取API密钥
    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        raise ValueError("未找到GEMINI_API_KEY环境变量")
    
    # 读取图片文件并转换为base64
    try:
        with open(image_path, "rb") as image_file:
            image_data = base64.b64encode(image_file.read()).decode('utf-8')
    except FileNotFoundError:
        raise FileNotFoundError(f"找不到图片文件: {image_path}")
    
    # 获取图片的MIME类型
    file_extension = Path(image_path).suffix.lower()
    mime_type_map = {
        '.jpg': 'image/jpeg',
        '.jpeg': 'image/jpeg',
        '.png': 'image/png',
        '.gif': 'image/gif',
        '.webp': 'image/webp'
    }
    mime_type = mime_type_map.get(file_extension, 'image/jpeg')
    
    # 修复提示词
    prompt = """修复这张受损的老照片，去除所有折痕、裂痕、污渍和划痕，补全缺失的细节，提升清晰度，让画面自然完整。人物的面貌务必保持不变，在修复的基础上进行彩色化处理：人物的肤色柔和自然，毛衣呈现温暖的深色调，肤色健康自然，头发为黑色，衣服保持复古质感；去掉人物手里的烟；背景调整为教室，后面是黑板，前面是课桌"""
    
    # 构建请求数据
    request_data = {
        "contents": [{
            "parts": [
                {"text": prompt},
                {
                    "inline_data": {
                        "mime_type": mime_type,
                        "data": image_data
                    }
                }
            ]
        }],
        "generationConfig": {
            "temperature": 0.7,
            "candidateCount": 1,
            "maxOutputTokens": 2048,
        }
    }
    
    # 尝试不同的模型，避免图像生成模型（它们有更严格的配额限制）
    models = [
        "gemini-1.5-flash-latest",
        "gemini-1.5-flash", 
        "gemini-1.5-pro-latest",
        "gemini-1.5-pro"
    ]
    
    headers = {
        "x-goog-api-key": api_key,
        "Content-Type": "application/json"
    }
    
    for model in models:
        url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent"
        
        try:
            print(f"正在尝试模型: {model}")
            response = requests.post(url, headers=headers, json=request_data, timeout=120)
            
            # 处理429错误（配额超限）
            if response.status_code == 429:
                error_data = response.json()
                print(f"模型 {model} 配额超限，错误信息:")
                print(json.dumps(error_data, indent=2, ensure_ascii=False))
                
                # 检查是否有重试建议
                if "error" in error_data and "details" in error_data["error"]:
                    for detail in error_data["error"]["details"]:
                        if detail.get("@type") == "type.googleapis.com/google.rpc.RetryInfo":
                            retry_delay = detail.get("retryDelay", "30s")
                            print(f"建议等待 {retry_delay} 后重试")
                
                continue
            
            # 处理其他HTTP错误
            if response.status_code != 200:
                print(f"HTTP状态码: {response.status_code}")
                print(f"响应内容: {response.text}")
                continue
            
            # 解析响应
            result = response.json()
            print("API响应成功!")
            
            # 提取文本内容（分析结果）
            if "candidates" in result and len(result["candidates"]) > 0:
                candidate = result["candidates"][0]
                if "content" in candidate and "parts" in candidate["content"]:
                    for part in candidate["content"]["parts"]:
                        if "text" in part:
                            analysis_text = part["text"]
                            print("\n=== 照片分析结果 ===")
                            print(analysis_text)
                            
                            # 保存分析结果到文件
                            analysis_file = output_path.replace('.png', '_analysis.txt')
                            with open(analysis_file, 'w', encoding='utf-8') as f:
                                f.write(analysis_text)
                            print(f"\n分析结果已保存到: {analysis_file}")
                            return True
            
            print("未找到有效的分析内容")
            return False
            
        except requests.exceptions.RequestException as e:
            print(f"模型 {model} 请求错误: {e}")
            continue
        except json.JSONDecodeError as e:
            print(f"模型 {model} JSON解析错误: {e}")
            continue
        except Exception as e:
            print(f"模型 {model} 未知错误: {e}")
            continue
    
    print("所有模型都尝试失败")
    return False

def check_api_status():
    """
    检查API状态和配额
    """
    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        print("未找到GEMINI_API_KEY环境变量")
        return False
    
    # 简单的测试请求
    test_data = {
        "contents": [{
            "parts": [{"text": "Hello, how are you?"}]
        }]
    }
    
    url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent"
    headers = {
        "x-goog-api-key": api_key,
        "Content-Type": "application/json"
    }
    
    try:
        print("检查API状态...")
        response = requests.post(url, headers=headers, json=test_data, timeout=30)
        
        if response.status_code == 200:
            print("✅ API工作正常")
            return True
        elif response.status_code == 429:
            print("⚠️ API配额超限")
            error_data = response.json()
            if "error" in error_data:
                print(f"错误信息: {error_data['error']['message']}")
            return False
        else:
            print(f"❌ API错误，状态码: {response.status_code}")
            print(f"响应: {response.text}")
            return False
            
    except Exception as e:
        print(f"❌ API测试失败: {e}")
        return False

def main():
    """主函数"""
    print("=== Gemini 照片分析工具 ===\n")
    
    # 检查API状态
    if not check_api_status():
        print("\n由于API配额限制，当前无法使用图像生成功能。")
        print("建议:")
        print("1. 等待一段时间后重试（通常几分钟到几小时）")
        print("2. 升级到付费版本获得更高的配额")
        print("3. 使用其他AI图像修复工具")
        return
    
    # 输入图片路径
    input_image = "./工大照片-1.jpg"
    output_analysis = "photo_analysis.txt"
    
    # 检查输入文件是否存在
    if not os.path.exists(input_image):
        print(f"错误: 找不到输入图片文件 {input_image}")
        print("请将您的老照片放在当前目录下，或修改input_image变量指向正确的路径")
        return
    
    print(f"正在分析图片: {input_image}")
    
    # 分析照片
    success = fix_old_photo(input_image, output_analysis)
    
    if success:
        print("\n✅ 照片分析完成!")
        print("由于API配额限制，当前只能提供分析报告。")
        print("您可以根据分析结果使用其他工具进行照片修复。")
    else:
        print("\n❌ 照片分析失败")

if __name__ == "__main__":
    main()