#!/usr/bin/env python3
"""
refine_metadata.py — 将现有 benchmark JSON 的 metadata.steps 精炼为原子级 GUI 操作步骤。
直接覆盖原文件，原始步骤保存到 metadata.steps_original。

用法:
  # 处理 avogadro 的所有任务
  python refine_metadata.py --software avogadro

  # 处理指定任务
  python refine_metadata.py --tasks avogadro/building-organic-molecules_task1.json

  # 全部 7 个软件
  python refine_metadata.py

  # 强制覆盖已精炼过的 (有 steps_original 的)
  python refine_metadata.py --force

  # 预览模式 (不调用 API)
  python refine_metadata.py --dry-run

环境变量:
  OPENAI_API_KEY     — API Key (必需)
  OPENAI_BASE_URL    — API Base URL (默认 https://api.openai.com/v1)
  EXTRACT_MODEL      — 模型名 (默认 gemini-3.1-pro-preview)
"""

import os
import sys
import asyncio
import aiohttp
import base64
import logging
import json
import re
import shutil
import tempfile
import argparse
from pathlib import Path
from typing import List, Optional, Dict, Tuple
from dataclasses import dataclass, field
from datetime import datetime
import io

# ─── Configuration ───────────────────────────────────────────────────────────

SCRIPT_DIR = Path(__file__).parent
PROJECT_ROOT = SCRIPT_DIR.parent

API_BASE_URL = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
API_URL = f"{API_BASE_URL}/chat/completions"
API_KEY = os.getenv("OPENAI_API_KEY")
MODEL_NAME = os.getenv("EXTRACT_MODEL", "gemini-3.1-pro-preview")

MAX_CONCURRENT_REQUESTS = 3
MAX_RETRY_ATTEMPTS = 3
RETRY_DELAY = 5
RETRY_BACKOFF = 2
MAX_IMAGES_PER_REQUEST = 20
MAX_TOKENS = 8192

EXAMPLES_FOLDER = SCRIPT_DIR / "examples"
INPUT_FOLDER = SCRIPT_DIR / "inputs"

# 只处理这 7 个软件
TARGET_SOFTWARE = {"avogadro", "imagej", "jade", "origin", "ovito", "pymol", "vesta"}

# ─── Prompt ──────────────────────────────────────────────────────────────────

SYSTEM_PROMPT = """你是一位桌面科研软件 GUI 自动化专家。你的任务是为一个 **已确定的具体任务** 生成极其详细的、原子级 GUI 操作步骤。

## 背景
这些步骤将作为"教程级操作指南"，直接注入给 AI Agent 去操控桌面软件。Agent 会严格按步骤执行，因此：
- **每一步必须是一个不可分割的原子 GUI 操作**（单击一个按钮、输入一个值、选择一个菜单项等）
- 不能跳步或合并操作
- 必须包含所有隐含的中间操作

## 你将收到
1. **教程文档的截图**（PDF 教程页面图片）— 作为操作上下文参考
2. **任务目标** (task_goal) — 一句话描述要完成的任务
3. **当前粗粒度步骤** (current_steps) — 现有的概要步骤，但粒度太粗

## 你需要输出
根据教程文档内容，将粗粒度步骤细化为**原子级 GUI 操作步骤**。

## 步骤编写规则

### 每步必须满足：
1. **只包含一个明确的、不可分割的 GUI 原子动作**
2. **必须指明**：操作的控件类型（菜单项/按钮/输入框/复选框/下拉菜单/树节点等）
3. **必须指明**：控件的标签文字或位置描述
4. **必须指明**：操作类型（单击/双击/右键/输入文字/勾选/选择等）
5. **包含所有隐含步骤**：不要遗漏任何中间操作

### ❌ 错误（太粗）：
- "搜索 benzene 并插入" — 缺少展开文件夹、选中文件、点击插入等中间步骤
- "设置平滑参数" — 没说具体修改哪个控件、输入什么值
- "在对话框中配置选项" — 没有分解为逐个控件操作

### ✅ 正确（原子级）：
- "在\"筛选\"输入框中，将光标定位到输入框，清空已有内容，输入 benzene"
- "在筛选结果的树形列表中，双击展开名为 aromatics 的文件夹"
- "在展开的列表中，单击选中 benzene.cjson 文件"
- "点击对话框底部的\"插入\"按钮"

## 输出格式
只返回一个 JSON 对象：
{
    "steps": "细化后的操作步骤字符串，每步带编号，用换行符分隔"
}

不要添加任何其他文本。只输出 JSON。"""


# ─── Logging ─────────────────────────────────────────────────────────────────

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger(__name__)


# ─── Stats ───────────────────────────────────────────────────────────────────

