"""
轨迹数据后处理
将录制的原始Host坐标转换为VM内坐标
"""
import json
import os
import sys
import argparse
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def process_trajectory(task_id, project_root=".", force=False, 
                       scale_x_adjust=1.0, scale_y_adjust=1.0,
                       offset_x=0, offset_y=0):
    """
    处理轨迹数据：将Host坐标转换为VM坐标
    
    Args:
        task_id: 任务ID
        project_root: 项目根目录
        force: 是否强制覆盖已有的处理结果
        scale_x_adjust: X轴缩放调整系数
        scale_y_adjust: Y轴缩放调整系数
        offset_x: X轴偏移调整
        offset_y: Y轴偏移调整
    """
    # 路径
    task_dir = os.path.join(project_root, "tasks", task_id)
    human_demo_dir = os.path.join(task_dir, "human_demo")
    raw_path = os.path.join(human_demo_dir, "actions_raw.json")
    processed_path = os.path.join(human_demo_dir, "actions.json")
    
    # 检查文件
    if not os.path.exists(raw_path):
        logger.error(f"❌ 原始轨迹文件不存在: {raw_path}")
        logger.info("   请先运行: python scripts/collect_task.py <task_id> --mode record")
        return False
    
    if os.path.exists(processed_path) and not force:
        logger.warning(f"⚠️  处理后的文件已存在: {processed_path}")
        logger.info("   使用 --force 参数强制覆盖")
        return False
    
    # 读取原始数据
    logger.info(f"读取原始轨迹: {raw_path}")
    with open(raw_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    metadata = data['metadata']
    actions = data['actions']
    
    logger.info(f"任务ID: {metadata['task_id']}")
    logger.info(f"动作数: {len(actions)}")
    
    # 获取分辨率信息
    if 'vm_resolution' in metadata and 'vm_screenshot_resolution' in metadata:
        vm_w, vm_h = metadata['vm_resolution']
        screenshot_w, screenshot_h = metadata['vm_screenshot_resolution']
        
        # 计算缩放比例
        # 注意：Host端的点击坐标对应截图坐标，需要转换为VM内实际坐标
        scale_x = (vm_w / screenshot_w) * scale_x_adjust
        scale_y = (vm_h / screenshot_h) * scale_y_adjust
        
        logger.info(f"VM分辨率: {vm_w}x{vm_h}")
        logger.info(f"截图分辨率: {screenshot_w}x{screenshot_h}")
        logger.info(f"转换比例: X={scale_x:.3f}, Y={scale_y:.3f}")
        
        if scale_x_adjust != 1.0 or scale_y_adjust != 1.0:
            logger.info(f"应用调整系数: X={scale_x_adjust}, Y={scale_y_adjust}")
        if offset_x != 0 or offset_y != 0:
            logger.info(f"应用偏移调整: X={offset_x}, Y={offset_y}")
    else:
        logger.warning("⚠️  元数据缺少分辨率信息，使用默认比例1.0")
        scale_x = 1.0 * scale_x_adjust
        scale_y = 1.0 * scale_y_adjust
    
    # 转换坐标
    converted_count = 0
    for action in actions:
        if 'pos_host' in action and action['pos_host']:
            host_x, host_y = action['pos_host']
            
            # 应用转换
            vm_x = int(host_x * scale_x + offset_x)
            vm_y = int(host_y * scale_y + offset_y)
            
            action['pos_vm'] = [vm_x, vm_y]
            converted_count += 1
    
    logger.info(f"✅ 坐标转换完成: {converted_count}/{len(actions)} 个动作")
    
    # 添加处理信息到元数据
    metadata['processed'] = {
        "processed_at": __import__('datetime').datetime.now().isoformat(),
        "scale_x": scale_x,
        "scale_y": scale_y,
        "offset_x": offset_x,
        "offset_y": offset_y,
        "converted_actions": converted_count
    }
    
    # 保存处理后的数据
    logger.info(f"保存处理后的轨迹: {processed_path}")
    with open(processed_path, 'w', encoding='utf-8') as f:
        json.dump(data, f, indent=2, ensure_ascii=False)
    
    logger.info("✅ 处理完成！")
    
    # 输出示例
    print("\n" + "=" * 60)
    print("📊 坐标转换示例（前5个点击）:")
    print("-" * 60)
    
    click_count = 0
    for action in actions:
        if action['type'] == 'click' and 'pos_host' in action:
            host_x, host_y = action['pos_host']
            vm_x, vm_y = action['pos_vm'] if action['pos_vm'] else (0, 0)
            # 转换为整数显示
            print(f"  Host({int(host_x):4d}, {int(host_y):4d}) → VM({int(vm_x):4d}, {int(vm_y):4d})")
            
            click_count += 1
            if click_count >= 5:
                break
    
    print("=" * 60)
    print("\n💡 下一步：可视化验证（可选）")
    print(f"   python scripts/visualize_trajectory.py {task_id}")
    print("=" * 60 + "\n")
    
    return True


def main():
    parser = argparse.ArgumentParser(description="处理轨迹数据，转换坐标")
    parser.add_argument("task_id", help="任务ID")
    parser.add_argument("--project-root", default=".", help="项目根目录")
    parser.add_argument("--force", action="store_true", help="强制覆盖已有文件")
    parser.add_argument("--scale-x", type=float, default=1.0, help="X轴缩放调整系数")
    parser.add_argument("--scale-y", type=float, default=1.0, help="Y轴缩放调整系数")
    parser.add_argument("--offset-x", type=int, default=0, help="X轴偏移调整")
    parser.add_argument("--offset-y", type=int, default=0, help="Y轴偏移调整")
    
    args = parser.parse_args()
    
    success = process_trajectory(
        task_id=args.task_id,
        project_root=args.project_root,
        force=args.force,
        scale_x_adjust=args.scale_x,
        scale_y_adjust=args.scale_y,
        offset_x=args.offset_x,
        offset_y=args.offset_y
    )
    
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()