import subprocess
subprocess.run(["pip", "install", "transformers==4.50.0"], check=True)
print("安装完成，继续执行后续代码...")

from transformers import AutoTokenizer, AutoModelForCausalLM
import argparse, os, json, base64, re, httpx, torch, pytz, io, requests
from openai import OpenAI
from collections import defaultdict
from datetime import datetime
from PIL import Image
from tqdm import tqdm
import numpy as np

# ========== 工具函数 ==========
def encode_image(image: Image.Image) -> str:
    """Convert PIL Image to base64-encoded string."""
    with io.BytesIO() as in_mem_file:
        image.save(in_mem_file, format="JPEG")
        in_mem_file.seek(0)
        return base64.b64encode(in_mem_file.read()).decode("utf-8")


def write_jsonl(data: list, path: str):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, 'a', encoding='utf-8') as f:
        for d in data:
            f.write(json.dumps(d, ensure_ascii=False) + '\n')

def __resize__(origin_img):
    resolution = origin_img.size
    w,h = resolution
    max_line_res = 1120
    if max_line_res is not None:
        max_line = max_line_res
        if h > max_line:
            w = int(w * max_line / h)
            h = max_line
        if w > max_line:
            h = int(h * max_line / w)
            w = max_line
    img = origin_img.resize((w,h),resample=Image.Resampling.LANCZOS)
    return img


def prepare_data(system_prompt: str, json_path: str, imgs_path: str, vllm_format=False) -> list:
    items = []
    with open(json_path, 'r', encoding='utf-8') as f:
        for line in f:
            items.append(json.loads(line))

    label_data = []
    for data in items:
        img_file = os.path.join(imgs_path, data['img_path'])
        image = Image.open(img_file).convert('RGB')
        w, h = image.size
        image = __resize__(image)
        bbox = [
            int(data['bounds'][0] / w * 1000),
            int(data['bounds'][1] / h * 1000),
            int(data['bounds'][2] / w * 1000),
            int(data['bounds'][3] / h * 1000)
        ]
        if not vllm_format:
        # 默认格式
            query = [{
                "role": "user",
                "content": [
                    f"<Question>{data['instruction']}</Question>\n当前屏幕截图：",
                    image
                ]
            }]
        else:
        # vLLM 专用格式
            query = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": [
                    {"type": "text", "text": f"<Question>{data['instruction']}</Question>\n当前屏幕截图：(<image>./</image>)"},
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(image)}"}}
                ]}
            ]

        label_data.append({
            'img_path': img_file,
            'query': query,
            'width': w,
            'height': h,
            'bbox': bbox,
            'device': data.get('device'),
            'data_type': data.get('data_type')
        })
    return label_data



def parse_output(text: str) -> list:
    try:
        data = json.loads(text)
        if all(k in data for k in ["x", "y", "width", "height"]):
            center_x = data["x"] + data["width"] // 2
            center_y = data["y"] + data["height"] // 2
            return [center_x, center_y]
    except json.JSONDecodeError:
        pass  # 继续执行正则解析

    match = re.search(r"\(\s*(\d+)\s*,\s*(\d+)\s*\)\s*,\s*\(\s*(\d+)\s*,\s*(\d+)\s*\)", text)
    if match:
        x1, y1, x2, y2 = map(int, match.groups())  # 转换为整数
        center_x = (x1 + x2) // 2
        center_y = (y1 + y2) // 2
        return [center_x, center_y]

    pattern = r"(\d+\.\d+|\d+)[, ]+(\d+\.\d+|\d+)(?:[, ]+(\d+\.\d+|\d+)[, ]+(\d+\.\d+|\d+))?"
    matches = re.findall(pattern, text)

    results = []
    for match in matches:
        numbers = list(filter(None, match))  # 过滤掉空匹配项
        numbers = list(map(float, numbers))  # 转换为浮点数
        numbers = list(map(int, numbers))    # 取整（四舍五入改为 `round()`）

        if len(numbers) == 4:
            x1, y1, x2, y2 = numbers
            center_x = (x1 + x2) // 2
            center_y = (y1 + y2) // 2
            results.append([center_x, center_y])  # 计算中心点
        elif len(numbers) == 2:
            results.append(numbers)
        else:
            results.append([0, 0])

    results = results[0] if len(results) != 0 else [0, 0]
    return results


def is_point_in_bbox(pt: list, bbox: list) -> bool:
    x, y = pt
    return bbox[0] <= x <= bbox[2] and bbox[1] <= y <= bbox[3]


def batch_data(data_list: list, bs: int):
    for i in range(0, len(data_list), bs):
        yield data_list[i:i+bs]