@dataclass
class Stats:
    total: int = 0
    completed: int = 0
    failed: int = 0
    skipped: int = 0
    retries: int = 0
    start_time: datetime = field(default_factory=datetime.now)
    failures: List[Tuple[str, str]] = field(default_factory=list)

    def log_progress(self):
        done = self.completed + self.failed + self.skipped
        pct = (done / self.total * 100) if self.total else 0
        logger.info(f"进度: {done}/{self.total} ({pct:.0f}%) | 成功={self.completed} 失败={self.failed} 跳过={self.skipped}")

    def summary(self):
        elapsed = (datetime.now() - self.start_time).total_seconds()
        logger.info("=" * 60)
        logger.info(f"完成! 总计={self.total} 成功={self.completed} 失败={self.failed} 跳过={self.skipped} 重试={self.retries}")
        logger.info(f"耗时: {int(elapsed//60)}m{int(elapsed%60)}s")
        if self.failures:
            logger.info("失败列表:")
            for path, err in self.failures:
                logger.info(f"  {path}: {err[:100]}")


# ─── Document → Images ──────────────────────────────────────────────────────

def convert_pdf_to_images(pdf_path: str) -> List[str]:
    """Convert PDF to base64 JPEG images."""
    try:
        from pdf2image import convert_from_path

        # Quick page count
        quick = convert_from_path(pdf_path, dpi=36, fmt='jpeg')
        total_pages = len(quick)
        del quick

        if total_pages > MAX_IMAGES_PER_REQUEST:
            dpi, quality = 100, 80
            step = total_pages / MAX_IMAGES_PER_REQUEST
            selected = [int(step * i) + 1 for i in range(MAX_IMAGES_PER_REQUEST)]
            logger.info(f"  大 PDF ({total_pages}页): 采样 {len(selected)} 页 @ {dpi} DPI")
            images = []
            for pn in selected:
                imgs = convert_from_path(pdf_path, dpi=dpi, fmt='jpeg', first_page=pn, last_page=pn)
                if imgs:
                    buf = io.BytesIO()
                    imgs[0].save(buf, format='JPEG', quality=quality)
                    images.append(base64.b64encode(buf.getvalue()).decode())
            return images
        else:
            dpi, quality = 150, 90
            imgs = convert_from_path(pdf_path, dpi=dpi, fmt='jpeg')
            result = []
            for img in imgs:
                buf = io.BytesIO()
                img.save(buf, format='JPEG', quality=quality)
                result.append(base64.b64encode(buf.getvalue()).decode())
            return result
    except Exception as e:
        logger.error(f"  PDF转图片失败: {e}")
        return []


# ─── Image Cache ─────────────────────────────────────────────────────────────

_pdf_image_cache: Dict[str, List[str]] = {}

def get_tutorial_images(pdf_path: str) -> List[str]:
    """Get images for a tutorial PDF, with caching to avoid re-converting."""
    key = str(pdf_path)
    if key not in _pdf_image_cache:
        logger.info(f"  正在转换教程 PDF: {Path(pdf_path).name}")
        _pdf_image_cache[key] = convert_pdf_to_images(pdf_path)
        logger.info(f"  得到 {len(_pdf_image_cache[key])} 张图片")
    return _pdf_image_cache[key]


# ─── API Call ────────────────────────────────────────────────────────────────

async def refine_steps_via_api(
    task_goal: str,
    current_steps: str,
    tutorial_images: List[str],
    session: aiohttp.ClientSession,
    stats: Stats
) -> Tuple[str, bool]:
    """Call LLM to refine coarse steps into fine-grained atomic GUI steps."""

    user_text = f"""## 任务目标
{task_goal}

## 当前粗粒度步骤（需要细化）
{current_steps}

请根据以上教程截图内容，将粗粒度步骤细化为原子级 GUI 操作步骤。只输出 JSON。"""

    content = [{"type": "text", "text": user_text}]
    for img_b64 in tutorial_images[:MAX_IMAGES_PER_REQUEST]:
        content.append({
            "type": "image_url",
            "image_url": {"url": f"data:image/jpeg;base64,{img_b64}"}
        })

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": content}
    ]

    for attempt in range(1, MAX_RETRY_ATTEMPTS + 1):
        try:
            headers = {
                "Authorization": f"Bearer {API_KEY}",
                "Content-Type": "application/json"
            }
            payload = {
                "model": MODEL_NAME,
                "messages": messages,
                "max_tokens": MAX_TOKENS,
            }

            async with session.post(API_URL, headers=headers, json=payload, timeout=240) as resp:
                if resp.status == 200:
                    data = await resp.json()
                    return data['choices'][0]['message']['content'], True
                else:
                    err = await resp.text()
                    logger.warning(f"  API错误 ({resp.status}): {err[:200]}")
                    if resp.status == 413:
                        return "Payload too large", False
        except asyncio.TimeoutError:
            logger.warning(f"  API超时 (attempt {attempt})")
        except Exception as e:
            logger.warning(f"  API异常: {e}")

        if attempt < MAX_RETRY_ATTEMPTS:
            delay = RETRY_DELAY * (RETRY_BACKOFF ** (attempt - 1))
            logger.info(f"  重试 {attempt}/{MAX_RETRY_ATTEMPTS} (等待 {delay}s)")
            stats.retries += 1
            await asyncio.sleep(delay)

    return "Max retries exceeded", False


