"""SagerNet UDP over TCP v2 帧编解码。

UoT 只负责把 UDP 数据报装入已有的可靠 TCP 字节流；外层连接仍由
HTTP CONNECT + TLS 建立。此模块不依赖 sing-box。
"""

from __future__ import annotations

import ipaddress
import socket
import struct
from dataclasses import dataclass

MAGIC_ADDRESS = "sp.v2.udp-over-tcp.arpa"
MAX_DATAGRAM = 65535


@dataclass(frozen=True)
class Address:
    host: str
    port: int


def encode_address(host: str, port: int) -> bytes:
    host = str(host).strip()
    if not 0 <= int(port) <= 65535:
        raise ValueError("端口超出范围")
    try:
        addr = ipaddress.ip_address(host)
    except ValueError:
        raw = host.encode("idna")
        if not 1 <= len(raw) <= 255:
            raise ValueError("域名长度超出 UoT 限制")
        return b"\x02" + bytes([len(raw)]) + raw + struct.pack("!H", int(port))
    if addr.version == 4:
        return b"\x00" + addr.packed + struct.pack("!H", int(port))
    return b"\x01" + addr.packed + struct.pack("!H", int(port))


def decode_address(data: bytes, offset: int = 0) -> tuple[Address, int]:
    if offset >= len(data):
        raise ValueError("UoT 地址头不完整")
    atyp = data[offset]
    offset += 1
    if atyp == 0:
        size = 4
    elif atyp == 1:
        size = 16
    elif atyp == 2:
        if offset >= len(data):
            raise ValueError("UoT 域名长度缺失")
        size = data[offset]
        offset += 1
    else:
        raise ValueError(f"未知 UoT 地址类型: {atyp}")
    end = offset + size
    if end + 2 > len(data):
        raise ValueError("UoT 地址头不完整")
    raw = data[offset:end]
    port = struct.unpack("!H", data[end:end + 2])[0]
    if atyp == 0:
        host = str(ipaddress.IPv4Address(raw))
    elif atyp == 1:
        host = str(ipaddress.IPv6Address(raw))
    else:
        host = raw.decode("idna")
    return Address(host, port), end + 2


def encode_request(is_connect: bool, destination: Address) -> bytes:
    return (b"\x01" if is_connect else b"\x00") + encode_address(
        destination.host, destination.port
    )


def encode_connect_frame(payload: bytes) -> bytes:
    if len(payload) > MAX_DATAGRAM:
        raise ValueError("UoT 数据报超过 65535 字节")
    return struct.pack("!H", len(payload)) + payload


def encode_datagram_frame(destination: Address, payload: bytes) -> bytes:
    if len(payload) > MAX_DATAGRAM:
        raise ValueError("UoT 数据报超过 65535 字节")
    return encode_address(destination.host, destination.port) + struct.pack("!H", len(payload)) + payload


def read_exact(stream, size: int) -> bytes:
    chunks = []
    remaining = size
    while remaining:
        chunk = stream.recv(remaining)
        if not chunk:
            raise ConnectionError("UoT TCP 流已关闭")
        chunks.append(chunk)
        remaining -= len(chunk)
    return b"".join(chunks)


def read_connect_frame(stream) -> bytes:
    size = struct.unpack("!H", read_exact(stream, 2))[0]
    return read_exact(stream, size)


def read_datagram_frame(stream) -> tuple[Address, bytes]:
    # 地址是变长的，只能先读 ATYP，再按类型补齐。
    head = read_exact(stream, 1)
    atyp = head[0]
    if atyp == 0:
        rest = read_exact(stream, 6)
    elif atyp == 1:
        rest = read_exact(stream, 18)
    elif atyp == 2:
        length = read_exact(stream, 1)
        rest = length + read_exact(stream, length[0] + 2)
    else:
        raise ValueError(f"未知 UoT 地址类型: {atyp}")
    address, _ = decode_address(head + rest)
    size = struct.unpack("!H", read_exact(stream, 2))[0]
    return address, read_exact(stream, size)