"""
    Drissionpage-mcp的基础功能
"""
import hashlib
import json
import os
import time
from typing import Any

from DrissionPage.errors import NoRectError
from fastmcp import FastMCP

from tools.browser_manager import BrowserManager
from tools.browser_proxy import DPProxyClientManager
from tools.tools import dp_mcp_message_pack, compress_html_js, compress_image_bytes


def register_visit_url(mcp: FastMCP, browser_manager: BrowserManager, client_manager: DPProxyClientManager):
    @mcp.tool(name="visit_url",
              description="使用Drissionpage打开url访问某个网站，并开始监听初始tab页的所有的XHR请求"
                          "当需要使用手机版浏览器Ua时use_mobile_user_agent为True"
                          "如果想要以域名对packet进行过滤，可以传入想要过滤的域名列表。默认是：None。"
                          "如果想要以method对packet进行过滤，可以传入想要过滤的method列表，默认是：['GET', 'POST']")
    async def visit_url(url: str, domain_filter: list = None, method_filter: list = ["GET", "POST"],
                        use_mobile_user_agent: bool = False) -> dict[str, Any]:
        mobile_user_agent = None
        if use_mobile_user_agent:
            mobile_user_agent = "Mozilla/5.0 (Linux; Android 6.0; Nexus 5 Build/MRA58N) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Mobile Safari/537.36"
        port, _browser = browser_manager.create_browser(mobile_user_agent)
        tab = _browser.get_tab()
        packet_filter = {
            "domain_filter": domain_filter,
            "method_filter": method_filter,
        }
        client_manager.create_client(tab, packet_filter)
        tab.get(url)
        tab_id = tab.tab_id
        return dp_mcp_message_pack(
            f"已在[{port}]端口创建浏览器对象，并已打开链接：{url}，打开的模式是：{'手机版' if use_mobile_user_agent else '电脑版'}",
            tab_id=tab_id,
            browser_port=port
        )


def register_get_new_tab(mcp: FastMCP, browser_manager, client_manager: DPProxyClientManager):
    @mcp.tool(name="get_new_tab",
              description="使用Drissionpage创建一个新的tab页，在新的tab页中打开url，并开始监听新的tab页的所有XHR请求"
                          "如果想要以域名对packet进行过滤，可以传入想要过滤的域名列表。默认是：None。"
                          "如果想要以method对packet进行过滤，可以传入想要过滤的method列表，默认是：['GET', 'POST']")
    async def get_new_tab(browser_port: int, url: str, domain_filter: list = None,
                          method_filter: list = ["GET", "POST"]) -> dict[str, Any]:
        _browser = browser_manager.get_browser(browser_port)
        tab = _browser.new_tab()
        packet_filter = {
            "domain_filter": domain_filter,
            "method_filter": method_filter,
        }
        client_manager.create_client(tab, packet_filter)
        tab.get(url)
        _browser.activate_tab(tab)
        tab_id = tab.tab_id
        return dp_mcp_message_pack(f"已创建新的tab页，并打开链接：{url}", tab_id=tab_id)


def register_pop_first_packet(mcp: FastMCP, browser_manager, client_manager: DPProxyClientManager):
    @mcp.tool(name="pop_first_packet",
              description="每调用一次就会弹出传入的tab页所监听到的数据包中的第一个packet_message，当一个packet_message的response body过长时会被切分成多个包，具体一个请求是否还有下一个包，可以参考body_completed字段")
    async def pop_first_packet(browser_port: int, tab_id: str) -> dict[str, Any]:
        _browser = browser_manager.get_browser(browser_port)
        client = client_manager.get_client(tab_id)
        current_queue_size, packet_message = client.pop_first_packet()
        message = f"tab页:【{tab_id}】，暂时没有监听到XHR数据包"
        if packet_message:
            message = f"tab页:【{tab_id}】，监听到XHR数据包，当前数据包队列中还剩 {current_queue_size} 条数据，如果还剩数据为0，可以暂时稍后再次调用该方法"
        if (packet_message is None) and current_queue_size:
            message = f"tab页:【{tab_id}】，当前弹出的第一个数据包不符合过滤条件，当前数据包队列中还剩 {current_queue_size} 条数据，请不要改变条件，继续弹出下一个数据包"
        return dp_mcp_message_pack(
            message,
            browser_port=browser_port,
            tab_id=tab_id,
            packet_message=packet_message,
            current_queue_size=current_queue_size,
        )