def parse_refined_steps(raw_content: str) -> Optional[str]:
    """Extract the 'steps' string from the LLM's JSON response."""
    # Try to find JSON block in markdown code fence
    m = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', raw_content, re.DOTALL)
    if m:
        json_str = m.group(1)
    else:
        # Try raw JSON
        m = re.search(r'\{.*\}', raw_content, re.DOTALL)
        json_str = m.group(0) if m else raw_content

    try:
        obj = json.loads(json_str)
        steps = obj.get("steps", "")
        if steps and len(steps) > 20:
            return steps
    except json.JSONDecodeError:
        pass

    # Fallback: if content looks like numbered steps directly
    if re.search(r'^\d+\.\s', raw_content, re.MULTILINE):
        lines = [l.strip() for l in raw_content.strip().split('\n') if re.match(r'^\d+\.', l.strip())]
        if len(lines) >= 2:
            return '\n'.join(lines)

    return None


# ─── Task Processing ────────────────────────────────────────────────────────

def find_tutorial_pdf(software: str, task_id: str) -> Optional[Path]:
    """Find the tutorial PDF for a given task.
    Task ID format: {tutorial_stem}_task{N}
    Tutorial PDF:   inputs/{software}/{tutorial_stem}.pdf
    """
    m = re.match(r'^(.+?)_task\d+$', task_id)
    if not m:
        return None
    tutorial_stem = m.group(1)
    pdf_path = INPUT_FOLDER / software / f"{tutorial_stem}.pdf"
    if pdf_path.exists():
        return pdf_path

    # Fallback: try to find any PDF with matching prefix
    sw_input_dir = INPUT_FOLDER / software
    if sw_input_dir.exists():
        for f in sw_input_dir.iterdir():
            if f.suffix.lower() == '.pdf' and f.stem == tutorial_stem:
                return f
    return None


async def process_task(
    task_json_path: Path,
    software: str,
    session: aiohttp.ClientSession,
    semaphore: asyncio.Semaphore,
    stats: Stats,
    force: bool = False,
    dry_run: bool = False
):
    """Process a single task JSON: refine its metadata.steps."""
    async with semaphore:
        rel_path = f"{software}/{task_json_path.name}"

        # Load task JSON
        try:
            with open(task_json_path, 'r', encoding='utf-8') as f:
                task_data = json.load(f)
        except Exception as e:
            stats.failed += 1
            stats.failures.append((rel_path, f"JSON读取失败: {e}"))
            stats.log_progress()
            return

        task_id = task_data.get("id", task_json_path.stem)
        instruction = task_data.get("instruction", "")
        metadata = task_data.get("metadata", {})
        current_steps = metadata.get("steps", "")

        # Check if already refined (has steps_original)
        if metadata.get("steps_original") and not force:
            logger.info(f"跳过 (已精炼): {rel_path}")
            stats.skipped += 1
            stats.log_progress()
            return

        if not instruction:
            logger.warning(f"跳过 (无instruction): {rel_path}")
            stats.skipped += 1
            stats.log_progress()
            return

        # Find tutorial PDF
        tutorial_pdf = find_tutorial_pdf(software, task_id)

        if dry_run:
            logger.info(f"[DRY RUN] 将处理: {rel_path}")
            logger.info(f"  任务: {instruction[:80]}...")
            logger.info(f"  当前步骤: {current_steps[:80]}...")
            logger.info(f"  教程PDF: {tutorial_pdf or '未找到'}")
            stats.completed += 1
            stats.log_progress()
            return

        if not tutorial_pdf:
            logger.warning(f"跳过 (无教程PDF): {rel_path} (task_id={task_id})")
            stats.skipped += 1
            stats.log_progress()
            return

        # Get tutorial images (cached)
        images = get_tutorial_images(str(tutorial_pdf))
        if not images:
            logger.warning(f"跳过 (PDF转图片失败): {rel_path}")
            stats.skipped += 1
            stats.log_progress()
            return

        # Call API to refine steps
        logger.info(f"正在精炼: {rel_path}")
        raw_response, success = await refine_steps_via_api(
            instruction, current_steps, images, session, stats
        )

        if not success:
            stats.failed += 1
            stats.failures.append((rel_path, raw_response[:200]))
            stats.log_progress()
            return

        # Parse refined steps
        refined_steps = parse_refined_steps(raw_response)
        if not refined_steps:
            stats.failed += 1
            stats.failures.append((rel_path, f"无法解析LLM返回: {raw_response[:200]}"))
            stats.log_progress()
            return

        # Save original steps, update with refined steps
        task_data["metadata"]["steps_original"] = current_steps
        task_data["metadata"]["steps"] = refined_steps

        # Write back to original file
        with open(task_json_path, 'w', encoding='utf-8') as f:
            json.dump(task_data, f, ensure_ascii=False, indent=2)

        # Log step count change
        step_count_old = len([l for l in current_steps.split('\n') if l.strip()])
        step_count_new = len([l for l in refined_steps.split('\n') if l.strip()])
        logger.info(f"  ✓ {rel_path}: {step_count_old}步 → {step_count_new}步")

        stats.completed += 1
        stats.log_progress()


