import os, math, time, json, base64, random, argparse, re, jieba, torch, httpx, pytz, io
from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
from qwen_vl_utils import process_vision_info
from openai import OpenAI
from datetime import datetime
from collections import defaultdict
from copy import deepcopy
from typing import Dict, List, Tuple, Union
import numpy as np
from PIL import Image
from tqdm import tqdm

random.seed(0)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)


# -------------------------- Helper Utilities ------------------------------ #
def clean_surrogates(obj):
    """递归清除 dict/list 中所有字符串里的 surrogate 字符"""
    if isinstance(obj, dict):
        return {k: clean_surrogates(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [clean_surrogates(v) for v in obj]
    elif isinstance(obj, str):
        return re.sub(r'[\ud800-\udfff]', '', obj)
    else:
        return obj


class AverageMeter:
    def __init__(self) -> None:
        self.reset()

    def reset(self) -> None:
        self.avg = 0.0
        self.sum = 0.0
        self.cnt = 0

    def update(self, val: float, n: int = 1) -> None:
        self.sum += val * n
        self.cnt += n
        self.avg = self.sum / self.cnt if self.cnt else 0.0


def write_jsonl(data: Union[dict, List[dict]], path: str) -> None:
    if isinstance(data, dict):
        data = [data]
    with open(path, "a", encoding="utf-8") as f:
        for d in data:
            f.write(json.dumps(clean_surrogates(d), ensure_ascii=False) + "\n")


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 calculate_f1_score(predicted: str, ground_truth: str) -> float:
    def tokenize(text: str):
        if not text:
            return set()
        # Chinese detection
        if re.search(r"[\u4e00-\u9fff]", text):
            return set(jieba.lcut(text))
        return set(text.lower().split())

    p, g = tokenize(predicted), tokenize(ground_truth)
    if not p or not g:
        return 0.0
    c = len(p & g)
    precision, recall = c / len(p), c / len(g)
    return 2 * precision * recall / (precision + recall) if precision + recall else 0.0


def is_click_accurate(pred: Tuple[int, int] | None, gt: List[float]) -> bool:
    if not pred or not gt:
        return False
    px, py = pred
    if len(gt) == 2:
        gx, gy = gt
        return math.hypot(px - gx, py - gy) <= 0.14 * 1000
    elif len(gt) == 4:
        x1, y1, x2, y2 = gt
        return x1 <= px <= x2 and y1 <= py <= y2
    else:
        raise ValueError(f"Invalid ground truth: {gt}")


# -------------------------- Dataset Parsing ------------------------------- #
# 动作空间：click， scroll， input， finished， home， back

def extract_action(content):
    answer_tag_pattern = r'<answer>(.*?)</answer>'
    action_pattern = r"'action':\s*'(\w+)'"
    content_answer_match = re.search(answer_tag_pattern, content, re.DOTALL)
    if content_answer_match:
        content_answer = content_answer_match.group(1).strip()
        action_match = re.search(action_pattern, content_answer)
        if action_match:
            return action_match.group(1)
    return None
def extract_coord(content):
    # Try to find the bbox within <answer> tags, if can not find, return [0, 0]
    answer_tag_pattern = r'<answer>(.*?)</answer>'
    bbox_pattern = r'\{.*\[(\d+),\s*(\d+)]\s*.*\}'
    content_answer_match = re.search(answer_tag_pattern, content, re.DOTALL)
    if content_answer_match:
        content_answer = content_answer_match.group(1).strip()
        coord_match = re.search(bbox_pattern, content_answer)
        if coord_match:
            coord = [int(coord_match.group(1)), int(coord_match.group(2))]
            x, y = coord
            return coord
    return [0, 0]


def parse_model_output(resp: str) -> Dict[str, Union[str, Tuple[int, int], None]]:
    """Extract last Action line from model response."""

    action_type = extract_action(resp)
    if action_type is None:
        return {"action_type": "finished", "action_param": None}

    if action_type in ['click', 'open_app']:
        coord = extract_coord(resp)
        return {"action_type": "click", "action_param": coord}

    elif action_type == "scroll":
        return {"action_type": "scroll", "action_param": "up"}

    elif action_type == "input_text":
        return {"action_type": "input", "action_param": "test"}

    elif action_type in ["navigate_back"]:
        return {"action_type": action_type, "action_param": None}

    return {"action_type": action_type, "action_param": None}


def data_action_to_prompt(step: dict, img_w: int, img_h: int) -> str:
    # 用于离线评估时拼接以前步骤的动作，需要将数据集格式转换为模型输入格式
    
    atype = step["action_type"]
    if atype == "scroll":
        direction = step["action_param"]
        action_str = f"scroll(direction='{direction}')"
    elif atype == "click":
        x = round((step["action_param"][0] + step["action_param"][2]) / 2 / img_w * 1000)
        y = round((step["action_param"][1] + step["action_param"][3]) / 2 / img_h * 1000)
        action_str = f"click(start_box='({x},{y})')"
    elif atype == "input":
        action_str = f"type(content='{step['action_param']}')"
    elif atype == "back":
        action_str = "press_back()"
    elif atype == "home":
        action_str = "press_home()"
    elif atype == "finished":
        action_str = "finished()"
    elif atype == "open":
        action_str = f"open_app(app_name='{step['action_param']}')"
    else:
        action_str = "finished()"
    return "Action: " + action_str


# ---------------------- Dialogue History Builder -------------------------- 
def prepare_data(instruction: str, images: List[str], history_responses: List[str], his_num: int, vllm_format: bool) -> List[dict]:

    image_path = images[-1]
    task_prompt = instruction

    question_template = (
        f"In this UI screenshot, I want to perform the command '{task_prompt}'.\n"
        "Please provide the action to perform (enumerate in ['click', 'open_app', 'scroll', 'navigate_back', 'input_text]')"
        "and the coordinate where the cursor is moved to(integer) if click is performed.\n"
        "Output the thinking process in <think> </think> and final answer in <answer> </answer> tags."
        "The output answer format should be as follows:\n"
        "<think> ... </think> <answer>[{'action': enum['click', 'open_app', 'scroll', 'navigate_back', 'input_text], 'coordinate': [x, y]}]</answer>\n"
        "Please strictly follow the format."
    )
    query = '<image>\n' + question_template

    if vllm_format:
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(image_path)}"}},
                    {"type": "text", "text": query}]
            }
    ]

    else:
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image_path},
                    {"type": "text", "text": query}]
            }
        ]
        
    return messages


