#
# Copyright (c) 2000, 2099, trustbe and/or its affiliates. All rights reserved.
# TRUSTBE PROPRIETARY/CONFIDENTIAL. Use is subject to license terms.
#
#

from typing import Dict, Any, Callable, Coroutine

import mesh.log as log
import mesh.tool as tool
from mesh.kinds import Principal
from mesh.macro import spi
from mesh.mpc import ServiceProxy, Mesh
from mesh.prsim import Transport, Session, Metadata


@spi("mesh")
class MeshTransport(Transport):
    """
    Mesh transport is stateless transport depends on mesh server.
    """

    def __init__(self):
        self.sessions: Dict[str, Session] = dict()

    async def open(self, session_id: str, metadata: Dict[str, str]) -> Session:
        skey = self.session_key(session_id, metadata)
        session = self.sessions.get(skey, None)
        if not session:
            attachments = Mesh.context().get_attachments()
            addr = Mesh.context().get_attribute(Mesh.REMOTE)
            principal = Mesh.context().get_principals().peek()
            session = MeshSession(session_id, metadata, attachments, addr, principal, lambda: self.finalize(skey))
            if self.sessions.__len__() < 500:
                self.sessions[skey] = session
        return session

    async def close(self, timeout: int):
        for k, v in self.sessions.items():
            try:
                await v.release(timeout, "")
            except BaseException as e:
                log.error("", e)
        self.sessions.clear()

    async def roundtrip(self, payload: bytes, metadata: Dict[str, str]) -> bytes:
        """ bid """
        pass

    def finalize(self, session_key: str):
        self.sessions.__delitem__(session_key)

    @staticmethod
    def session_key(session_id: str, metadata: Dict[str, str]) -> str:
        target_node_id = Metadata.MESH_TARGET_NODE_ID.get(metadata)
        target_inst_id = Metadata.MESH_TARGET_INST_ID.get(metadata)
        return f"{session_id}.{target_node_id}.{target_inst_id}"


@spi("mesh")
class MeshSession(Session):

    def __init__(self, session_id: str, metadata: Dict[str, str], attachments: Dict[str, str], address: str,
                 principal: Principal, finalizer: Callable):
        self.session_id = session_id
        self.metadata = metadata
        self.attachments = attachments
        self.address = address
        self.principal = principal
        self.session = ServiceProxy.default_proxy(Session)
        self.finalizer = finalizer

    async def with_context(self, fn: Coroutine, remote: bool, timeout: int) -> Any:
        return await Mesh.context_safe(self.context_safe(fn, remote, timeout))

    async def context_safe(self, fn: Any, remote: bool, timeout: int) -> Any:
        if self.attachments:
            for k, v in self.attachments.items():
                Mesh.context().get_attachments()[k] = v

        if self.metadata:
            for k, v in self.metadata.items():
                Mesh.context().get_attachments()[k] = v

        if tool.required(self.address):
            Mesh.context().set_attribute(Mesh.REMOTE, self.address)

        if timeout > 0:
            Mesh.context().set_attribute(Mesh.TIMEOUT, timeout)

        Metadata.MESH_SESSION_ID.set(Mesh.context().get_attachments(), self.session_id)

        if not remote:
            return await fn

        target_node_id = Metadata.MESH_TARGET_NODE_ID.get(Mesh.context().get_attachments())
        target_inst_id = Metadata.MESH_TARGET_INST_ID.get(Mesh.context().get_attachments())
        if '' == target_node_id and '' == target_inst_id:
            return await fn
        try:
            principal = Principal()
            principal.node_id = target_node_id
            principal.inst_id = target_inst_id
            Mesh.context().get_principals().append(principal)
            return await fn
        finally:
            Mesh.context().get_principals().pop()

    async def peek(self, topic: str = "") -> bytes:
        return await self.with_context(self.session.peek(topic), False, 0)

    async def pop(self, timeout: int, topic: str = "") -> bytes:
        return await self.with_context(self.session.pop(timeout, topic), False, timeout)

    async def push(self, payload: bytes, metadata: Dict[str, str], topic: str = ""):
        return await self.with_context(self.session.push(payload, metadata, topic), True, 0)

    async def release(self, timeout: int, topic: str = ""):
        self.finalizer()
        return await self.with_context(self.session.release(timeout, topic), False, timeout)
