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


# ========== 工具函数 ==========
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 = [
            int(data['bounds'][0] / w * 1000),
            int(data['bounds'][1] / h * 1000),
            int(data['bounds'][2] / w * 1000),
            int(data['bounds'][3] / h * 1000)
        ]
        query = [{
            "role": "user",
            "content": [
                {"type": 'image', "image": img_file},
                {"type": 'text',  "text": system_prompt + data['instruction']}
            ]
        }]
        if vllm_format:
            query = [{
                "role": "user",
                "content": [
                    {"type": 'image_url', "image_url": {"url": f"data:image/jpeg;base64,{encode_image(img_file)}"}},
                    {"type": 'text',  "text": system_prompt + 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:
    m = re.match(r"\((\d+),\s*(\d+)\)", text)
    if m:
        return [int(m.group(1)), int(m.group(2))]
    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]


# ========== Transformers 模型实现 ==========
class TransformersVisionModel():
    def __init__(self, args):
        super().__init__()
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.processor = AutoProcessor.from_pretrained(args.model_path)
        if args.batch_size != 1:
            self.processor.tokenizer.padding_side = 'left'
        self.model = Qwen2VLForConditionalGeneration.from_pretrained(
            args.model_path,
            torch_dtype=torch.bfloat16,
            attn_implementation="flash_attention_2",
            device_map="auto"
        ).eval().to(self.device)

    def infer_batch(self, batch: list) -> list:
        texts, imgs, vids = [], [], []
        for d in batch:
            txt = self.processor.apply_chat_template(
                d['query'], tokenize=False, add_generation_prompt=True)
            texts.append(txt)
            img_in, vid_in = process_vision_info(d['query'])
            imgs.append(img_in)
            vids.append(vid_in)

        if all(v is None for v in vids):
            vids = None

        inputs = self.processor(
            text=texts,
            images=imgs,
            videos=vids,
            padding=True,
            return_tensors="pt"
        )
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        with torch.no_grad():
            out_ids = self.model.generate(**inputs, max_new_tokens=128)

        outs = []
        for inp, gen in zip(inputs['input_ids'], out_ids):
            trim = gen[len(inp):]
            text = self.processor.decode(
                trim,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False
            )
            outs.append(text)
        return outs

# ========== vLLM API 实现 ==========
class VLLMVisionModel():
    def __init__(self, args):
        super().__init__()
        proxies = {"http://": args.proxy, "https://": args.proxy}
        client = httpx.Client(proxies=proxies)
        self.openai = OpenAI(base_url=args.vllm_url,
                             api_key="empty",
                             http_client=client)
        self.deployment = args.deployment_name

    def infer_batch(self, batch: list) -> list:
        results = []
        for d in batch:
            try:
                res = self.openai.chat.completions.create(
                    model=self.deployment,
                    messages=d['query'],
                    frequency_penalty=1,
                    max_tokens=2048,
                    temperature=0.0,
                    top_p=0.7
                )
                results.append(res.choices[0].message.content)
            except Exception as e:
                print(f"⚠️ vLLM 调用失败: {e}")
                results.append(None)
        return results

# ========== 评估函数 ==========
def evaluate(model, args):
    system_prompt = r"""Output only the coordinate of one point in your response. What element matches the following task: """
    data = prepare_data(system_prompt, args.json_path, args.img_path, args.vllm_url is not None)
    all_scores, grouped = [], defaultdict(list)
    for batch in tqdm(batch_data(data, args.batch_size)):
        outputs = model.infer_batch(batch)
        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

            if d['query'][0]['content'][0].get('image_url'):
                d['query'][0]['content'][0]['image_url'] = 'masked'
            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='your_path/UI-TARS-72B-SFT')
    parser.add_argument('--json_path',  '-json', type=str, default='your_path/GUI-CEval/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=None)
    parser.add_argument('--deployment_name', type=str, default='ui-tars-72b-sft')
    parser.add_argument('--proxy',   type=str, default='None')
    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)