import base64
import json
import os
os.system("pip install openai")
from openai import OpenAI
from openai import APIError, AuthenticationError, NotFoundError

def image_to_base64(image_path):
    """将图像文件转换为base64编码字符串"""
    try:
        with open(image_path, 'rb') as image_file:
            image_data = image_file.read()
        return base64.b64encode(image_data).decode('utf-8')
    except FileNotFoundError:
        print(f"错误：未找到图像文件 '{image_path}'")
        return None
    except Exception as e:
        print(f"转换图像时出错：{str(e)}")
        return None

def clean_json_response(raw_response):
    """清理模型返回的响应，移除 ```json 和 ``` 代码块标记"""
    # 移除开头的 ```json（不区分大小写）
    if raw_response.startswith('```json') or raw_response.startswith('```JSON'):
        raw_response = raw_response[7:]
    # 移除结尾的 ```
    if raw_response.endswith('```'):
        raw_response = raw_response[:-3]
    # 去除首尾多余空格/换行
    return raw_response.strip()

def parse_analysis_result(result):
    """解析分析结果并以结构化方式展示"""
    if not result:
        return
    
    print("\n" + "="*50)
    print("📊 图像分析结果解析")
    print("="*50)
    
    # 解析页面评分
    print("\n【页面评分】")
    scoring = result.get("页面评分", {})
    for key, value in scoring.items():
        if key.endswith("_comment"):
            # 处理评论（对应前面的评分项）
            score_key = key.replace("_comment", "")
            print(f"  - {score_key}: {scoring.get(score_key, 'N/A')}分 - {value}")
    
    # 解析智能体能力
    print("\n【智能体能力拆解】")
    capabilities = result.get("智能体能力拆解", {})
    
    # 核心功能
    print("\n核心功能:")
    for i, func in enumerate(capabilities.get("core_functions", []), 1):
        print(f"  {i}. {func}")
    
    # 优势
    print("\n优势:")
    for i, strength in enumerate(capabilities.get("strengths", []), 1):
        print(f"  {i}. {strength}")
    
    # 劣势
    print("\n劣势:")
    for i, weakness in enumerate(capabilities.get("weaknesses", []), 1):
        print(f"  {i}. {weakness}")
    
    # 潜在用途
    print("\n潜在用途:")
    for i, use in enumerate(capabilities.get("potential_uses", []), 1):
        print(f"  {i}. {use}")
    
    # 改进方向
    print("\n改进方向:")
    for i, area in enumerate(capabilities.get("improvement_areas", []), 1):
        print(f"  {i}. {area}")
    
    # 详细分析
    print("\n【详细分析】")
    print(capabilities.get("detailed_analysis", "无详细分析信息"))
    
    print("\n" + "="*50)

def analyze_image_with_ernie(api_key, image_path):
    """使用百度ERNIE-VL模型分析图像并返回结果"""
    # 转换图像为base64
    base64_string = image_to_base64(image_path)
    if not base64_string:
        return None

    try:
        # 初始化客户端
        client = OpenAI(
            base_url='https://qianfan.baidubce.com/v2',
            api_key=api_key
        )

        # 构建消息
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "请分析这张智能体网页的截图，并完成以下任务：\n\
                        1. 页面评分（每项1-10分，并给出简短评论）：\n\
                           - 整体评价（overall）\n\
                           - 设计美感（design）\n\
                           - 易用性（usability）\n\
                           - 功能完整性（functionality）\n\
                           - 响应式设计（responsiveness）\n\
                        2. 智能体能力拆解：\n\
                           - 核心功能（core_functions）：列出3-5个主要功能\n\
                           - 优势（strengths）：列出3-4个主要优势\n\
                           - 劣势（weaknesses）：列出3-4个潜在劣势\n\
                           - 潜在用途（potential_uses）：列出3-4个可能的应用场景\n\
                           - 改进方向（improvement_areas）：列出2-3个可以改进的地方\n\
                           - 详细分析（detailed_analysis）：一段详细的综合分析\n\
                        【输出要求】：仅返回纯JSON字符串，不要包含任何代码块标记（如```json、```），确保JSON结构可直接被Python的json.loads()解析。"
                    },
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/jpeg;base64,{base64_string}"}
                    }
                ]
            }
        ]

        # 发送请求
        print("正在分析图像，请稍候...")
        response = client.chat.completions.create(
            model="ernie-4.5-turbo-vl", 
            messages=messages, 
            temperature=0.8, 
            top_p=0.2,
            extra_body={"penalty_score": 1}
        )

        # 处理响应
        raw_content = response.choices[0].message.content
        cleaned_content = clean_json_response(raw_content)

        try:
            result = json.loads(cleaned_content)
            return result
        except json.JSONDecodeError as e:
            print(f"错误：清理后的内容仍不是有效JSON格式，错误详情：{str(e)}")
            print(f"清理后内容：{cleaned_content}")
            return None

    except AuthenticationError:
        print("错误：API密钥验证失败，请检查您的百度智能云API密钥是否正确")
        return None
    except NotFoundError:
        print("错误：未找到指定的模型，请确认模型名称是否正确")
        return None
    except APIError as e:
        print(f"API请求失败：{str(e)}")
        return None
    except Exception as e:
        print(f"发生未知错误：{str(e)}")
        return None

def save_analysis_result(result, output_file="analysis_result.json"):
    """保存分析结果到JSON文件"""
    if result:
        try:
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(result, f, ensure_ascii=False, indent=2)
            print(f"\n分析结果已保存到 {output_file}")
        except Exception as e:
            print(f"保存结果时出错：{str(e)}")

if __name__ == "__main__":
    # 配置
    API_KEY = os.getenv('baidu_api_key')  # 替换为实际API密钥
    IMAGE_PATH = "demo.png"  # 替换为实际图像路径
    OUTPUT_FILE = "analysis_result.json"  # 结果保存文件名

    # 检查API密钥
    if API_KEY in ["您的百度智能云API密钥", "请替换为您的百度API密钥"]:
        print("⚠️ 请先替换代码中的API_KEY为您自己的百度智能云API密钥！")
        print("获取地址：https://console.bce.baidu.com/iam/#/iam/apikey/list")
    else:
        # 执行分析
        analysis_result = analyze_image_with_ernie(API_KEY, IMAGE_PATH)
        
        # 解析并展示结果
        if analysis_result:
            # 打印原始JSON（便于调试）
            print("\n📄 原始分析结果(JSON):")
            print(json.dumps(analysis_result, ensure_ascii=False, indent=2))
            
            # 解析并格式化展示
            parse_analysis_result(analysis_result)
            
            # 保存结果到文件
            save_analysis_result(analysis_result, OUTPUT_FILE)