from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation import GenerationConfig

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
        bbox = [
            float(data['bounds'][0] / w),
            float(data['bounds'][1] / h),
            float(data['bounds'][2] / w),
            float(data['bounds'][3] / h)
        ]
        # 默认格式
        query = [
            {'image': img_file},
            {'text': system_prompt.format(data['instruction'])}
           ]
        # vLLM 专用格式
        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.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):
    match = re.search(r"\(\s*([\d.]+)\s*,\s*([\d.]+)\s*\)", text)
    if match:
        x, y = map(float, match.groups())  # 转换为浮点数
        return [x, 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.tokenizer_path, trust_remote_code=True)
        self.model = AutoModelForCausalLM.from_pretrained(args.model_path, device_map="cuda", trust_remote_code=True, bf16=True).eval().to(self.device)
        self.model.generation_config = GenerationConfig.from_pretrained(args.tokenizer_path, trust_remote_code=True)

    def infer_batch(self, batch: list) -> list:
        assert len(batch) == 1

        messages = batch[0]['query']

        query = self.tokenizer.from_list_format(messages)
        output_text, history = self.model.chat(self.tokenizer, query=query, history=None)
        
        return [output_text]

# ========== vLLM API 实现 ==========
class VLLMVisionModel:
    def __init__(self, args):
        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 = "In this UI screenshot, what is the position of the element corresponding to the command \"{}\" (with point)?"
    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)
        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 d['query'][0]['content'][0].get('image_url'):
            #     d['query'][0]['content'][0]['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/SeeClick/models_seeclick')
    parser.add_argument('--tokenizer_path', '-t',    type=str, default='your_path/SeeClick/models_qwen_vl_chat')
    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-2b-sft')
    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)