# ========== Transformers 模型实现 ==========
class TransformersVisionModel:
    def __init__(self, args):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
        self.model = AutoModelForCausalLM.from_pretrained(args.model_path, trust_remote_code=True, torch_dtype=torch.bfloat16).to(self.device)

    def infer_batch(self, batch: list, system_prompt: str) -> list:
        assert len(batch) == 1
        batch = batch[0]

        outs = self.model.chat(
            image=None,
            msgs=batch['query'],
            system_prompt=system_prompt,
            tokenizer=self.tokenizer,
            temperature=0.1,
            top_p=0.3,
            n=1,
        )
        return [outs]

# ========== vLLM API 实现 ==========
class VLLMVisionModel:
    def __init__(self, args):
        self.proxies = {"http://": args.proxy, "https://": args.proxy}
        self.END_POINT = args.vllm_url + '/chat/completions'
        self.headers = {
            "Content-Type": "application/json",
        }
    

    def infer_batch(self, batch: list, system_prompt: str=None) -> list:
        payload = {
            "model": args.deployment_name,
            "temperature": 0.1,
            "messages": batch[0]['query'],
            "max_tokens": 2048,
        }

        response = requests.post(self.END_POINT, headers=self.headers, json=payload, proxies=self.proxies) 
        assistant_msg = response.json()["choices"][0]["message"]["content"]
        return [assistant_msg]

# ========== 评估函数 ==========
def evaluate(model, args):
    ACTION_SCHEMA = json.load(open('./schema.json', encoding="utf-8"))
    items = list(ACTION_SCHEMA.items())
    insert_index = 3
    items.insert(insert_index, ("required", ["thought"])) # enable/disable thought by setting it to "required"/"optional"
    ACTION_SCHEMA = dict(items)
    
    system_prompt = f'''# Role
你是一名熟悉安卓系统触屏GUI操作的智能体，将根据用户的问题，分析当前界面的GUI元素和布局，生成相应的操作。

# Task
针对用户问题，根据输入的当前屏幕截图，输出下一步的操作。

# Rule
- 以紧凑JSON格式输出
- 输出操作必须遵循Schema约束

# Schema
{json.dumps(ACTION_SCHEMA, indent=None, ensure_ascii=False, separators=(',', ':'))}'''
    data = prepare_data(
        system_prompt, args.json_path, args.img_path,
        vllm_format=bool(args.vllm_url)
    )
    all_scores, grouped = [], defaultdict(list)
    for batch in tqdm(batch_data(data, args.batch_size)):
        outputs = model.infer_batch(batch, system_prompt)
        for d, txt in zip(batch, outputs):
            d['response'] = txt
            d['predict'] = parse_output(txt)
            correct = is_point_in_bbox(d['predict'], d['bbox'])
            d['correct'] = correct
            # mask base64 in logs
            # if type(d['query']) is list and type(d['query'][1]['content'][1]) is dict and d['query'][1]['content'][1].get('image_url'):
            #     d['query'][1]['content'][1]['image_url'] = 'masked'
            d.pop('query')
            key = f"{d['device']}_{d['data_type']}"
            grouped[key].append(float(correct))
            all_scores.append(float(correct))
        write_jsonl(data=batch, path=args.result_path_file)
        for itm in batch:
            print(itm)

    args_dict = vars(args)
    write_jsonl([args_dict], args.result_path_file)
    stats = [{"item": k, "score": np.mean(v)} for k, v in grouped.items()]
    stats.append({"item": "all", "score": np.mean(all_scores)})
    write_jsonl(stats, args.result_path_file)
    for s in stats:
        print(s)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_path',    '-m',    type=str, default='your_path/AgentCPM-GUI')
    parser.add_argument('--json_path',     '-json', type=str, default='your_path/pub_version/grounding.jsonl')
    parser.add_argument('--img_path',      '-img',  type=str, default='your_path/GUI-CEval/img/grounding')
    parser.add_argument('--result_path',   '-o',    type=str, default='your_path/gui-ceval/grounding')
    parser.add_argument('--vllm_url',      '-url',  type=str, default="your_vllm_url")  # 如果使用 vLLM API，设置此参数
    parser.add_argument('--deployment_name',        type=str, default='agentcpm')
    parser.add_argument('--proxy',         type=str, default='if needed')
    parser.add_argument('--batch_size',    type=int, default=1)
    args = parser.parse_args()

    # 构造输出文件路径
    now = datetime.now(pytz.timezone("Asia/Shanghai")).strftime("%Y%m%d%H%M%S")
    fname = f"{now}_{os.path.basename(args.model_path)}_guiceval_grounding.jsonl"
    args.result_path_file = os.path.join(args.result_path, fname)

    # 选择模型实现
    vision_model = VLLMVisionModel(args) if args.vllm_url else TransformersVisionModel(args)

    evaluate(vision_model, args)