import threading
from collections import deque
import time
from DrissionPage import ChromiumPage, ChromiumOptions
from DrissionPage._pages.chromium_tab import ChromiumTab
from DrissionPage._units.listener import DataPacket
from typing import Tuple, Optional
import json
from urllib.parse import urlparse, urlunparse

one_turn_max_token = 16000


class DPProxyClient:
    def __init__(self, driver: ChromiumTab, packet_filter: dict, self_kill=False):
        self.tab_id = driver.tab_id
        self.driver = ChromePageProxy(driver, self)
        self.thread = None
        self.self_kill = self_kill
        self.packet_filter = packet_filter
        self.packet_queue = deque()

    def get_driver(self, start_listen, count=None, timeout=10) -> ChromiumTab:
        """
        获取代理后的driver、tab
        :param start_listen: 若你自己写的代码里已经使用自动化框架监听发包的功能了，则该值应该置为False。若没监听，则必须将该值置为True
        :param count: 需捕获的数据包总数，为None表示无限
        :param timeout: 两个数据包之间等待的最大时长（秒），为None表示无限，默认为10秒
        :return:
        """
        if start_listen:
            self.driver.listen.set_targets(res_type=('xhr', 'fetch'))
            self.driver.listen.start()
            self.thread = threading.Thread(target=self.start_listen, args=(count, timeout,))
            self.thread.start()
        return self.driver

    def start_listen(self, count=None, timeout=10):
        for _ in self.driver.listen.steps(count=count, timeout=timeout, gap=1):
            pass

    # 每次调用函数，都从队列的左端弹出一个数据包
    def pop_first_packet(self):
        if self.packet_queue:
            result = self.packet_queue.popleft()
            current_queue_size = len(self.packet_queue)
            return current_queue_size, json.dumps(result, ensure_ascii=False, separators=(',', ':')).replace("\\", "")
        else:
            return 0, None


class DPProxyClientManager:
    """浏览器池管理器 - 使用单例模式"""
    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance.tab_pool = {}
        return cls._instance

    def create_client(self, tab: ChromiumTab, packet_filter: dict, self_kill=False) -> Tuple[
        str, DPProxyClient, ChromiumTab]:
        """创建新的tab页面代理实例"""
        client = DPProxyClient(tab, packet_filter, self_kill=self_kill)
        tab = client.get_driver(True, timeout=60 * 10)
        tab_id = tab.tab_id
        self.tab_pool[tab_id] = {"client": client, "driver": tab}
        return tab_id, client, tab

    def get_client(self, tab_id: str) -> Optional[DPProxyClient]:
        """根据端口获取浏览器实例"""
        return self.tab_pool.get(tab_id).get("client", None)

    def remove_client(self, tab_id: str) -> Tuple[bool, Optional[ChromiumPage]]:
        """根据端口移除浏览器实例"""
        client = self.tab_pool.pop(tab_id, None)
        return client is not None, client

    def list_clients(self) -> list[int]:
        """列出所有活跃的浏览器端口"""
        return list(self.tab_pool.keys())


class ChromePageProxy:
    def __init__(self, page, client=None):
        self.__dict__['page'] = page
        self.__dict__['client'] = client

    def __getattr__(self, item):
        attr = getattr(self.page, item)
        print(item, attr)
        if item == 'listen':
            listen_proxy = DrissionPageListenerProxy(attr, self.__dict__['client'])
            return listen_proxy
        return attr


class DrissionPageListenerProxy:
    def __init__(self, listener, client=None):
        self.listener = listener
        self.client = client

    def __getattr__(self, item):
        attr = getattr(self.listener, item)
        # 当监听到wait被调用的时候
        if item == "wait":
            def wrapper(*args, **kwargs):
                result = attr(*args, **kwargs)
                check_data_packet(result, self.client)
                return result

            return wrapper
        # 当监听到steps被调用的时候
        if item == "steps":
            def wrapper(*args, **kwargs):
                if kwargs.get("gap", 1) > 1:
                    raise Exception("暂不支持多包监控")
                result = attr(*args, **kwargs)
                if attr.__name__ == "steps":
                    for step in result:
                        check_data_packet(step, self.client)
                        yield step

            return wrapper
        return attr


def check_data_packet(packet: DataPacket, client: DPProxyClient):
    """
    封装监听到的数据包，并将其存放在client的packet_queue中
    :param packet:
    :param client:
    :return:
    """
    url = packet.url
    method = packet.request.method
    data = None
    if packet.request.hasPostData:
        data = packet.request.postData
    domain = urlparse(url).netloc
    body = packet.response.body
    if isinstance(body, dict):
        body_str = json.dumps(body, ensure_ascii=False, separators=(',', ':'))
    else:
        body_str = str(body)
    body_str_list = [body_str[i:i + one_turn_max_token] for i in range(0, len(body_str), one_turn_max_token)]
    body_completed = True
    packet_filter = client.packet_filter
    domain_filter = packet_filter.get("domain_filter", None)
    method_filter = packet_filter.get("method_filter", ["GET", "POST"])
    for index, body_str in enumerate(body_str_list):
        # 如果给了domain_filter并且domain没有在domain_filter中时跳过该数据包
        if domain_filter and domain not in domain_filter:
            continue
        # 如果method没有在method_filter中，则跳过该数据包
        if method not in method_filter:
            continue
        if (index + 1) != len(body_str_list):
            body_completed = False
        try:
            response_headers = packet.response.headers
        except TypeError:
            response_headers = {}
        temp_dict = {
            "url": url,
            "body_completed": body_completed,
            "method": method,
            "request_data": data,
            "request_headers": dict(packet.request.headers),
            "response_headers": dict(response_headers),
            "response_body_segment": body_str.replace("\\", ""),
        }
        client.packet_queue.append(temp_dict)


client_manager = DPProxyClientManager()
