import os, math, time, json, base64, random, argparse, re, jieba, torch, httpx, pytz, io
from transformers import Qwen2_5_VLForConditionalGeneration, 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 infer_direction(start: List[float], end: List[float]) -> str:
    dx, dy = end[0] - start[0], end[1] - start[1]
    return ("up" if dy < 0 else "down") if abs(dy) > abs(dx) else ("left" if dx < 0 else "right")

def parse_model_output(response_text: str) -> Dict[str, Union[str, list, None]]:
    try:
        # 仅支持 <tool_call> 包裹的 JSON
        start = response_text.find("{")
        end = response_text.rfind("}")
        if start == -1 or end == -1:
            raise ValueError("No valid JSON found")

        json_part = response_text[start:end+1]
        obj = json.loads(json_part)

        args = obj.get("arguments", {})
        raw_type = args.get("action")
        if not raw_type:
            return {"action_type": None, "action_param": None}

        raw_type = raw_type.lower()
        if raw_type == "click":
            return {"action_type": "click", "action_param": args.get("coordinate")}

        elif raw_type == "type":
            return {"action_type": "input", "action_param": args.get("text")}

        elif raw_type == "swipe":
            p1, p2 = args.get("coordinate"), args.get("coordinate2")
            if isinstance(p1, list) and isinstance(p2, list):
                direction = infer_direction(p1, p2)
                return {"action_type": "scroll", "action_param": direction}
            return {"action_type": "scroll", "action_param": None}

        elif raw_type == "terminate":
            return {"action_type": "finished", "action_param": args.get("status")}

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

    except Exception as e:
        print(f"[parse_model_output] Error parsing response: {e}")
        return {"action_type": None, "action_param": None}

def data_action_to_prompt(step: dict, img_w: int, img_h: int) -> str:
    atype = step["action_type"]
    action_json = {"action": None, "text": None, "coordinate": None, "coordinate2": None, "status": None}

    if atype == "click":
        if isinstance(step["action_param"], str):
            step["action_param"] = json.loads(step["action_param"])
        box = step["action_param"]
        x = round((box[0] + box[2]) / 2)
        y = round((box[1] + box[3]) / 2)
        action_json["action"] = "click"
        action_json["coordinate"] = [x, y]

    elif atype == "input":
        action_json["action"] = "type"
        action_json["text"] = step["action_param"]

    elif atype == "scroll":
        note = step.get("action_note")
        if isinstance(note, str):
            note = json.loads(note)
        if note and isinstance(note, list) and len(note) == 2:
            x1 = int(note[0][0])
            y1 = int(note[0][1])
            x2 = int(note[1][0])
            y2 = int(note[1][1])
            action_json["action"] = "swipe"
            action_json["coordinate"] = [x1, y1]
            action_json["coordinate2"] = [x2, y2]
        else:
            action_json["action"] = "swipe"

    elif atype == "finished":
        action_json["action"] = "terminate"
        action_json["status"] = step.get("action_param", "success")

    else:
        action_json["action"] = "terminate"
        action_json["status"] = "failure"

    return json.dumps({"name": "mobile_use", "arguments": {k: v for k, v in action_json.items() if v is not None}}, ensure_ascii=False)


