#!/usr/bin/env python3
"""
SVG Review Script - 将 SVG 转为 PNG 并调用多模态大模型评价
"""

import argparse
import json
import os
import subprocess
import sys
from pathlib import Path

# 配置
DEFAULT_MODEL = "MiniMax-M2.7-BF16"
API_URL = "http://llm.bnuzh.edu.cn:9180/v1/chat/completions"
API_KEY = "sk-AeGYLnhYWUEifWq7fedjCioDZssxEVraGzEcyhej9gQ1um6s"


def svg_to_png(svg_path: str, png_path: str, width: int = 1200) -> bool:
    """将 SVG 转换为 PNG"""
    try:
        # 使用 ImageMagick convert
        cmd = ["convert", "-background", "white", "-density", "300", "-resize", f"{width}", svg_path, png_path]
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
        if result.returncode != 0:
            print(f"[错误] SVG转PNG失败: {result.stderr}", flush=True)
            return False
        return True
    except Exception as e:
        print(f"[错误] SVG转PNG异常: {e}", flush=True)
        return False


def read_image_base64(image_path: str) -> str:
    """读取图片并转为 base64"""
    import base64
    try:
        with open(image_path, 'rb') as f:
            return base64.b64encode(f.read()).decode('utf-8')
    except Exception as e:
        print(f"[错误] 读取图片失败: {e}", flush=True)
        return None


def call_multimodal_model(image_base64: str, desc_content: str, model: str = DEFAULT_MODEL) -> str:
    """调用多模态大模型评价图片"""
    import urllib.request
    import urllib.parse
    
    # 构建 prompt
    prompt = f"""Please analyze the following academic diagram and provide evaluation and improvement suggestions based on the following criteria:

Evaluation Criteria:
1. Route arrows through gaps between boxes, not through box interiors
2. Move arrow labels 6-8px away from the arrow line (offset-first); add background rects only when offset is insufficient
3. Widen inter-row/inter-column gutters so same-layer arrows have clear corridors
4. Collapse repeated cross-layer arrows into a single "delegates down" rail outside the content area
5. Move legend/notes out of any region where arrows or labels land
6. Increase viewBox height/width rather than packing elements tighter
7. If a filtered element (drop-shadow, blur) is missing one side of its border, move it ≥30px away from that viewBox edge, or remove the filter and rely on color/contrast for visual separation

Please return evaluation results in JSON format with the following fields:
- overall_score: 0-10 overall score
- issues: list of issues found, each containing "criterion" (criterion number), "description" (issue description), "suggestion" (improvement suggestion)
- summary: overall evaluation summary

If unable to evaluate (e.g., image quality issues), return empty list.

Diagram description: {desc_content}
"""

    # 构建请求
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/png;base64,{image_base64}"
                    }
                },
                {
                    "type": "text",
                    "text": prompt
                }
            ]
        }
    ]
    
    data = {
        "model": model,
        "messages": messages,
        "temperature": 0.3,
        "max_tokens": 2000
    }
    
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }
    
    try:
        req = urllib.request.Request(
            API_URL,
            data=json.dumps(data).encode('utf-8'),
            headers=headers,
            method='POST'
        )
        with urllib.request.urlopen(req, timeout=120) as response:
            result = json.loads(response.read().decode('utf-8'))
            return result['choices'][0]['message']['content']
    except Exception as e:
        print(f"[错误] 调用大模型失败: {e}", flush=True)
        return None


