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 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 resize_if_needed(img: Image.Image, max_side: int = 1120) -> Image.Image:
    w, h = img.size
    if max(w, h) <= max_side:
        return img
    scale = max_side / max(w, h)
    return img.resize((int(w * scale), int(h * scale)),
                      resample=Image.Resampling.LANCZOS)


def encode_image(image):
    
    img = resize_if_needed(Image.open(image))
    buf = io.BytesIO()
    img.save(buf, format="png")
    
    return base64.b64encode(buf.getvalue()).decode()


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, gt) -> 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(resp: str) -> Dict[str, Union[str, list, None]]:
    try:
        resp_dict = json.loads(resp)
    except json.JSONDecodeError:
        return {"action_type": None, "action_param": None}

    # 优先判断是否为 finished 类型
    if resp_dict.get("STATUS") == "finish":
        return {"action_type": "finished", "action_param": None}

    # 滑动动作
    if "to" in resp_dict and resp_dict["to"] in ["up", "down", "left", "right"]:
        return {"action_type": "scroll", "action_param": resp_dict["to"]}

    # 点击动作
    if "POINT" in resp_dict and isinstance(resp_dict["POINT"], list) and len(resp_dict["POINT"]) >= 2:
        return {
            "action_type": "click",
            "action_param": tuple(resp_dict["POINT"][:2])
        }

    # 输入动作
    if "TYPE" in resp_dict:
        return {
            "action_type": "input",
            "action_param": resp_dict["TYPE"]
        }
    
    if "HOME" in resp_dict:
        return {"action_type": "press_home", "action_param": None}

    if "BACK" in resp_dict:
        return {"action_type": "press_back", "action_param": None}

    # 兜底：如果都没有匹配，可能是等待、无效响应或需要进一步思考
    return {"action_type": None, "action_param": None}



def data_action_to_prompt(action, w, h):
    """
    根据动作数据（如 action_type, action_param）反向生成模型 response 字符串。
    生成格式应与 parse_model_output() 中的解析逻辑保持一致。
    """
    resp = {}

    atype = action.get("action_type")
    param = action.get("action_param")

    # 点击类：生成 POINT 坐标
    if atype == "click" and isinstance(param, (list, tuple)) and len(param) == 2:
        resp["thought"] = "执行点击操作"
        resp["POINT"] = [int(param[0]), int(param[1])]

    # 输入类：生成 TYPE
    elif atype == "input" and isinstance(param, str):
        resp["thought"] = f"输入关键词 {param}"
        resp["TYPE"] = param

    # 滑动类：生成方向
    elif atype == "scroll" and param in ["up", "down", "left", "right"]:
        resp["thought"] = f"执行向{param}滑动操作"
        resp["to"] = param

    # 结束类动作：生成状态标记
    elif atype == "finished":
        resp["thought"] = "任务已完成"
        resp["STATUS"] = "finish"

    # 无效动作或未预测到动作
    elif atype is None:
        resp["thought"] = "未识别到动作"
        resp["STATUS"] = "impossible"

    else:
        resp["thought"] = "未知动作类型"
        # 不设置 POINT/TYPE/to，保持 minimal

    return json.dumps(resp, 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]:
    ACTION_SCHEMA = json.load(open('./schema.json', encoding="utf-8"))
    items = list(ACTION_SCHEMA.items())
    insert_index = 3
    items.insert(insert_index, ("required", ["thought"])) # enable/disable thought by setting it to "required"/"optional"
    ACTION_SCHEMA = dict(items)
    
    SYSTEM_PROMPT = f'''# Role
你是一名熟悉安卓系统触屏GUI操作的智能体，将根据用户的问题，分析当前界面的GUI元素和布局，生成相应的操作。

# Task
针对用户问题，根据输入的当前屏幕截图，输出下一步的操作。

# Rule
- 以紧凑JSON格式输出
- 输出操作必须遵循Schema约束

# Schema
{json.dumps(ACTION_SCHEMA, indent=None, ensure_ascii=False, separators=(',', ':'))}'''

    messages = [{"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]}]
    if vllm_format:
        img_idx = 0
        total_history = len(history_responses)

        for h_idx, hist in enumerate(history_responses):
            if h_idx + his_num + 1 > total_history:
                messages.append(
                    {"role": "user", 
                    "content": [{"type": "text", "text": f"<Question>{instruction}</Question>\n当前屏幕截图：(<image>./</image>)"}, 
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(images[img_idx])}"}}]})
                img_idx += 1
                messages.append({"role": "assistant", "content": hist})

        # 当前截图
        messages.append(
            {"role": "user", "content": [
                {"type": "text", "text": f"<Question>{instruction}</Question>\n当前屏幕截图：(<image>./</image>)"}, 
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(images[img_idx])}"}}
                ]
                })

    return messages


# -------------------------- Model Wrappers ------------------------------- #

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,
                    top_p=0.7,
                    frequency_penalty=1.0,
                )
                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



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])
        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/pub_version/offline_agent.jsonl')
    parser.add_argument("--img_path", "-img", type=str, default="your_path/GUI-CEval/img/offline_agent") 
    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="agentcpm")
    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'])
    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)
