"""
虎牙 Wup 协议 Python 实现
基于前端 WUP 实现逆向

Wup 包结构:
  [4字节长度][Wup body]
  Wup body = tag1:iVersion, tag2:cPacketType, tag3:iMessageType,
             tag4:iRequestId, tag5:sServantName, tag6:sFuncName,
             tag7:sBuffer(bytes), tag8:iTimeout, tag9:context(map), tag10:status(map)
  sBuffer = Map<"tReq", 编码后的请求结构体>
"""
import struct
from typing import Any, Dict, Optional
from .taf_protocol import TafOutputStream, TafInputStream, TafType, TafStruct


class WupRequest:
    """Wup 请求对象"""

    def __init__(self):
        self.iVersion: int = 3          # tag 1
        self.cPacketType: int = 0       # tag 2
        self.iMessageType: int = 0      # tag 3
        self.iRequestId: int = 0        # tag 4
        self.sServantName: str = ""     # tag 5
        self.sFuncName: str = ""        # tag 6
        self.sBuffer: bytes = b''       # tag 7
        self.iTimeout: int = 3000       # tag 8
        self.context: Dict[str, str] = {}   # tag 9
        self.status: Dict[str, str] = {}    # tag 10
        self.newdata: Dict[str, bytes] = {}

    def setServant(self, name: str):
        self.sServantName = name

    def setFunc(self, name: str):
        self.sFuncName = name

    def setRequestId(self, req_id: int):
        self.iRequestId = req_id

    def writeStruct(self, key: str, struct_data):
        """
        写入请求参数到 newdata[key]

        Args:
            key: 键名（通常是 "tReq"）
            struct_data: TafStruct 对象 / dict / list
        """
        os = TafOutputStream()

        if isinstance(struct_data, TafStruct) or hasattr(struct_data, 'write_to'):
            # 结构体对象：STRUCT_BEGIN + 内容 + STRUCT_END
            os.write_struct(0, struct_data)
        elif isinstance(struct_data, dict):
            os.write_struct_begin(0)
            for tag_id, (field_name, field_value) in enumerate(struct_data.items()):
                self._write_field(os, tag_id, field_value)
            os.write_struct_end()
        elif isinstance(struct_data, (list, tuple)):
            # wsLaunch 等用 list 参数
            os.write_head(0, TafType.LIST)
            os.write_int32(0, len(struct_data))
            for i, item in enumerate(struct_data, start=1):
                self._write_field(os, i, item)
        else:
            raise TypeError(f"不支持的 struct_data 类型: {type(struct_data)}")

        self.newdata[key] = os.get_bytes()

    def _write_field(self, os: TafOutputStream, tag: int, value: Any):
        if isinstance(value, bool):
            os.write_boolean(tag, value)
        elif isinstance(value, int):
            os.write_int64(tag, value)
        elif isinstance(value, float):
            os.write_double(tag, value)
        elif isinstance(value, str):
            os.write_string(tag, value)
        elif isinstance(value, bytes):
            os.write_bytes(tag, value)
        elif isinstance(value, dict):
            os.write_map(tag, value)
        elif isinstance(value, (list, tuple)):
            os.write_list(tag, list(value))
        elif hasattr(value, 'write_to'):
            os.write_struct(tag, value)
        else:
            raise TypeError(f"不支持的字段类型: {type(value)}")

    def encode(self) -> bytes:
        """编码为完整 Wup 包（含长度前缀）"""
        # 1. newdata -> Map (tag 0)
        data_os = TafOutputStream()
        data_os.write_head(0, TafType.MAP)
        data_os.write_int32(0, len(self.newdata))
        for k, v in self.newdata.items():
            data_os.write_string(0, k)
            data_os.write_bytes(1, v)
        self.sBuffer = data_os.get_bytes()

        # 2. Wup 头
        wup_os = TafOutputStream()
        wup_os.write_int16(1, self.iVersion)
        wup_os.write_int8(2, self.cPacketType)
        wup_os.write_int32(3, self.iMessageType)
        wup_os.write_int32(4, self.iRequestId)
        wup_os.write_string(5, self.sServantName)
        wup_os.write_string(6, self.sFuncName)
        wup_os.write_bytes(7, self.sBuffer)
        wup_os.write_int32(8, self.iTimeout)
        wup_os.write_map(9, self.context)
        wup_os.write_map(10, self.status)
        wup_body = wup_os.get_bytes()

        # 3. 长度前缀
        length = 4 + len(wup_body)
        return struct.pack('>I', length) + wup_body


def normalize_wup_payload(data: bytes) -> bytes:
    """去掉可选的 4 字节 WUP 长度前缀，返回裸 WUP body"""
    if len(data) < 4:
        return data
    declared_len = struct.unpack('>I', data[0:4])[0]
    if declared_len == len(data) or declared_len + 4 == len(data):
        return data[4:]
    return data


