from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info, smart_resize

import argparse, os, json, base64, re, httpx, torch, pytz, io
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 encode_image(image_path, w=None, h=None):
    with Image.open(image_path) as img:
        if img.mode != 'RGB':
            img = img.convert('RGB')

        if w is not None and h is not None:
            img = img.resize((w, h), Image.BILINEAR)

        buf = io.BytesIO()
        img.save(buf, format='JPEG', optimize=True)

        base64_bytes = base64.b64encode(buf.getvalue())
        return base64_bytes.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

        MAX_PIXELS = 4096 * 28 * 28
        h_b, w_b = smart_resize(h, w, max_pixels=MAX_PIXELS)
        bbox = [
            int(data['bounds'][0] * w_b / w),
            int(data['bounds'][1] * h_b / h),
            int(data['bounds'][2] * w_b / w),
            int(data['bounds'][3] * h_b / h)
        ]
        instruction = data['instruction']
        prompt = "Locate UI components that match the command: \"{}\". ".format(instruction)
        prompt+="Output a JSON in the format [{{\"bbox_2d\": [...], \"label\": \"{{the_whole_command}}\"}}, ...]."

        if not vllm_format:
            # 默认格式
            query = [
                {
                    "role": "system",
                    "content": "You are a helpful assistant."
                },
                {
                    "role": "user",
                    "content": [
                        {"type": 'image',    "image": img_file, "max_pixels": MAX_PIXELS},
                        {"type": 'text',     "text": prompt}
                    ]
            }]
        else:
            # vLLM 专用格式
            query = [
                {
                    "role": "system",
                    "content": "You are a helpful assistant."
                },
                {
                    "role": "user",
                    "content": [
                        {"type": 'image_url', "image_url": {"url": f"data:image/jpeg;base64,{encode_image(img_file, w_b, h_b)}"}},
                        {"type": 'text',      "text": prompt}
                    ]
            }]

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


def parse_output(text: str) -> list:
    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))   

        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.processor = AutoProcessor.from_pretrained(args.model_path)
        if args.batch_size != 1:
            self.processor.tokenizer.padding_side = 'left'
        self.model = Qwen2_5_VLForConditionalGeneration.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):
        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 = "Outline the position of \"{}\" and output all the coordinates in JSON format."
    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/MiMo-VL-7B-RL')
    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=None)
    parser.add_argument('--deployment_name',        type=str, default='mimovl-rl')
    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")
    m_name = os.path.basename(args.model_path) if not args.vllm_url else args.deployment_name
    fname = f"{now}_{m_name}_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)