"""
这个文件中提供的工具作为独立的Drissionpage mcp工具
"""
import hashlib
import json
import os
import time
from typing import Any

from fastmcp import FastMCP

from tools.browser_manager import BrowserManager
from tools.tools import compress_html, requests_html, dp_headless_html, assert_waf_cookie, dp_mcp_message_pack, \
    compress_html_js, compress_image_bytes
from tools.browser_proxy import DPProxyClient, DPProxyClientManager

html_source_code_local_save_path = os.path.join(os.getcwd(), "html-source-code")
waf_status_code_dict = {
    412: "瑞数",
    521: "加速乐"
}
# 一轮最大输入，以免单个html最大长度超过ai最大输入
one_turn_max_token = 8000


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,
        )


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()
        if not os.path.exists(html_source_code_local_save_path):
            os.makedirs(html_source_code_local_save_path)
        # 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(html_source_code_local_save_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页中是否包含元素，并返回元素attr_name所对应的值。"
                          "当要选择的元素包含过多元素时，需要传入offset和page_size来分批查看元素，一般不建议调整page_size，更推荐你调整offset"
                          "同时如果单个元素属性值太长，函数会进行截断。一般的单个元素的属性值超过300个字符的就会触发截断，截断后会在最后拼接'...'")
    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)
        css_selector = css_selector
        if "css:" not in css_selector:
            css_selector = "css:" + css_selector
        target_eles = target_tab.eles(css_selector)
        exist_flag = False
        if len(target_eles) != 0:
            exist_flag = True
        if len(target_eles) > page_size:
            target_eles = target_eles[offset:offset + page_size]
        slice_seg = max(300, one_turn_max_token // (page_size + 6))
        if attr_name == "text":
            ele_attr_list = [i.text.replace("\n", "") for i in target_eles]
            ele_attr_list = [attr_str[:slice_seg] for attr_str in ele_attr_list]
            # 如果经过截断遍历后的字符串长度与截断长度相等，则默认截断了
            ele_attr_list = [attr_str + "..." if len(attr_str) == slice_seg else attr_str for attr_str in ele_attr_list]
            attr_output = "\n".join(ele_attr_list)
        else:
            ele_attr_list = [i.attr(attr_name) for i in target_eles]
            ele_attr_list = [attr_str[:slice_seg] for attr_str in ele_attr_list if attr_str]
            ele_attr_list = [attr_str + "..." if len(attr_str) == slice_seg else attr_str for attr_str in ele_attr_list]
            attr_output = json.dumps(ele_attr_list, ensure_ascii=False)
        # 对attr_output逐个截断，截断的长度为：一轮最大token除以元素个数+3个点+两个引号和逗号
        return dp_mcp_message_pack(
            f"已完成tab页:【{tab_id}】对：【{css_selector}】的检查",
            tab_id=tab_id,
            selector=css_selector,
            selector_ele_exist=exist_flag,
            page_size=page_size,
            offset=offset,
            attr_output=attr_output
        )


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_assert_waf(mcp: FastMCP, browser_manager):
    @mcp.tool(name="assert_waf",
              description="通过对比requests、有头浏览器、无头浏览器获取到的html，判断网页是否使用了waf以及是否为动态渲染的网页")
    async def assert_waf(browser_port: int, tab_id: str) -> dict[str, Any]:
        _browser = browser_manager.get_browser(browser_port)
        target_tab = _browser.get_tab(tab_id)
        recommend_team = "drissionpage_head"
        head_cookies = target_tab.cookies()
        # 通过cookie判断是否有waf
        waf_flag, waf_type = assert_waf_cookie(head_cookies)
        head_html = target_tab.html
        min_head_html, head_rate = compress_html(head_html, only_text=True)
        raw_html, status_code = requests_html(target_tab.url)
        min_raw_html, raw_rate = compress_html(raw_html, only_text=True)
        r_h_rate_diff = abs(head_rate - raw_rate)
        # 如果有已知的防火墙，则不浪费时间使用无头获取html和压缩比了
        if waf_flag or status_code in waf_status_code_dict.keys():
            return dp_mcp_message_pack(
                f"已完成tab页:【{tab_id}】的分析，该tab页存在waf",
                tab_id=tab_id,
                recommend_team=recommend_team,
                raw_head_rate_difference=r_h_rate_diff,
                raw_headless_rate_difference=0,
                head_headless_rate_difference=0
            )

        headless_html = dp_headless_html(target_tab.url)
        min_headless_html, headless_rate = compress_html(headless_html, only_text=True)
        r_hless_rate_diff = abs(raw_rate - headless_rate)
        h_hless_rate_diff = abs(head_rate - headless_rate)
        # 最优情况：requests，dp有头和无头拿到的结果基本一致，认定为没有防护的静态网页
        if r_h_rate_diff < 40 and r_hless_rate_diff < 40 and h_hless_rate_diff < 40:
            recommend_team = "requests"
        # 最差情况：requests，dp有头和无头拿到的结果差距都很大，认定为有浏览器无头检测+动态网页
        # if r_h_rate_diff < 40 and r_hless_rate_diff < 40 and h_hless_rate_diff < 40:
        # 较差1：dp有头和无头差距很小，但是requests拿不到正确结果，认定为有requests防护 or 动态网页
        elif h_hless_rate_diff < 30 and r_hless_rate_diff > 40:
            recommend_team = "drissionpage_headless"
        # 较差2：有头和无头差距很大，但是requests和有头拿到的结果基本一致，认定为要么有别的没有防护requests的waf，或者间歇性的瑞数【此时应该拿有头的cookie去判断其中是否有瑞数特征，上面已经做了】
        # if r_h_rate_diff < 15 and h_hless_rate_diff > 40:
        return dp_mcp_message_pack(
            f"已完成tab页:【{tab_id}】的分析，该tab页存在waf",
            tab_id=tab_id,
            recommend_team=recommend_team,
            raw_head_rate_difference=r_h_rate_diff,
            raw_headless_rate_difference=h_hless_rate_diff,
            head_headless_rate_difference=h_hless_rate_diff
        )


def register_click_action(mcp: FastMCP, browser_manager):
    @mcp.tool(name="click_action", description="尝试点击tab页中的元素，返回元素是否可以被点击，以及是否点击成功。")
    async def click_action(browser_port: int, tab_id: str, css_selector: str) -> 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)
        click_success = False
        element_clickable = False
        if len(target_eles) == 1:
            target_element = target_eles[0]
            element_clickable = target_element.states.is_clickable
            try:
                target_element.click()
                click_success = True
            except Exception as e:
                click_success = False
            message = f"tab页:【{tab_id}】点击【{css_selector}】 {'成功' if click_success else '失败'} 了"
        else:
            message = f"tab页:【{tab_id}】传入的css_selector找到了{len(target_eles)}个元素，请确保传入的css_selector可以找到唯一的一个元素"
        return dp_mcp_message_pack(
            message=message,
            browser_port=browser_port,
            tab_id=tab_id,
            css_selector=css_selector,
            element_clickable=element_clickable,
            click_success=click_success,
            extra_message="点击成功，页面可能有更新，请重新获取页面html，并重新分析页面Selector" if click_success else ""
        )