# ---------------------- Dialogue History Builder -------------------------- 
def prepare_data(instruction: str, images: List[str], history_responses: List[str], his_num: int, vllm_format: bool) -> List[dict]:
    width, height = Image.open(images[-1]).convert("RGB").size
    system_prompt_template = (
    "You are a helpful assistant.\n\n"
    "# Tools\n\n"
    "You may call one or more functions to assist with the user query.\n\n"
    "You are provided with function signatures within <tools></tools> XML tags:\n"
    "<tools>\n"
    "{\"type\": \"function\", \"function\": {\"name_for_human\": \"mobile_use\", \"name\": \"mobile_use\", "
    "\"description\": \"Use a touchscreen to interact with a mobile device, and take screenshots.\\n"
    "* This is an interface to a mobile device with touchscreen. You can perform actions like clicking, typing, swiping, etc.\\n"
    "* Some applications may take time to start or process actions, so you may need to wait and take successive screenshots to see the results of your actions.\\n")
    system_prompt_template += f"* The screen's resolution is {width}x{height}.\\n"
    system_prompt_template += (
    "* Make sure to click any buttons, links, icons, etc with the cursor tip in the center of the element. Don't click boxes on their edges unless asked.\", "
    "\"parameters\": {\"type\": \"object\", \"properties\": {"
    "\"action\": {\"type\": \"string\", \"description\": \"The action to perform. The available actions are:\\n"
    "* `click`: Click the point on the screen with coordinate (x, y).\\n"
    "* `type`: Input the specified text into the activated input box.\\n"
    "* `swipe`: Swipe from the starting point with coordinate (x, y) to the end point with coordinates2 (x2, y2).\\n"
    "* `terminate`: Terminate the current task and report its completion status.\", "
    "\"enum\": [\"click\", \"type\", \"swipe\", \"terminate\"]}, "
    "\"coordinate\": {\"type\": \"array\", \"description\": \"(x, y): The coordinate to click or swipe from.\"}, "
    "\"coordinate2\": {\"type\": \"array\", \"description\": \"(x, y): The end coordinate of swipe.\"}, "
    "\"text\": {\"type\": \"string\", \"description\": \"Required only by `action=type`.\"}, "
    "\"status\": {\"type\": \"string\", \"enum\": [\"success\", \"failure\"], \"description\": \"Required only by `action=terminate`.\"}}, "
    "\"required\": [\"action\"]}, \"args_format\": \"Format the arguments as a JSON object.\"}}\n"
    "</tools>\n\n"
    "For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n"
    "<tool_call>\n"
    "{\"name\": \"mobile_use\", \"arguments\": { ... }}\n"
    "</tool_call>"
)
   
    # system_prompt = system_prompt_template.format(width=width, height=height)

    if vllm_format:
        user_prompt = "The user query:  {}\nTask progress (You have done the following operation on the current device):".format(instruction)
        for h_idx, hist in enumerate(history_responses):
            user_prompt +=f"Step {h_idx+1}: {hist};"

        messages = [
            {"role": "system", "content": system_prompt_template},
            {"role": "user", "content": [
                {"type": "text", "text": user_prompt}, 
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(images[-1])}"}}
            ]}
        ]

    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 = (
            Qwen2_5_VLForConditionalGeneration.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, gy1 / h, gx2 / w, gy2 / h]
            # 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_all", "episode_1e", "episode_3e",
        "click_action", "click_type",
        "input_action", "input_type",
        "scroll_action", "scroll_type",
        "finished_action"
    ]}

    app_episode_stats = defaultdict(lambda: {"total": 0, "all": 0, "1e": 0, "3e": 0})

    start_time = time.time()
    total_steps = 0
    total_episodes = 0

    for episode in tqdm(dataset, desc="Episodes"):
        total_episodes += 1
        # if total_episodes > 120:
        #     break
        goal = episode["instruction"]
        history_resp: List[str] = []
        images: List[str] = []
        wrong_count = 0

        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
            w_o, h_o = episode['resolution']
            w = w_o / w
            h = h_o / h
            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)
            if not acc_action:
                wrong_count += 1

            # 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_all"].update(int(wrong_count == 0))
        meters["episode_1e"].update(int(wrong_count <= 1))
        meters["episode_3e"].update(int(wrong_count <= 3))

        # 更新 app 层级统计
        app_name = episode["app_name"]
        app_episode_stats[app_name]["total"] += 1
        app_episode_stats[app_name]["all"] += int(wrong_count == 0)
        app_episode_stats[app_name]["1e"] += int(wrong_count <= 1)
        app_episode_stats[app_name]["3e"] += int(wrong_count <= 3)

    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}")

    # 输出每个app的 episode all / 1e / 3e
    # write_jsonl([{"item": k, "score": v} for k, v in stats.items()], args.result_path_file)
    # print("\nPer-app episode stats:")
    # for app, stat in app_episode_stats.items():
    #     total = stat["total"]
    #     all_acc = stat["all"] / total
    #     acc_1e = stat["1e"] / total
    #     acc_3e = stat["3e"] / total
    #     print(f"{app:20s} | episode_all: {all_acc:.3f} | episode_1e: {acc_1e:.3f} | episode_3e: {acc_3e:.3f}")


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/pub_version/offline_agent.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="qwen25vl-7b")
    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=4)
    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)