def parse_review_response(response: str) -> dict:
    """解析大模型返回的评价"""
    import re
    
    try:
        # 尝试提取 JSON
        # 先去除 markdown 代码块
        response = re.sub(r'^```json\s*', '', response)
        response = re.sub(r'^```\s*', '', response)
        response = re.sub(r'\s*```$', '', response)
        response = re.sub(r'^```', '', response)
        response = re.sub(r'```$', '', response)
        
        # 尝试找到 JSON 对象（从第一个 { 到最后一个 }）
        json_match = re.search(r'\{[\s\S]*\}', response)
        if json_match:
            return json.loads(json_match.group())
        
        # 如果没有找到 JSON，返回原始响应
        return {"raw_response": response, "error": "无法解析JSON"}
    except json.JSONDecodeError as e:
        # 如果 JSON 解析失败，尝试提取部分信息
        result = {"error": str(e)}
        
        # 尝试提取 score
        score_match = re.search(r'"overall_score"\s*:\s*(\d+)', response)
        if score_match:
            result['overall_score'] = int(score_match.group(1))
        
        # 尝试提取 summary
        summary_match = re.search(r'"summary"\s*:\s*"([^"]+)"', response)
        if summary_match:
            result['summary'] = summary_match.group(1)
        
        result['raw_response'] = response
        return result
    except Exception as e:
        return {"raw_response": response, "error": str(e)}


def main():
    parser = argparse.ArgumentParser(description="SVG Review - 评价生成的SVG图片")
    parser.add_argument("-s", "--svg", required=True, help="SVG文件路径")
    parser.add_argument("-d", "--desc", required=True, help="描述JSON文件路径")
    parser.add_argument("-o", "--output", required=True, help="输出评价结果JSON文件路径")
    parser.add_argument("-w", "--width", type=int, default=1200, help="PNG宽度")
    parser.add_argument("-m", "--model", default=DEFAULT_MODEL, help="多模态模型名称")
    args = parser.parse_args()
    
    # 检查文件存在
    if not os.path.exists(args.svg):
        print(f"[错误] SVG文件不存在: {args.svg}", flush=True)
        sys.exit(1)
    
    if not os.path.exists(args.desc):
        print(f"[错误] 描述文件不存在: {args.desc}", flush=True)
        sys.exit(1)
    
    # 读取描述文件
    try:
        with open(args.desc, 'r', encoding='utf-8') as f:
            desc_data = json.load(f)
        
        # 提取关键信息（兼容两种格式：直接数据 或 target_description 包装）
        if 'target_description' in desc_data:
            desc_content = json.dumps(desc_data.get('target_description', {}), ensure_ascii=False)
        else:
            desc_content = json.dumps(desc_data, ensure_ascii=False)
    except Exception as e:
        print(f"[错误] 读取描述文件失败: {e}", flush=True)
        sys.exit(1)
    
    # 生成临时 PNG 路径
    svg_path = Path(args.svg)
    png_path = svg_path.with_suffix('.png')
    
    # SVG 转 PNG
    print(f"[信息] 转换SVG到PNG: {svg_path} -> {png_path}", flush=True)
    if not svg_to_png(str(svg_path), str(png_path), args.width):
        print("[错误] SVG转PNG失败，跳过review", flush=True)
        sys.exit(1)
    
    # 读取图片 base64
    print("[信息] 读取图片...", flush=True)
    image_base64 = read_image_base64(str(png_path))
    if not image_base64:
        print("[错误] 读取图片失败，跳过review", flush=True)
        sys.exit(1)
    
    # 调用多模态模型
    print(f"[信息] 调用多模态模型: {args.model}", flush=True)
    response = call_multimodal_model(image_base64, desc_content, args.model)
    if not response:
        print("[错误] 调用大模型失败，跳过review", flush=True)
        sys.exit(1)
    
    # 解析响应
    review_result = parse_review_response(response)
    review_result['source_svg'] = str(svg_path.absolute())
    review_result['source_desc'] = str(Path(args.desc).absolute())
    review_result['model'] = args.model
    
    # 保存结果
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(review_result, f, indent=2, ensure_ascii=False)
    
    print(f"[完成] 评价结果已保存到: {output_path}", flush=True)
    
    # 打印摘要
    if 'overall_score' in review_result:
        print(f"[评分] 整体评分: {review_result['overall_score']}/10", flush=True)
    if 'issues' in review_result and review_result['issues']:
        print(f"[问题] 发现 {len(review_result['issues'])} 个问题", flush=True)


if __name__ == "__main__":
    main()