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 parse_model_output(resp: str) -> Dict[str, Union[str, Tuple[int, int], None]]:
    """Extract last Action line from model response."""
    import re

    action_line_match = re.findall(r"Action:\s*([^\n]+)", resp.strip())
    if not action_line_match:
        return {"action_type": None, "action_param": None}
    action_line = action_line_match[-1]

    m = re.match(r"(\w+)\((.*)\)", action_line)
    if not m:
        return {"action_type": None, "action_param": None}
    action_type, param_str = m.groups()

    if action_type == "click":
        m2 = re.search(r"start_box=['\"]?\((\d+),(\d+)\)['\"]?", param_str)
        if m2:
            return {"action_type": "click", "action_param": (int(m2[1]), int(m2[2]))}

    elif action_type == "scroll":
        m2 = re.search(r"direction=['\"](down|up|left|right)['\"]", param_str)
        if m2:
            return {"action_type": "scroll", "action_param": m2[1]}

    elif action_type == "type":
        m2 = re.search(r"(content|app_name)=['\"]([^'\"]+)['\"]", param_str)
        if m2:
            return {"action_type": "input", "action_param": m2[2]}

    elif action_type in ["finished", "press_back", "press_home"]:
        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":
        if isinstance(step["action_param"], str):
            step["action_param"] = json.loads(step["action_param"])
        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]:

    prompt_origin = """You are a GUI agent. You are given a task and your action history, with screenshots. You need to perform the next action to complete the task. 
## Output Format
```
Action: ...
```
## Action Space
{action_space}

## User Instruction
{instruction}
    """

    action_space = """
    click(start_box='<|box_start|>(x1,y1)<|box_end|>')
    type(content='')
    scroll(start_box='<|box_start|>(x1,y1)<|box_end|>', direction='down or up or right or left')
    press_back()
    press_home()
    finished()
    """

    if vllm_format:
        system_msg = {
            "role": "system",
            "content": [{"type": "text", "text": "You are a helpful assistant."}],
        }
        user_prompt = prompt_origin.format(instruction=instruction, action_space=action_space)
        user_msg = {"role": "user", "content": [{"type": "text", "text": user_prompt}]}

        messages = [system_msg, user_msg]

        # attach history (limited by his_num images)
        img_idx = 0
        for h_idx, hist in enumerate(history_responses):
            if h_idx + his_num + 1 > len(history_responses):
                messages.append({
                    "role": "user",
                    "content": [{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(images[img_idx])}"}}]
                })
                img_idx += 1
            messages.append({"role": "assistant", "content": [{"type": "text", "text": hist}]})

        # add the latest screenshot (current state)
        messages.append({"role": "user", "content": [{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(images[img_idx])}"}}]})  

    else:
        system_msg = {
            "role": "system",
            "content": [{"type": "text", "text": "You are a helpful assistant."}],
        }
        user_prompt = prompt_origin.format(instruction=instruction, action_space=action_space)
        user_msg = {"role": "user", "content": [{"type": "text", "text": user_prompt}]}

        messages = [system_msg, user_msg]

        # attach history (limited by his_num images)
        img_idx = 0
        for h_idx, hist in enumerate(history_responses):
            if h_idx + his_num + 1 > len(history_responses):
                messages.append({
                    "role": "user",
                    "content": [{"type": "image", "image": images[img_idx]}]
                })
                img_idx += 1
            messages.append({"role": "assistant", "content": [{"type": "text", "text": hist}]})

        # add the latest screenshot (current state)
        messages.append({"role": "user", "content": [{"type": "image", "image": images[img_idx]}]})
        
    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]
            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

    for episode in tqdm(dataset, desc="Episodes"):
        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, h = episode['resolution']
            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="uitars-7b-dpo")
    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)