# todo：这个获取html的方式后续需要修改，后面采用的思路应该是：获取全部html，然后提取一个高度抽象化的html结构化对象，然后提供一系列的工具对这个对象进行查询，对象至少包含如下属性和工具
#  属性：
#   1. id【用于通过id查询特定的html元素】
#  工具：
#   1. 根据id查询特定的html元素
#   2. 根据传入的正则表达式去html全文检索，如果匹配到则返回所有匹配到的元素id【最好可以给出除id以外的其他更准确的信息，否则ai在查询出来一大堆元素后需要逐个元素调用1.的方法会大幅拖慢处理速度】
def register_get_html(mcp: FastMCP, browser_manager):
    @mcp.tool(name="get_html", description="使用Drissionpage获取某一个tab页的html")
    async def get_html(browser_port: int, tab_id: str) -> dict[str, Any]:
        _browser = browser_manager.get_browser(browser_port)
        tab = _browser.get_tab(tab_id)
        file_name_prefix = hashlib.md5(str(tab.title).encode('utf-8')).hexdigest()

        # min_html, compress_rate = compress_html(tab.html)
        min_html = tab.run_js(compress_html_js)
        # html_str_list = [min_html[i:i + one_turn_max_token] for i in range(0, len(min_html), one_turn_max_token)]
        html_file_list = []
        for index, html_str in enumerate([min_html]):
            file_name = file_name_prefix + f"_{tab_id}_segment{index}.html"
            abs_path = os.path.join(browser_manager.get_cache_path(), file_name)
            with open(abs_path, "w", encoding="utf-8") as f:
                f.write(html_str)
            html_file_list.append(abs_path)
        message = f"已保存tab页：【{tab_id}】的html源码片段共{len(html_file_list)}个"
        return dp_mcp_message_pack(message, tab_id=tab_id, htmls_local_path=html_file_list)


def register_switch_tab(mcp: FastMCP, browser_manager):
    @mcp.tool(name="switch_tab", description="根据传入的tab_id切换到对应的tab页", )
    async def switch_tab(browser_port: int, tab_id: str) -> dict[str, Any]:
        _browser = browser_manager.get_browser(browser_port)
        _browser.activate_tab(tab_id)
        return dp_mcp_message_pack(f"已将tab页:【{tab_id}】切换至最前端")


def register_close_tab(mcp: FastMCP, browser_manager):
    @mcp.tool(name="close_tab", description="根据传入的tab_id关闭tab页", )
    async def close_tab(browser_port: int, tab_id: str) -> dict[str, Any]:
        _browser = browser_manager.get_browser(browser_port)
        _browser.close_tabs(tab_id)
        return dp_mcp_message_pack(f"已将tab页:【{tab_id}】关闭")


