import json
import os
import re
import requests
import logging
import base64
from openai import OpenAI
from dotenv import load_dotenv
from PIL import Image
import random
import matplotlib.pyplot as plt
from collections import Counter
from datetime import datetime
import shutil
import torch
from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM
import io
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
test_dir = 'test'
# 创建test目录
os.makedirs(test_dir, exist_ok=True)

# 在全局范围内定义模型和tokenizer变量
tokenizer = None
model = None

# 在全局范围内定义模型和处理器变量
qwen_model = None
processor = None

# 添加模型加载函数
def load_json_data(file_path):
    """
    从指定路径加载JSON数据
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            return json.load(file)
    except Exception as e:
        logger.error(f"加载JSON文件失败: {e}")
        return None

def encode_image_to_base64(image_path):
    """
    将图片编码为base64字符串
    """
    try:
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode('utf-8')
    except Exception as e:
        logger.error(f"图片编码失败: {e}")
        return None

def load_qwen_model():
    """
    加载Qwen2VL模型和处理器
    """
    global qwen_model, processor
    
    # 设置本地模型路径
    model_path = "/data1/yuyr/models--bytedance-research--UI-TARS-7B-DPO/snapshots/727b0df39207dafc6cf211a61f29d84b7659c39c/"
    
    try:
        logger.info("正在加载Qwen2VL模型和处理器...")
        start_time = datetime.now()
        
        # 加载模型
        qwen_model = Qwen2VLForConditionalGeneration.from_pretrained(
            model_path, torch_dtype="auto", device_map="cuda:1"
        )
        logger.info("Qwen2VL模型加载完成")
        
        # 加载processor
        processor = AutoProcessor.from_pretrained(model_path)
        logger.info("处理器加载完成")
        
        end_time = datetime.now()
        load_time = (end_time - start_time).total_seconds()
        logger.info(f"Qwen2VL模型和处理器加载完成，耗时: {load_time:.2f}秒")
        
        return True
    except Exception as e:
        logger.error(f"Qwen2VL模型加载失败: {e}")
        return False

def call_gpt4o_mini(image_path, title, messages=None, max_retries=3):
    """
    使用本地HuggingFace模型代替OpenAI API，获取下一步应该点击的坐标
    """
    global qwen_model, processor
    
    # 使用Qwen2VL模型进行推理
    try:
        # 检查模型是否已加载
        if qwen_model is None or processor is None:
            logger.warning("模型或处理器未加载，尝试加载...")
            if not load_qwen_model():
                raise ValueError("模型加载失败")
        
        # 导入必要的库
        from qwen_vl_utils import process_vision_info
        
        # 如果是新会话，初始化消息列表
        if messages is None:
            messages = [
                {
                    "role": "system",
                    "content": """You are a GUI agent. You are given a task and your action history, with screenshots. You need to perform the next action to complete the task. 

## Output Format
```
Thought: ...
Action: ...
```

## Action Space
click(start_box='<|box_start|>(x1,y1)<|box_end|>')
left_double(start_box='<|box_start|>(x1,y1)<|box_end|>')
right_single(start_box='<|box_start|>(x1,y1)<|box_end|>')
drag(start_box='<|box_start|>(x1,y1)<|box_end|>', end_box='<|box_start|>(x3,y3)<|box_end|>')
hotkey(key='')
type(content='') #If you want to submit your input, use \"\" at the end of `content`.
scroll(start_box='<|box_start|>(x1,y1)<|box_end|>', direction='down or up or right or left')
wait() #Sleep for 5s and take a screenshot to check for any changes.
finished()
call_user() # Submit the task and call the user when the task is unsolvable, or when you need the user's help.