def register_scroll_action(mcp: FastMCP, browser_manager):
    @mcp.tool(name="scroll_action", description="尝试滚动tab页"
                                                "forward参数是滚动的方向：down、up、left、right"
                                                "pixel参数是滚动的像素值，默认为None。"
                                                "当forward为down且pixel为None，则将页面滚动到垂直中间位置，水平位置不变"
                                                "当forward为up且pixel为None，则将页面滚动到顶部，水平位置不变"
                                                "当forward为left且pixel为None，则将页面滚动到最左边，垂直位置不变"
                                                "当forward为right且pixel为None，则将页面滚动到最右边，垂直位置不变")
    async def scroll_action(browser_port: int, tab_id: str, forward: str = "down", pixel: int = None) -> dict[str, Any]:
        _browser = browser_manager.get_browser(browser_port)
        target_tab = _browser.get_tab(tab_id)
        if forward == "down":
            if pixel is None:
                target_tab.scroll.to_half()
            target_tab.scroll.down(pixel)
        elif forward == "up":
            if pixel is None:
                target_tab.scroll.to_top()
            target_tab.scroll.up(pixel)
        elif forward == "left":
            if pixel is None:
                target_tab.scroll.to_leftmost()
            target_tab.scroll.left(pixel)
        elif forward == "right":
            if pixel is None:
                target_tab.scroll.to_rightmost()
            target_tab.scroll.right(pixel)
        else:
            if pixel is None:
                target_tab.scroll.to_half()
            target_tab.scroll.down()
        message = f"已完成对tab页:【{tab_id}】forward={forward} 的滑动"
        return dp_mcp_message_pack(
            message=message,
            browser_port=browser_port,
            tab_id=tab_id,
        )


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)
        if not os.path.exists(html_source_code_local_save_path):
            os.makedirs(html_source_code_local_save_path)
        timestamp = int(time.time() * 1000)
        origin_png = target_tab.get_screenshot(as_bytes="png")
        compress_png = compress_image_bytes(origin_png)
        image_path = os.path.join(html_source_code_local_save_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
        )
