import asyncio
from types import FunctionType
from typing import Collection, Any
from uuid import UUID

import msgspec


class Package(msgspec.Struct):
    func: FunctionType
    args: Collection
    kwargs: dict


class ResourceManager:
    def __init__(self):
        self.resource_map: dict[UUID, Any] = {}
        self.resource_remap: dict[Any, UUID] = {}
        self.sub_managers: list[ResourceManager] = []

    def register(self, obj: Any, uuid: UUID):
        if obj in self.resource_remap:
            raise ValueError("请勿重复注册资源")
        self.resource_map[uuid] = obj
        self.resource_remap[obj] = uuid
        for man in self.sub_managers:
            man.register(obj, uuid)

    def create_sub_manager(self):
        man = ResourceManager()
        self.sub_managers.append(man)
        man.resource_map = self.resource_map.copy()
        man.resource_remap = self.resource_remap.copy()
        return man


class PackageSender:
    def __init__(self,
                 reader: asyncio.StreamReader,
                 writer: asyncio.StreamWriter,
                 resource_manager: ResourceManager):
        self.reader = reader
        self.writer = writer
        self.resource_manager = resource_manager

    def enc_hook(self, obj):
        uuid = self.resource_manager.resource_remap.get(obj, None)
        if uuid is None:
            raise RuntimeError("请先注册资源")
        return msgspec.msgpack.Ext(0, uuid.bytes)

    def ext_hook(self, code, data):
        if code == 0:
            uuid = UUID(bytes=bytes(data))
            obj = self.resource_manager.resource_map.get(uuid, None)
            if obj is None:
                raise RuntimeError("未定义的资源")
            return obj

    async def _read_data(self):
        size = int.from_bytes(await self.reader.readexactly(4), "big")
        data = await self.reader.readexactly(size)
        return data

    async def _send_data(self, data: bytes):
        size = len(data).to_bytes(4, "big")
        self.writer.write(size + data)
        await self.writer.drain()

    async def read_package(self) -> Package:
        data = await self._read_data()
        return msgspec.msgpack.decode(data, type=Package, ext_hook=self.ext_hook)

    async def send_package(self, package: Package):
        data = msgspec.msgpack.encode(package, enc_hook=self.enc_hook)
        await self._send_data(data)


def get_namespace(obj):
    return obj.__module__ + "." + obj.__name__


def namespace_to_uuid(namespace: str) -> UUID:
    import uuid
    return uuid.uuid5(uuid.NAMESPACE_DNS, namespace)
