from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info

import argparse, os, json, base64, re, httpx, torch, pytz
from openai import OpenAI
from collections import defaultdict
from datetime import datetime
from PIL import Image
from tqdm import tqdm
import numpy as np
from openai import AzureOpenAI


# ========== 工具函数 ==========
def encode_image(image_path: str) -> str:
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_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 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'])
        # w, h = Image.open(img_file).size
        w, h = data['size']
        bbox = [
            float(data['bounds'][0] / w),
            float(data['bounds'][1] / h),
            float(data['bounds'][2] / w),
            float(data['bounds'][3] / h)
        ]
        query = [
            {
                "role": "system",
                "content": system_prompt
            },
            {
                "role": "user",
                "content": [
                    {"type": 'image_url', "image_url": {"url": f"data:image/jpeg;base64,{encode_image('./gpt-example.jpeg')}"}},
                    {"type": 'text',      "text": "For example, the (x1, y1, x2, y2) coordinates for the region with the beer bottle are (0.47, 0.48, 0.55, 0.87)."}
                ]
            },
            {
                "role": "user",
                "content": [
                    {"type": 'image_url', "image_url": {"url": f"data:image/jpeg;base64,{encode_image(img_file)}"}},
                    {"type": 'text',      "text": "For this image below, what are the (x1, y1, x2, y2) coordinates for the region with the following task or description: {}".format(data['instruction'])}
                ]
        }]

        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:
    match = re.search(r"\(([\d.]+),\s*([\d.]+),\s*([\d.]+),\s*([\d.]+)\)", text)
    if match:
        x1, y1, x2, y2 = map(float, match.groups())
        center_x = round((x1 + x2) / 2, 3)
        center_y = round((y1 + y2) / 2, 3)
        return [center_x, center_y]
    else:
        return [0, 0]


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]


class VLLMVisionModel:
    def __init__(self, args):
        api_base = args.api_base
        api_key = args.api_key
        api_version = '2025-01-01-preview'
        self.deployment_name = args.model_path
        
        self.client = AzureOpenAI(
            api_key=api_key,
            api_version=api_version,
            azure_endpoint=api_base,
        )


    def infer_batch(self, chat):
        data = []
        for msg in chat:
            try:
                res = self.client.chat.completions.create(
                    model=self.deployment_name,
                    messages=msg,
                    max_tokens=2048
                )
                res = res.choices[0].message.content
            except Exception as e:
                res = "(0.0, 0.0)"
        return [res]


# ========== 评估函数 ==========
def evaluate(model, args):
    system_prompt = """
    Locate the referred object and represent the location of the region.
Regions are represented by (x1, y1, x2, y2) coordinates.
x1 and x2 are the leftmost and rightmost positions, normalized into [0, 1], where 0 is the left and 1 is the right.
y1 and y2 are the topmost and bottommost positions, normalized into [0, 1], where 0 is the top and 1 is the bottom.


"""
    data = prepare_data(system_prompt, args.json_path, args.img_path, True)
    all_scores, grouped = [], defaultdict(list)
    for batch in tqdm(batch_data(data, 1)):
        outputs = model.infer_batch([batch[0]['query']])
        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
            d.pop('query')
            key = f"{d['device']}_{d['data_type']}"
            grouped[key].append(float(correct))
            all_scores.append(float(correct))
        write_jsonl(batch, 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='gpt-4o')
    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/img/grounding')
    parser.add_argument('--result_path','-o',   type=str, default='your_path/gui-ceval/grounding')
    parser.add_argument('--api_base',   '-api_base', type=str, default='https://your_endpoint.openai.azure.com')
    parser.add_argument('--api_key',    '-api_key', type=str, default='your_api_key')
    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)

    evaluate(vision_model, args)