def register_check_selector(mcp: FastMCP, browser_manager):
    @mcp.tool(name="check_selector",
              description="查找tab页中是否包含元素，并以列表的形式返回元素信息"
                          "当要选择的元素包含过多元素时，需要传入offset和page_size来分批查看元素，一般不建议调整page_size，更推荐你调整offset"
                          "同时如果单个元素属性值太长，函数会进行截断。一般的单个元素的属性值超过100个字符的就会触发截断，截断后会在最后拼接'...'")
    async def check_selector(browser_port: int, tab_id: str, css_selector: str, attr_name: str = "text",
                             offset: int = 0, page_size: int = 10) -> dict[str, Any]:
        _browser = browser_manager.get_browser(browser_port)
        target_tab = _browser.get_tab(tab_id)
        slice_seg_len = 100
        exist_flag = False
        output_list = []
        if "css:" not in css_selector:
            css_selector = "css:" + css_selector
        if "css:*" == css_selector:
            return dp_mcp_message_pack(
                message="禁止使用全局匹配的Selector：*，请重新生成更精确的Selector",
                tab_id=tab_id,
                selector=css_selector,
            )
        target_eles = target_tab.eles(css_selector)
        total_ele_len = len(target_eles)
        if len(target_eles) != 0:
            exist_flag = True
        if len(target_eles) > page_size:
            target_eles = target_eles[offset:offset + page_size]
        for index, target_ele in enumerate(target_eles):
            temp_dict = {}
            attr_value = "".join(target_ele.texts()).replace("\n", "")
            if attr_name != "text":
                attr_value = target_ele.attr(attr_name) if target_ele.attr(attr_name) else ""
            is_truncate = len(attr_value) > slice_seg_len
            attr_value = attr_value[:slice_seg_len]
            if is_truncate:
                attr_value += "..."
            temp_dict["element_index"] = offset + index
            temp_dict["attr_name"] = attr_name
            temp_dict["attr_value"] = attr_value
            temp_dict["is_truncate"] = is_truncate
            try:
                has_rect = target_ele.states.has_rect
                is_covered = target_ele.states.is_covered
            except NoRectError:
                has_rect = False
                is_covered = "元素没有大小和位置信息"
            temp_dict["ele_info"] = {
                "元素tag": target_ele.tag,
                "元素是否拥有大小和位置信息": has_rect,
                "元素是否可被模拟点击": target_ele.states.is_clickable,
                "元素是否可见": target_ele.states.is_displayed,
                "元素是否可用": target_ele.states.is_enabled,
                "元素是否被其它元素覆盖": is_covered,
            }
            output_list.append(temp_dict)

        return dp_mcp_message_pack(
            f"已完成tab页:【{tab_id}】对：【{css_selector}】的检查，Selector共选中了 {total_ele_len} 个元素",
            tab_id=tab_id,
            selector=css_selector,
            selector_ele_exist=exist_flag,
            page_size=page_size,
            offset=offset,
            ele_info=output_list,
        )


def register_quit_browser(mcp: FastMCP, browser_manager):
    @mcp.tool(name="quit_browser", description="退出浏览器会话，关闭浏览器")
    async def quit_browser(browser_port: int) -> dict[str, Any]:
        flag, _browser = browser_manager.remove_page(browser_port)
        if flag:
            _browser.quit()
        return dp_mcp_message_pack(
            f"浏览器[{browser_port}]，退出会话，关闭浏览器{'成功' if flag else '失败'}",
            browser_port=browser_port,
            quit_flag=flag
        )


