import os, math, time, json, base64, random, argparse, re, jieba, torch, httpx, pytz, io
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
from openai import AzureOpenAI
from adb_utils import *

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)
        m2 = re.search(r"\((\d+),\s*(\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. All coordinates are normalized to the range of 0 to 1000.
## Output Format
```Thought: ...\n
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

class VLLMVisionModel:
    def __init__(self, args):
        api_base = args.api_base
        api_key = args.api_key
        api_version = '2025-01-01-preview'
        self.deployment_name = args.deployment_name
        
        self.client = AzureOpenAI(
            api_key=api_key,
            api_version=api_version,
            azure_endpoint=api_base,
        )


    def infer_batch(self, chat):
        outs = []
        for msg in chat:
            try:
                res = self.client.chat.completions.create(
                    model=self.deployment_name,
                    messages=msg['messages'],
                    max_tokens=2048
                )
                usage = {
                    "prompt_tokens": res.usage.prompt_tokens if res.usage else 0,
                    "completion_tokens": res.usage.completion_tokens if res.usage else 0,
                }
                outs.append([res.choices[0].message.content, usage])
            except Exception as e:
                print(f"⚠️ vLLM 调用失败: {e}")
                outs.append(["", {"prompt_tokens": 0, "completion_tokens": 0}])
        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 execute_actions(device_id, extracted_actions, w, h):
    for action in extracted_actions:
        if action["action_type"] == "click":
            click_point = action["action_param"]
            click_point = [click_point[0] / 1000 * w, click_point[1] / 1000 * h]
            tap(device_id, click_point[0], click_point[1])
            print(f"Click at: {click_point[0]}, {click_point[1]}")
        elif action["action_type"] == "scroll":
            scroll_direction = action["action_param"]
            if scroll_direction == "up":
                swipe(device_id, 0.5*w, 0.7*h, 0.5*w, 0.3*h)
            elif scroll_direction == "down":
                swipe(device_id, 0.5*w, 0.3*h, 0.5*w, 0.7*h)
            elif scroll_direction == "left":
                swipe(device_id, 0.7*w, 0.7*h, 0.3*w, 0.7*h)
            elif scroll_direction == "right":
                swipe(device_id, 0.3*w, 0.7*h, 0.7*w, 0.7*h)
        elif action["action_type"] == "input":
            input_content = action["action_param"]
            type_text(device_id, input_content)
        elif action["action_type"] == "press_back":
            press_back(device_id)
        elif action["action_type"] == "press_home":
            press_home(device_id)
        elif action["action_type"] == "finished":
            return
        

# ----------------------------- Evaluation -------------------------------- #
def evaluate(model, args):
        # 读取所有 jsonl 文件
    dataset = []
    # for root, _, files in os.walk('./guiceval_online'):
    #     for file in files:
    #         if file.endswith(".jsonl"):
    #             with open(os.path.join(root, file), "r", encoding="utf-8") as f:
    #                 dataset += [json.loads(l) for l in f]
    with open('./guiceval_online/guiceval_online.jsonl', "r", encoding="utf-8") as f:
        dataset += [json.loads(l) for l in f]

    test_time = datetime.now(pytz.timezone("Asia/Shanghai")).strftime("%Y%m%d-%H%M%S")
    test_file_path = os.path.join(args.result_path, f"{test_time}_{args.deployment_name}_{args.mode}")
    os.makedirs(test_file_path, exist_ok=True)
    args.result_path_file = os.path.join(test_file_path, "result.jsonl")

    for ii, episode in tqdm(enumerate(dataset), total=len(dataset)):
        # if ii > 2:
        #     break

        goal = episode["instruction"]
        app_name = episode.get("app_name", "unknown")

        w, h = get_resolution(args.device_id)
        home_clear(args.device_id)

        if args.mode == "home":
            time.sleep(2)
        elif args.mode == "other":
            open_app(args.device_id, "com.mi.car.mobile")
            time.sleep(5)
        elif args.mode == "intra":
            open_app(args.device_id, episode["package_name"])
            time.sleep(5)
            tap(args.device_id, w//2, h//2)
            time.sleep(1)

        task_file_path = os.path.join(test_file_path, f"trace/{goal}")
        os.makedirs(task_file_path, exist_ok=True)

        _logger = setup_default_logging(os.path.join(task_file_path, 'log.txt'))
        _logger.info(f"\n\n####第{ii + 1}个任务###\n{goal}\n")

        json_output = {
            'instruction': goal,
            'device_id': args.device_id,
            'package_name': episode["package_name"],
            'screen_resolution': [w, h],
            'golden_steps': episode['gold_step'],
            'app_name': app_name,
            'steps': [],
        }

        history_resp: List[str] = []
        images: List[str] = []
        step = 1
        last_action = None
        start_time = time.time()
        step_times = []
        get_xml_times = []
        get_screenshot_times = []
        prepare_data_times = []
        infer_times = []

        total_input_tokens = 0
        total_output_tokens = 0

        while last_action != "finished":
            try:
                if step > args.max_steps:
                    _logger.info("任务执行超出最大步数")
                    break

                _logger.info(f"\n####第{step}步###")
                time.sleep(random.uniform(1, 5))  # 防止大规模测试导致账号被封禁

                # 记录 step 开始时间
                step_start = time.time()

                xml_path = os.path.join(task_file_path, f"{step}.xml")
                get_xml(args.device_id, f"tmp_{args.device_id}.xml")
                if not os.path.exists(xml_path):
                    os.rename(f"tmp_{args.device_id}.xml", xml_path)

                time1 = time.time()
                get_xml_time = time1 - step_start
                _logger.info(f"获取 XML 耗时: {get_xml_time:.2f} 秒")
                get_xml_times.append(get_xml_time)


                screenshot_path = os.path.join(task_file_path, f"{step}.jpg")
                get_screenshot(args.device_id, f"tmp_{args.device_id}.jpg")
                if not os.path.exists(screenshot_path):
                    os.rename(f"tmp_{args.device_id}.jpg", screenshot_path)

                time2 = time.time()
                get_screenshot_time = time2 - time1
                _logger.info(f"获取截图耗时: {get_screenshot_time:.2f} 秒")
                get_screenshot_times.append(get_screenshot_time)

                images.append(screenshot_path)
                msgs = prepare_data(goal, images, history_resp, args.his_num, vllm_format=bool(args.vllm_url))
                batch_item = {"messages": msgs}

                time3 = time.time()
                prepare_data_time = time3 - time2
                _logger.info(f"准备数据耗时: {prepare_data_time:.2f} 秒")
                prepare_data_times.append(prepare_data_time)

                pred_txt, usage = model.infer_batch([batch_item])[0]
                _logger.info("Response: %s", pred_txt)
                history_resp.append(pred_txt)

                time4 = time.time()
                infer_time = time4 - time3
                _logger.info(f"推理耗时: {infer_time:.2f} 秒")
                infer_times.append(infer_time)

                # token统计（假设模型返回有 token 数）
                total_input_tokens += usage.get("prompt_tokens", 0)
                total_output_tokens += usage.get("completion_tokens", 0)

                pred_action = parse_model_output(pred_txt)
                _logger.info("Action: %s", pred_action)
                last_action = pred_action["action_type"]

                execute_actions(args.device_id, [pred_action], w, h)
                step_end = time.time()
                step_time = step_end - step_start
                step_times.append(step_time)
                _logger.info(f"第{step}步耗时: {step_time:.2f} 秒")

                json_output['steps'].append({
                    'step': step,
                    'pred_action': pred_action,
                    'response': pred_txt,
                    'xml_time': get_xml_time,
                    'screenshot_time': get_screenshot_time,
                    'prepare_data_time': prepare_data_time,
                    'infer_time': infer_time,
                    'step_time': step_time,
                    'prompt_tokens': usage.get("prompt_tokens", 0),
                    'completion_tokens': usage.get("completion_tokens", 0),
                })

                with open(os.path.join(task_file_path, "data.json"), 'w', encoding='utf-8') as file:
                    json.dump(json_output, file, ensure_ascii=False)

            except KeyboardInterrupt:
                _logger.info("手动终止任务")
                break
            except Exception as e:
                _logger.info(f"执行异常: {e}")
                step += 1
                continue

            step += 1

        # 清理 logger handler 避免重复打印
        logging.getLogger().handlers.clear()

        # 是否完成判断
        success = False
        verified_elements = episode.get('verified_elements', [])
        success = verify_ui_elements_exist_adb(args.device_id, verified_elements, f"tmp_{args.device_id}.xml")

        json_output.update({
            'model_steps': step - 1,
            'model_self_finish': last_action == "finished",
            'model_success': success,
            'get_xml_times': get_xml_times,
            'get_screenshot_times': get_screenshot_times,
            'prepare_data_times': prepare_data_times,
            'infer_times': infer_times,
            'step_times': step_times,
            'total_time': time.time() - start_time,
            'total_input_tokens': total_input_tokens,
            'total_output_tokens': total_output_tokens
        })

        # 最终写入 result.jsonl
        with open(args.result_path_file, 'a', encoding='utf-8') as fout:
            fout.write(json.dumps(json_output, ensure_ascii=False) + "\n")



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/online_test_trajectory/online_test_trajectory.jsonl')
    parser.add_argument("--img_path", "-img", type=str, default="your_path/online_test_trajectory/data") 
    parser.add_argument("--result_path", "-o", type=str, default='./results')
    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="gpt-4o")
    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)
    parser.add_argument("--device_id", type=str, default="6a2591bc")
    parser.add_argument("--max_steps", type=int, default=15)
    parser.add_argument("--mode", type=str, default='home', choices=['home', 'other', 'intra'])
    parser.add_argument("--api_base", type=str, default="https://api.openai.com/v1")
    parser.add_argument("--api_key", type=str, default="your key")

    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 None
    evaluate(model, args)