# -------------------------- Model Wrappers ------------------------------- #
class TransformersVisionModel:
    def __init__(self, args: argparse.Namespace):
        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",
            ).to(self.device).eval()
        )

    def infer_batch(self, batch: List[dict]) -> List[str]:
        texts, imgs, vids = [], [], []
        for d in batch:
            txt = self.processor.apply_chat_template(
                d["messages"], tokenize=False, add_generation_prompt=True
            )
            img_in, vid_in = process_vision_info(d["messages"])
            texts.append(txt)
            imgs.append(img_in)
            vids.append(vid_in)

        vids = None if all(v is None for v in vids) else vids
        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():
            gen_ids = self.model.generate(**inputs, max_new_tokens=512)

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


class VLLMVisionModel:
    def __init__(self, args: argparse.Namespace):
        proxies = {"http://": args.proxy, "https://": args.proxy} if args.proxy else None
        self.client = OpenAI(
            base_url=args.vllm_url, api_key="empty", http_client=httpx.Client(proxies=proxies)
        )
        self.deployment = args.deployment_name

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


def match_and_update_metrics(meters, pred_action, ref_action, w, h):
    acc_type = int(pred_action["action_type"] == ref_action["action_type"])
    acc_action = False

    meters["type"].update(acc_type)

    if ref_action["action_type"] == "click":
        meters["click_type"].update(acc_type)
        if acc_type:
            gx1, gy1, gx2, gy2 = ref_action["action_param"]
            # gt_point = [gx1 / w * 1000, gy1 / h * 1000, gx2 / w * 1000, gy2 / h * 1000]
            gt_point = [gx1, gy1, gx2, gy2]
            acc_action = is_click_accurate(pred_action["action_param"], gt_point)
            meters["click_action"].update(acc_action)
    elif ref_action["action_type"] == "scroll":
        meters["scroll_type"].update(acc_type)
        if acc_type:
            acc_action = pred_action["action_param"] == ref_action["action_param"]
            meters["scroll_action"].update(acc_action)
    elif ref_action["action_type"] == "input":
        meters["input_type"].update(acc_type)
        if acc_type:
            score = calculate_f1_score(ref_action["action_param"], pred_action["action_param"] or "")
            acc_action = score >= 0.5
            meters["input_action"].update(acc_action)
    elif ref_action["action_type"] == "finished":
        meters["finished_action"].update(acc_type)
        acc_action = acc_type
    else:
        acc_action = acc_type

    meters["action"].update(acc_action)
    return acc_type, acc_action