## Note
- Use Chinese in `Thought` part.
- Summarize your next action (with its target element) in one sentence in `Thought` part."""
                }
            ]
        else:
            # 把messages中的role为user的content中的image删掉
            for message in messages:
                if message["role"] == "user":
                    message["content"] = [content for content in message["content"] if content["type"] != "image"]
        
        # 构建当前步骤的用户消息
        user_message = {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image_url": f"file://{image_path}"
                },
                {"type": "text", "text": title},
            ],
        }
        
        # 添加用户消息到会话历史
        messages.append(user_message)
        
        # 准备推理输入
        text = processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        image_inputs, video_inputs = process_vision_info(messages)
        inputs = processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        )
        inputs = inputs.to(qwen_model.device)
        
        # 推理：生成输出
        logger.info("正在生成模型回复...")
        generated_ids = qwen_model.generate(**inputs, max_new_tokens=1024)
        generated_ids_trimmed = [
            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
        ]
        output_text = processor.batch_decode(
            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
        )[0]
        
        logger.info(f"模型输出: {output_text}")
        
        # 添加助手的回复到会话历史
        assistant_message = {
            "role": "assistant",
            "content": output_text
        }
        messages.append(assistant_message)
        
        return output_text, messages
    except Exception as e:
        logger.error(f"模型推理失败: {e}")
        if max_retries > 0:
            logger.warning(f"重试中，剩余次数: {max_retries}")
            return call_gpt4o_mini(image_path, title, messages, max_retries - 1)
        else:
            logger.error("重试次数已用完，推理失败")
            return None, messages

def extract_coordinates(response_text, image_path):
    """
    从API响应中提取坐标，并将1000x1000的坐标映射到实际图片尺寸
    """
    logger.info(f"API响应: {response_text}")
    try:
        # 使用正则表达式提取坐标
        pattern = r'\((\d+),(\d+)\)'  # 匹配格式 (x,y)
        match = re.search(pattern, response_text)
        if match:
            # 获取1000x1000下的坐标
            x_1000 = int(match.group(1))
            y_1000 = int(match.group(2))
            
            # 获取实际图片尺寸
            with Image.open(image_path) as img:
                actual_width, actual_height = img.size
            
            # 映射坐标到实际尺寸
            x = round(actual_width * x_1000 / 1000)
            y = round(actual_height * y_1000 / 1000)
            
            logger.info(f"坐标映射: 从 ({x_1000}, {y_1000}) 映射到 ({x}, {y})")
            return {
                "raw_x": x_1000,
                "raw_y": y_1000,
                "adjust_x": x,
                "adjust_y": y
            }
        return None
    except Exception as e:
        logger.error(f"坐标提取或映射失败: {e}")
        return None

def is_point_in_box(point, box):
    """
    判断点是否在边界框内
    """
    return (box["x"] <= point["adjust_x"] <= box["x"] + box["width"] and
            box["y"] <= point["adjust_y"] <= box["y"] + box["height"])

def test_path(url, path_data, use_title=True):
    """
    测试单个路径
    
    Args:
        url: 路径URL
        path_data: 路径数据
        use_title: 是否使用title作为任务描述，False则使用description
    """
    logger.info(f"开始测试路径: {url}, 使用{'标题' if use_title else '描述'}")
    meta = path_data["shortestPathsMeta"][0]
    title = meta["task_summaries"][0]["question"] if use_title else meta["raw_result"]
    task_type = "title" if use_title else "raw_result"
    chain_ids = meta["chainIDs"]
    child_nums = meta["chainChildNum"]
    viewport_boxes = meta["chainViewportBoundingBoxes"]
    
    total_steps = len(chain_ids) - 1
    success_steps = 0
    steps_data = []
    
    logger.info(f"任务内容: {title}")
    logger.info(f"总步骤数: {total_steps}")
    
    # 初始化会话历史
    messages = None
    
    # 遍历每个步骤
    for i in range(total_steps):
        current_id = chain_ids[i]
        child_num = child_nums[i] if i < len(child_nums) else None
        
        # 构建图片路径
        if child_num is not None:
            image_path = f"screenshots/{current_id}_{child_num}.png"
        else:
            image_path = f"screenshots/{current_id}.png"
        
        logger.info(f"步骤 {i+1}/{total_steps}: 处理图片 {image_path}")
        
        # 检查图片是否存在
        if not os.path.exists(image_path):
            logger.error(f"图片不存在: {image_path}")
            break
        
        # 调用模型获取点击坐标，传入会话历史
        response, messages = call_gpt4o_mini(image_path, title, messages)
        if not response:
            logger.error("模型调用失败")
            break
        
        # 提取坐标
        coordinates = extract_coordinates(response, image_path)
        if not coordinates:
            logger.error("无法从响应中提取坐标")
            break
        
        logger.info(f"模型返回坐标: {coordinates}")
        
        # 判断坐标是否在边界框内
        target_box = viewport_boxes[i] if i < len(viewport_boxes) else None
        step_success = False
        
        if target_box and is_point_in_box(coordinates, target_box):
            logger.info("坐标在边界框内，步骤成功!")
            success_steps += 1
            step_success = True
        else:
            logger.warning(f"坐标不在边界框内，步骤失败。目标边界框: {target_box}")
        
        # 记录步骤数据
        step_data = {
            "model_output": response,
            "raw_x": coordinates["raw_x"],
            "raw_y": coordinates["raw_y"],
            "adjust_x": coordinates["adjust_x"],
            "adjust_y": coordinates["adjust_y"],
            "bounding_box": target_box,
            "is_success": step_success
        }
        steps_data.append(step_data)
        
        # 如果步骤失败，终止后续步骤
        if not step_success:
            break
    
    # 总结测试结果
    is_success = (success_steps == total_steps)
    logger.info(f"路径测试完成: 总步骤 {total_steps}, 成功步骤 {success_steps}, 路径是否完全成功: {is_success}")
    
    return {
        "steps_data": steps_data,
        "success_steps": success_steps,
        "is_success": is_success
    }

def main():
    """
    主函数，从path数据中随机选择10条路径进行测试
    """
    # 加载Qwen2VL模型
    logger.info("开始加载Qwen2VL模型...")
    if not load_qwen_model():
        logger.error("Qwen2VL模型加载失败，程序退出")
        return
    
    # 加载原始路径数据
    path_json = "path/processed_3_with_analysis.json"
    data = load_json_data(path_json)
    if not data:
        logger.error("数据加载失败，程序退出")
        return
    
    # 随机选择10条路径
    urls = list(data.keys())
    # if len(urls) > 10:
    #     selected_urls = random.sample(urls, 10)
    # else:
    #     selected_urls = urls
    selected_urls = urls
    
    # 准备结果数据
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S,%f")[:-3]
    result_data = {
        "timestamp": timestamp,
        "num": len(selected_urls),
        "list": []
    }
    
    # 测试每条路径
    for url in selected_urls:
        path_data = data[url]
        meta = path_data["shortestPathsMeta"][0]
        title = meta["task_summaries"][0]["question"]
        description = meta["raw_result"]
        total_steps = len(meta["chainIDs"]) - 1
        
        # 记录路径长度，用于文件名
        path_length = total_steps
        
        logger.info(f"测试路径: {url}")
        logger.info(f"标题: {title}")
        logger.info(f"描述: {description}")
        
        # 使用标题测试
        title_result = test_path(url, path_data, use_title=True)
        
        # 使用描述测试
        desc_result = test_path(url, path_data, use_title=False)
        
        # 记录结果
        path_result = {
            "url": url,
            "title": title,
            "description": description,
            "total_steps": total_steps,
            "title_steps": title_result["steps_data"],
            "desc_steps": desc_result["steps_data"],
            "title_success_steps": title_result["success_steps"],
            "title_is_success": title_result["is_success"],
            "desc_success_step": desc_result["success_steps"],
            "desc_is_success": desc_result["is_success"]
        }
        
        result_data["list"].append(path_result)
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    # 保存结果
    result_file = f"{test_dir}/sample_path_{path_length + 1}_{timestamp}.json"
    with open(result_file, 'w', encoding='utf-8') as f:
        json.dump(result_data, f, indent=2, ensure_ascii=False)
    
    logger.info(f"结果已保存到 {result_file}")
    
    # 统计成功步数的分布并生成图表
    generate_statistics(result_data)

def generate_statistics(result_data):
    """
    生成统计图表
    """
    # 统计标题测试的成功步数分布
    title_success_steps = [item["title_success_steps"] for item in result_data["list"]]
    title_success_counts = Counter(title_success_steps)
    
    # 统计描述测试的成功步数分布
    desc_success_steps = [item["desc_success_step"] for item in result_data["list"]]
    desc_success_counts = Counter(desc_success_steps)
    
    # 绘制标题测试的成功步数分布图
    plt.figure(figsize=(10, 6))
    title_steps = sorted(title_success_counts.keys())
    title_counts = [title_success_counts[step] for step in title_steps]
    
    plt.bar(title_steps, title_counts, color='blue', alpha=0.7)
    plt.title('标题测试成功步数分布')
    plt.xlabel('成功步数')
    plt.ylabel('URL数量')
    
    # 在每个柱子上方显示具体数值
    for i, count in enumerate(title_counts):
        plt.text(title_steps[i], count, str(count), ha='center', va='bottom')
    
    # 保存图片
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    title_chart_path = f'{test_dir}/title_success_steps_distribution_{timestamp}.png'
    plt.savefig(title_chart_path)
    plt.close()
    
    # 绘制描述测试的成功步数分布图
    plt.figure(figsize=(10, 6))
    desc_steps = sorted(desc_success_counts.keys())
    desc_counts = [desc_success_counts[step] for step in desc_steps]
    
    plt.bar(desc_steps, desc_counts, color='green', alpha=0.7)
    plt.title('描述测试成功步数分布')
    plt.xlabel('成功步数')
    plt.ylabel('URL数量')
    
    # 在每个柱子上方显示具体数值
    for i, count in enumerate(desc_counts):
        plt.text(desc_steps[i], count, str(count), ha='center', va='bottom')
    
    # 保存图片
    desc_chart_path = f'{test_dir}/desc_success_steps_distribution_{timestamp}.png'
    plt.savefig(desc_chart_path)
    plt.close()
    
    # 绘制对比图
    plt.figure(figsize=(12, 7))
    
    # 合并所有可能的步数
    all_steps = sorted(set(title_steps + desc_steps))
    
    # 获取每个步数对应的计数
    title_all_counts = [title_success_counts.get(step, 0) for step in all_steps]
    desc_all_counts = [desc_success_counts.get(step, 0) for step in all_steps]
    
    # 设置柱状图的位置
    x = range(len(all_steps))
    width = 0.35
    
    # 绘制柱状图
    plt.bar([i - width/2 for i in x], title_all_counts, width, label='brief', color='blue', alpha=0.7)
    plt.bar([i + width/2 for i in x], desc_all_counts, width, label='detail', color='green', alpha=0.7)
    
    # 设置图表标题和标签
    plt.title('brief vs detail success steps distribution')
    plt.xlabel('success steps')
    plt.ylabel('URL count')
    plt.xticks(x, all_steps)
    plt.legend()
    
    # 在每个柱子上方显示具体数值
    for i, count in enumerate(title_all_counts):
        if count > 0:
            plt.text(i - width/2, count, str(count), ha='center', va='bottom')
    
    for i, count in enumerate(desc_all_counts):
        if count > 0:
            plt.text(i + width/2, count, str(count), ha='center', va='bottom')
    
    # 保存图片
    compare_chart_path = f'{test_dir}/title_vs_desc_distribution_{timestamp}.png'
    plt.savefig(compare_chart_path)
    plt.close()
    
    logger.info(f"统计图表已保存: {title_chart_path}, {desc_chart_path}, {compare_chart_path}")

if __name__ == "__main__":
    main()