class WupResponse:
    """Wup 响应对象"""

    def __init__(self):
        self.iVersion: int = 0
        self.cPacketType: int = 0
        self.iMessageType: int = 0
        self.iRequestId: int = 0
        self.sServantName: str = ""
        self.sFuncName: str = ""
        self.sBuffer: bytes = b''
        self.iTimeout: int = 0
        self.context: Dict[str, str] = {}
        self.status: Dict[str, str] = {}
        self.newdata: Dict[str, bytes] = {}

    def decode(self, data: bytes):
        """解码响应（不包含长度前缀；若含前缀会自动跳过）"""
        data = normalize_wup_payload(data)

        ins = TafInputStream(data)
        # 按字段读，直到结束
        try:
            while True:
                tag, dtype = ins.peek_head()
                if tag > 10:
                    break
                if tag == 1:
                    ins.read_head()
                    self.iVersion = ins._read_int_value(dtype)
                elif tag == 2:
                    ins.read_head()
                    self.cPacketType = ins._read_int_value(dtype)
                elif tag == 3:
                    ins.read_head()
                    self.iMessageType = ins._read_int_value(dtype)
                elif tag == 4:
                    ins.read_head()
                    self.iRequestId = ins._read_int_value(dtype)
                elif tag == 5:
                    ins.read_head()
                    self.sServantName = _read_string_value(ins, dtype)
                elif tag == 6:
                    ins.read_head()
                    self.sFuncName = _read_string_value(ins, dtype)
                elif tag == 7:
                    ins.read_head()
                    self.sBuffer = _read_bytes_value(ins, dtype)
                    self._decode_buffer()
                elif tag == 8:
                    ins.read_head()
                    self.iTimeout = ins._read_int_value(dtype)
                elif tag == 9:
                    # context map
                    ins.read_head()
                    self.context = _read_map_value(ins, dtype)
                elif tag == 10:
                    ins.read_head()
                    self.status = _read_map_value(ins, dtype)
                    break
                else:
                    ins.read_head()
                    ins.skip_field(dtype)
        except EOFError:
            pass

    def _decode_buffer(self):
        """解析 sBuffer 里的 newdata Map"""
        if not self.sBuffer:
            return
        ins = TafInputStream(self.sBuffer)
        try:
            tag, dtype = ins.read_head()
            if tag == 0 and dtype == TafType.MAP:
                count = ins._read_int_len()
                for _ in range(count):
                    _, kt = ins.read_head()
                    key = _read_string_value(ins, kt)
                    _, vt = ins.read_head()
                    val = _read_bytes_value(ins, vt)
                    self.newdata[key] = val
        except Exception as e:
            print(f"[WupResponse] 解析 newdata 失败: {e}")

    def readStruct(self, key: str, struct_class=None):
        """
        读取响应结构体

        Args:
            key: "tRsp" / "tResp" / "tReq"（响应里通常是 tRsp）
            struct_class: 结构体类（需实现 read_from），None 则返回原始 bytes
        """
        data = self.newdata.get(key)
        if not data and key == "tRsp":
            data = self.newdata.get("tResp")
        if not data and key == "tResp":
            data = self.newdata.get("tRsp")
        if not data:
            return None

        if struct_class is None:
            return data

        ins = TafInputStream(data)
        # newdata 里的结构体以 STRUCT_BEGIN 开头
        try:
            tag, dtype = ins.peek_head()
            if dtype == TafType.STRUCT_BEGIN:
                ins.read_head()  # 消费 STRUCT_BEGIN
                obj = struct_class()
                obj.read_from(ins)
                return obj
        except Exception as e:
            print(f"[WupResponse] 解析 {struct_class.__name__} 失败: {e}")
        return None


# ============================================================
# 辅助：按已知 dtype 读取值
# ============================================================

def _read_string_value(ins: TafInputStream, dtype: int) -> str:
    if dtype == TafType.STRING1:
        length = struct.unpack('B', ins.buf.read(1))[0]
    elif dtype == TafType.STRING4:
        length = struct.unpack('>I', ins.buf.read(4))[0]
    else:
        return ""
    return ins.buf.read(length).decode('utf-8', errors='replace')


def _read_bytes_value(ins: TafInputStream, dtype: int) -> bytes:
    if dtype != TafType.SIMPLE_LIST:
        return b''
    ins.read_head()  # 元素类型 INT8
    length = ins._read_int_len()
    return ins.buf.read(length)


def _read_map_value(ins: TafInputStream, dtype: int) -> Dict:
    if dtype != TafType.MAP:
        return {}
    count = ins._read_int_len()
    result = {}
    for _ in range(count):
        _, kt = ins.read_head()
        k = _read_string_value(ins, kt)
        _, vt = ins.read_head()
        v = _read_string_value(ins, vt) if vt in (TafType.STRING1, TafType.STRING4) else ""
        result[k] = v
    return result