def register_get_ele_info(mcp: FastMCP, browser_manager):
    @mcp.tool(name="get_ele_info", description="返回元素有关一系列的信息。"
                                               "参数说明："
                                               "browser_port：浏览器端口号。"
                                               "tab_id：tab页的id。"
                                               "css_selector：目标选择器"
                                               "element_index：如果选择器定位到的是一个列表【或者可以定位到多个元素】，则该值用于定位这个列表中具体元素，默认值为0【列表中的第一个元素】"
                                               "返回值说明："
                                               "element_tag：此属性返回元素的标签名。"
                                               "element_attrs_key：此属性以list的形式返回元素所有属性的key。"
                                               "element_rect_size：此属性以元组形式返回元素的大小【如果元素没有位置及大小，则返回空元组】。"
                                               "is_in_viewport：此属性以布尔值方式返回元素是否在视口中，以元素可以接受点击的点为判断。"
                                               "is_whole_in_viewport：此属性以布尔值方式返回元素是否整个在视口中。"
                                               "is_alive：此属性以布尔值形式返回当前元素是否仍可用。用于判断是否因页面刷新而导致元素失效。"
                                               "is_checked：此属性以布尔值返回表单单选或多选元素是否选中。"
                                               "is_selected：此属性以布尔值返回<select>元素中的项是否选中。"
                                               "is_enabled：此属性以布尔值返回元素是否可用。"
                                               "is_displayed：此属性以布尔值返回元素是否可见。"
                                               "is_covered：此属性返回元素是否被其它元素覆盖/遮挡。如被覆盖/遮挡，返回覆盖元素的 id【无障碍树的id，无法用于Selector】，否则返回False。"
                                               "is_clickable：此属性返回元素是否可被模拟点击，从是否有大小、是否可用、是否显示、是否响应点击判断，不判断是否被遮挡。"
                                               "has_rect：此属性返回元素是否拥有大小和位置信息，有则返回四个角在页面上的坐标组成的列表，没有则返回False。")
    async def get_ele_info(browser_port: int, tab_id: str, css_selector: str, element_index: int = 0) -> dict[
        str, Any]:
        _browser = browser_manager.get_browser(browser_port)
        target_tab = _browser.get_tab(tab_id)
        css_selector = css_selector
        if "css:" not in css_selector:
            css_selector = "css:" + css_selector
        target_eles = target_tab.eles(css_selector)
        try:
            target_element = target_eles[element_index]
        except IndexError:
            return dp_mcp_message_pack(
                message="报错：IndexError: list index out of range。请检查Selector中是否包含了如：first-child、nth-child等字段，如果有则去掉。必须使用element_index来控制元素的选择",
                browser_port=browser_port,
                tab_id=tab_id,
                css_selector=css_selector,
                element_index=element_index,
            )

        try:
            has_rect = target_element.states.has_rect
            is_covered = target_element.states.is_covered
            element_rect_size = target_element.rect.size
            is_in_viewport = target_element.states.is_in_viewport
            is_whole_in_viewport = target_element.states.is_whole_in_viewport
            rect_x_center = (has_rect[0][0] + has_rect[1][0]) / 2
            rect_y_center = (has_rect[1][1] + has_rect[2][1]) / 2
            rect_center = (rect_x_center, rect_y_center)
        except NoRectError:
            has_rect = False
            is_covered = "元素没有大小和位置信息"
            element_rect_size = "元素没有大小和位置信息"
            rect_center = "元素没有大小和位置信息"
            is_in_viewport = "元素没有大小和位置信息"
            is_whole_in_viewport = "元素没有大小和位置信息"
        try:
            child_count = target_element.child_count
        except Exception:
            child_count = 0
        return dp_mcp_message_pack(
            message="元素可以被正常的选择到，以下是元素的信息",
            browser_port=browser_port,
            tab_id=tab_id,
            css_selector=css_selector,
            element_index=element_index,
            element_tag=target_element.tag,
            element_attrs_key=list(target_element.attrs.keys()),
            element_child_count=child_count,
            element_rect_size=element_rect_size,
            element_rect_center=rect_center,
            is_in_viewport=is_in_viewport,
            is_whole_in_viewport=is_whole_in_viewport,
            is_alive=target_element.states.is_alive,
            is_checked=target_element.states.is_checked,
            is_selected=target_element.states.is_selected,
            is_enabled=target_element.states.is_enabled,
            is_displayed=target_element.states.is_displayed,
            is_covered=is_covered,
            is_clickable=target_element.states.is_clickable,
            has_rect=has_rect
        )


def register_get_screenshot(mcp: FastMCP, browser_manager):
    @mcp.tool(name="get_tab_screenshot",
              description="尝试对传入tab页进行截图，并将截图压缩为1M大小png图片，会返回截图保存路径")
    async def get_tab_screenshot(browser_port: int, tab_id: str) -> dict[str, Any]:
        _browser = browser_manager.get_browser(browser_port)
        target_tab = _browser.get_tab(tab_id)
        target_tab.wait.doc_loaded()
        timestamp = int(time.time() * 1000)
        time.sleep(3)
        origin_png = target_tab.get_screenshot(as_bytes="png")
        compress_png = compress_image_bytes(origin_png)
        image_path = os.path.join(browser_manager.get_cache_path(), f"{browser_port}_{tab_id}_{timestamp}.png")
        with open(image_path, "wb") as f:
            f.write(compress_png)
        return dp_mcp_message_pack(
            message=f"已完成对browser_port={browser_port},tab_id={tab_id}的截屏",
            browser_port=browser_port,
            tab_id=tab_id,
            screenshot_path=image_path
        )
