import base64
import json
from unittest.mock import MagicMock

import pytest

from skytale_sdk.bridge import ProtocolBridge
from skytale_sdk.envelope import Envelope, Protocol


def mock_manager():
    mgr = MagicMock()
    mgr.send_envelope = MagicMock()
    mgr.receive_envelopes = MagicMock(return_value=[])
    return mgr


def make_envelope(protocol, payload, content_type="application/json", metadata=None):
    if isinstance(payload, str):
        payload = payload.encode("utf-8")
    return Envelope(protocol, content_type, payload, metadata=metadata)


class TestTranslateIdentity:
    """Same-protocol translation should just retag."""

    def test_same_protocol_passthrough(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        env = make_envelope(Protocol.A2A, '{"parts":[]}')
        result = bridge._translate(env, Protocol.A2A)
        assert result.protocol == Protocol.A2A
        assert result.payload == env.payload

    def test_unknown_pair_passthrough(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        env = make_envelope(Protocol.ACP, '{"data":1}')
        result = bridge._translate(env, Protocol.ANP)
        assert result.protocol == Protocol.ANP
        assert result.payload == env.payload


class TestA2AToSlim:
    def test_wraps_with_slim_content_type(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        env = make_envelope(Protocol.A2A, '{"parts":[]}')
        result = bridge._a2a_to_slim(env)
        assert result.protocol == Protocol.SLIM
        assert result.content_type == "application/a2a+json"
        assert result.payload == env.payload

    def test_adds_payload_type_metadata(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        env = make_envelope(Protocol.A2A, '{"parts":[]}')
        result = bridge._a2a_to_slim(env)
        assert result.metadata["payload_type"] == "a2a"

    def test_preserves_existing_metadata(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        env = make_envelope(Protocol.A2A, '{}', metadata={"custom": "val"})
        result = bridge._a2a_to_slim(env)
        assert result.metadata["custom"] == "val"
        assert result.metadata["payload_type"] == "a2a"


class TestSlimToA2A:
    def test_passthrough_a2a_json(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        env = make_envelope(
            Protocol.SLIM, '{"parts":[]}',
            content_type="application/a2a+json",
        )
        result = bridge._slim_to_a2a(env)
        assert result.protocol == Protocol.A2A
        assert result.content_type == "application/json"
        assert result.payload == env.payload

    def test_wraps_non_a2a_content(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        env = make_envelope(
            Protocol.SLIM, '{"key":"value"}',
            content_type="application/json",
        )
        result = bridge._slim_to_a2a(env)
        assert result.protocol == Protocol.A2A
        a2a = json.loads(result.payload)
        assert a2a["role"] == "agent"
        assert len(a2a["parts"]) == 1
        assert a2a["parts"][0]["type"] == "data"


class TestMCPToSlim:
    def test_wraps_with_slim_content_type(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        rpc = json.dumps({"jsonrpc": "2.0", "method": "tools/list", "id": 1})
        env = make_envelope(Protocol.MCP, rpc)
        result = bridge._mcp_to_slim(env)
        assert result.protocol == Protocol.SLIM
        assert result.content_type == "application/mcp+json"

    def test_extracts_method_metadata(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        rpc = json.dumps({"jsonrpc": "2.0", "method": "tools/call", "id": 1})
        env = make_envelope(Protocol.MCP, rpc)
        result = bridge._mcp_to_slim(env)
        assert result.metadata["mcp_method"] == "tools/call"


class TestSlimToMCP:
    def test_passthrough_mcp_json(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        rpc = json.dumps({"jsonrpc": "2.0", "method": "tools/list"})
        env = make_envelope(
            Protocol.SLIM, rpc,
            content_type="application/mcp+json",
        )
        result = bridge._slim_to_mcp(env)
        assert result.protocol == Protocol.MCP
        assert result.payload == env.payload

    def test_wraps_non_mcp_as_notification(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        env = make_envelope(
            Protocol.SLIM, '{"data":1}',
            content_type="application/json",
        )
        result = bridge._slim_to_mcp(env)
        rpc = json.loads(result.payload)
        assert rpc["jsonrpc"] == "2.0"
        assert rpc["method"] == "message"
        assert "id" not in rpc  # notification, no id
        # params.data is base64 of the original payload
        decoded = base64.b64decode(rpc["params"]["data"])
        assert json.loads(decoded) == {"data": 1}


class TestA2AToMCP:
    def test_creates_tools_call(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        a2a = json.dumps({
            "messageId": "msg-1",
            "role": "agent",
            "parts": [{"type": "text", "text": "hello"}],
        })
        env = make_envelope(Protocol.A2A, a2a)
        result = bridge._a2a_to_mcp(env)
        rpc = json.loads(result.payload)
        assert rpc["method"] == "tools/call"
        assert rpc["id"] == "msg-1"
        assert rpc["params"]["name"] == "a2a_message"

    def test_non_json_wraps_as_data(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        env = make_envelope(Protocol.A2A, b"\x00\x01\x02",
                           content_type="application/octet-stream")
        result = bridge._a2a_to_mcp(env)
        rpc = json.loads(result.payload)
        assert rpc["method"] == "tools/call"


class TestMCPToA2A:
    def test_result_to_a2a(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        rpc = json.dumps({
            "jsonrpc": "2.0",
            "id": 1,
            "result": {
                "content": [{"type": "text", "text": "hello world"}],
            },
        })
        env = make_envelope(Protocol.MCP, rpc)
        result = bridge._mcp_to_a2a(env)
        a2a = json.loads(result.payload)
        assert a2a["role"] == "agent"
        assert a2a["parts"][0]["type"] == "text"
        assert a2a["parts"][0]["text"] == "hello world"

    def test_request_to_a2a(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        rpc = json.dumps({
            "jsonrpc": "2.0",
            "method": "tools/call",
            "params": {"name": "search"},
        })
        env = make_envelope(Protocol.MCP, rpc)
        result = bridge._mcp_to_a2a(env)
        a2a = json.loads(result.payload)
        assert a2a["parts"][0]["data"]["mcp_method"] == "tools/call"


class TestANPSlimBridge:
    def test_anp_to_slim(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        env = make_envelope(Protocol.ANP, '{"agent":"a1"}')
        result = bridge._anp_to_slim(env)
        assert result.protocol == Protocol.SLIM
        assert result.content_type == "application/anp+json"
        assert result.metadata["payload_type"] == "anp"

    def test_slim_to_anp_passthrough(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        env = make_envelope(
            Protocol.SLIM, '{"agent":"a1"}',
            content_type="application/anp+json",
        )
        result = bridge._slim_to_anp(env)
        assert result.protocol == Protocol.ANP
        assert result.payload == env.payload

    def test_slim_to_anp_wraps(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        env = make_envelope(
            Protocol.SLIM, '{"data":1}',
            content_type="application/json",
        )
        result = bridge._slim_to_anp(env)
        parsed = json.loads(result.payload)
        assert parsed["type"] == "bridged"


class TestLMOSSlimBridge:
    def test_lmos_to_slim(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        env = make_envelope(Protocol.LMOS, '{"task":"run"}')
        result = bridge._lmos_to_slim(env)
        assert result.protocol == Protocol.SLIM
        assert result.content_type == "application/lmos+json"
        assert result.metadata["payload_type"] == "lmos"

    def test_slim_to_lmos_passthrough(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        env = make_envelope(
            Protocol.SLIM, '{"task":"run"}',
            content_type="application/lmos+json",
        )
        result = bridge._slim_to_lmos(env)
        assert result.protocol == Protocol.LMOS
        assert result.payload == env.payload


class TestNLIPSlimBridge:
    def test_nlip_to_slim(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        env = make_envelope(Protocol.NLIP, '{"msg":"hi"}')
        result = bridge._nlip_to_slim(env)
        assert result.protocol == Protocol.SLIM
        assert result.content_type == "application/nlip+json"
        assert result.metadata["payload_type"] == "nlip"

    def test_slim_to_nlip_passthrough(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        env = make_envelope(
            Protocol.SLIM, '{"msg":"hi"}',
            content_type="application/nlip+json",
        )
        result = bridge._slim_to_nlip(env)
        assert result.protocol == Protocol.NLIP
        assert result.payload == env.payload

    def test_slim_to_nlip_wraps_with_submessages(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        env = make_envelope(
            Protocol.SLIM, '{"data":1}',
            content_type="application/json",
        )
        result = bridge._slim_to_nlip(env)
        parsed = json.loads(result.payload)
        assert parsed["type"] == "bridged"
        assert len(parsed["submessages"]) == 1
        assert parsed["submessages"][0]["subMessageType"] == "content"


class TestTranslateDispatch:
    """Test that _translate dispatches to correct translator."""

    def test_all_registered_pairs(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        expected_pairs = [
            (Protocol.A2A, Protocol.SLIM),
            (Protocol.SLIM, Protocol.A2A),
            (Protocol.MCP, Protocol.SLIM),
            (Protocol.SLIM, Protocol.MCP),
            (Protocol.A2A, Protocol.MCP),
            (Protocol.MCP, Protocol.A2A),
            (Protocol.ANP, Protocol.SLIM),
            (Protocol.SLIM, Protocol.ANP),
            (Protocol.LMOS, Protocol.SLIM),
            (Protocol.SLIM, Protocol.LMOS),
            (Protocol.NLIP, Protocol.SLIM),
            (Protocol.SLIM, Protocol.NLIP),
        ]
        assert len(bridge._translators) == len(expected_pairs)
        for pair in expected_pairs:
            assert pair in bridge._translators


class TestBridgeLifecycle:
    def test_stop_sets_flag(self):
        mgr = mock_manager()
        bridge = ProtocolBridge(mgr)
        assert bridge._running is True
        bridge.stop()
        assert bridge._running is False