# ─── Main ────────────────────────────────────────────────────────────────────

def collect_tasks(software_filter: List[str] = None, task_filter: List[str] = None) -> List[Tuple[Path, str]]:
    """Collect all (task_json_path, software_name) pairs to process."""
    results = []

    if task_filter:
        for t in task_filter:
            parts = t.split('/')
            if len(parts) == 2:
                sw, fname = parts
                path = EXAMPLES_FOLDER / sw / fname
                if path.exists():
                    results.append((path, sw))
                else:
                    logger.warning(f"任务文件不存在: {t}")
            else:
                logger.warning(f"任务格式应为 software/filename.json: {t}")
        return results

    softwares = software_filter if software_filter else sorted(TARGET_SOFTWARE)

    for sw in softwares:
        sw_dir = EXAMPLES_FOLDER / sw
        if not sw_dir.exists():
            logger.warning(f"软件目录不存在: {sw}")
            continue
        for f in sorted(sw_dir.glob("*.json")):
            results.append((f, sw))

    return results


async def main():
    parser = argparse.ArgumentParser(description="Refine benchmark metadata.steps to fine-grained GUI operations")
    parser.add_argument('--software', nargs='*', help='Only process these software names (e.g. avogadro jade)')
    parser.add_argument('--tasks', nargs='*', help='Only process specific tasks (e.g. avogadro/building-organic-molecules_task1.json)')
    parser.add_argument('--force', action='store_true', help='Overwrite already-refined files (those with steps_original)')
    parser.add_argument('--dry-run', action='store_true', help='Preview mode, no API calls')
    args = parser.parse_args()

    if not API_KEY and not args.dry_run:
        logger.error("请设置 OPENAI_API_KEY 环境变量")
        return

    # Collect tasks
    tasks = collect_tasks(args.software, args.tasks)
    if not tasks:
        logger.error("未找到需要处理的任务文件")
        return

    stats = Stats(total=len(tasks))

    logger.info("=" * 60)
    logger.info("Refine Metadata — 细粒度步骤生成")
    logger.info("=" * 60)
    logger.info(f"模型: {MODEL_NAME}")
    logger.info(f"数据目录: {EXAMPLES_FOLDER}")
    logger.info(f"教程目录: {INPUT_FOLDER}")
    logger.info(f"任务数: {len(tasks)}")
    logger.info(f"并发数: {MAX_CONCURRENT_REQUESTS}")
    if args.force:
        logger.info("模式: 强制覆盖")
    if args.dry_run:
        logger.info("模式: DRY RUN (不调用API)")
    logger.info("=" * 60)

    # Show distribution
    from collections import Counter
    dist = Counter(sw for _, sw in tasks)
    for sw, cnt in sorted(dist.items()):
        logger.info(f"  {sw}: {cnt} tasks")

    semaphore = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS)

    async with aiohttp.ClientSession() as session:
        coros = [
            process_task(path, sw, session, semaphore, stats, args.force, args.dry_run)
            for path, sw in tasks
        ]
        await asyncio.gather(*coros)

    stats.summary()

    # Save processing report
    report = {
        "timestamp": datetime.now().isoformat(),
        "model": MODEL_NAME,
        "total": stats.total,
        "completed": stats.completed,
        "failed": stats.failed,
        "skipped": stats.skipped,
        "failures": [{"file": f, "error": e} for f, e in stats.failures]
    }
    report_path = EXAMPLES_FOLDER / "refine_report.json"
    with open(report_path, 'w', encoding='utf-8') as f:
        json.dump(report, f, ensure_ascii=False, indent=2)
    logger.info(f"处理报告: {report_path}")


if __name__ == "__main__":
    asyncio.run(main())