# ----------------------------- Evaluation -------------------------------- #
def evaluate(model, args):
    with open(args.json_path, "r", encoding="utf-8") as f:
        dataset = [json.loads(l) for l in f]

    # --- metrics ---
    meters = {k: AverageMeter() for k in [
        "action", "type", "episode",
        "click_action", "click_type",
        "input_action", "input_type",
        "scroll_action", "scroll_type",
        "finished_action"
    ]}

    start_time = time.time()
    total_steps = 0

    for episode in tqdm(dataset, desc="Episodes"):
        goal = episode["instruction"]
        history_resp: List[str] = []
        images: List[str] = []
        episode_correct = True

        for step in episode["actions"]:
            total_steps += 1
            img_path = step["img_path"]
            img_path = os.path.join(args.img_path, img_path)
            images.append(img_path)
            if len(images) > args.his_num + 1:
                images = images[-(args.his_num + 1) :]

            w, h = Image.open(img_path).size
            ref_action = step
            # history_resp.append(data_action_to_prompt(step, w, h))

            # Build model messages
            msgs = prepare_data(goal, images, history_resp[:-1], args.his_num, vllm_format=bool(args.vllm_url))
            batch_item = {"messages": msgs}
            pred_txt = model.infer_batch([batch_item])[0]
            pred_action = parse_model_output(pred_txt)

            acc_type, acc_action = match_and_update_metrics(meters, pred_action, ref_action, w, h)
            episode_correct = episode_correct and bool(acc_action)

            # record step info
            info = {
                "goal": goal,
                "app": episode["app_name"],
                "img_path": img_path,
                "action_ref": ref_action,
                "action_pred": pred_action,
                "response": pred_txt,
                "correct_action": acc_action,
                "correct_type": acc_type,
            }
            write_jsonl(info, args.result_path_file)
            print(info)

        meters["episode"].update(int(episode_correct))

    elapsed = time.time() - start_time
    stats = {k: v.avg for k, v in meters.items()}
    stats.update({"avg_time_per_step": elapsed / total_steps})

    write_jsonl([{"item": k, "score": v} for k, v in stats.items()], args.result_path_file)
    for k, v in stats.items():
        print(f"{k}: {v:.4f}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_path", "-m", type=str, default="your_path/UI-TARS-7B-SFT")
    parser.add_argument("--json_path", "-json", type=str, default='your_path/offline_agent_short.jsonl')
    parser.add_argument("--img_path", "-img", type=str, default="your_path/img/offline_agent") 
    parser.add_argument("--result_path", "-o", type=str, default='your_path/gui-ceval/offline_agent')
    parser.add_argument("--vllm_url", "-url", type=str, default="your url")
    # parser.add_argument("--vllm_url", "-url", type=str, default=None)
    parser.add_argument("--deployment_name", type=str, default="uir1-3b")
    parser.add_argument("--proxy", type=str, default='if needed')
    parser.add_argument("--close_proxy", action="store_true")
    parser.add_argument("--batch_size", type=int, default=1)
    parser.add_argument("--his_num", type=int, default=0)
    args = parser.parse_args()
    if args.close_proxy:
        args.proxy = None

    # resolve output filename
    os.makedirs(args.result_path, exist_ok=True)
    now=datetime.now(pytz.timezone("Asia/Shanghai")).strftime("%Y%m%d%H%M%S")
    m_name = os.path.basename(args.model_path) if args.vllm_url is None else args.deployment_name
    fname = f"{now}_{m_name}_guiceval_offline_agent.jsonl"
    args.result_path_file = os.path.join(args.result_path, fname)

    model = VLLMVisionModel(args) if args.vllm_url is not None else TransformersVisionModel(args)
    evaluate(model, args)