"""Inline Protobuf schema support for the ServiceBridge Python SDK.

Allows handlers to declare typed input/output schemas that are compiled at
runtime into binary Protobuf message descriptors.  Payloads are then encoded
as binary Protobuf instead of JSON, and the schema JSON is published to the
control plane for registry introspection.

Usage::

    from service_bridge import ServiceBridge, RpcSchemaOpts, RpcFieldDef

    svc = ServiceBridge(...)
    svc.register_schema("order.create", RpcSchemaOpts(
        input={
            "order_id": RpcFieldDef(type="string", id=1),
            "amount":   RpcFieldDef(type="int64",  id=2),
        },
        output={
            "success": RpcFieldDef(type="bool",   id=1),
            "message": RpcFieldDef(type="string", id=2),
        },
    ))
"""

from __future__ import annotations

import json
import threading
from typing import Any, Optional

# ─── Type-to-descriptor mapping ───────────────────────────────────────────────

_PROTO_FIELD_TYPES: dict[str, int] = {
    # Values match google.protobuf.descriptor_pb2.FieldDescriptorProto.Type.*
    "double":   1,
    "float":    2,
    "int64":    3,
    "uint64":   4,
    "int32":    5,
    "fixed64":  6,
    "fixed32":  7,
    "bool":     8,
    "string":   9,
    "bytes":    12,
    "uint32":   13,
    "sint32":   17,
    "sint64":   18,
}

_LABEL_OPTIONAL = 1
_LABEL_REPEATED = 3


# ─── Schema cache ─────────────────────────────────────────────────────────────

_cache_lock = threading.Lock()
_descriptor_cache: dict[str, Any] = {}  # cache_key -> MessageDescriptor


# ─── Public compile/encode/decode API ────────────────────────────────────────

def compile_schema(msg_name: str, schema: dict) -> Any:
    """Compile an RpcSchema dict into a protobuf MessageDescriptor.

    The descriptor is cached by ``msg_name`` to avoid repeated compilation.

    Args:
        msg_name: Unique name for the message (used as the protobuf message name).
        schema: Mapping of field names to RpcFieldDef instances.

    Returns:
        A ``google.protobuf.descriptor.Descriptor`` instance.

    Raises:
        ImportError: If ``protobuf`` is not installed.
        ValueError: If a field type is unknown.
    """
    with _cache_lock:
        if msg_name in _descriptor_cache:
            return _descriptor_cache[msg_name]

    descriptor = _build_descriptor(msg_name, schema)
    with _cache_lock:
        _descriptor_cache[msg_name] = descriptor
    return descriptor


def encode_with_schema(descriptor: Any, payload: Any) -> bytes:
    """Encode a Python dict/object to binary Protobuf using ``descriptor``.

    Args:
        descriptor: A compiled ``google.protobuf.descriptor.Descriptor``.
        payload: A JSON-serialisable Python object (dict expected).

    Returns:
        Binary Protobuf bytes.
    """
    from google.protobuf import message_factory as _mf
    from google.protobuf import json_format

    msg_class = _get_message_class(descriptor)
    msg = json_format.ParseDict(payload if isinstance(payload, dict) else {}, msg_class())
    return msg.SerializeToString()


def decode_with_schema(descriptor: Any, data: bytes) -> Any:
    """Decode binary Protobuf bytes to a Python dict using ``descriptor``.

    Args:
        descriptor: A compiled ``google.protobuf.descriptor.Descriptor``.
        data: Binary Protobuf bytes.

    Returns:
        Python dict with the decoded message fields.
    """
    from google.protobuf import json_format

    msg_class = _get_message_class(descriptor)
    msg = msg_class()
    msg.ParseFromString(data)
    return json_format.MessageToDict(msg, preserving_proto_field_name=True)


# ─── Internal helpers ─────────────────────────────────────────────────────────

def _build_descriptor(msg_name: str, schema: dict) -> Any:
    """Construct a FileDescriptorProto and register it, returning the message Descriptor."""
    from google.protobuf import descriptor_pb2
    from google.protobuf import descriptor_pool as _pool_mod
    from google.protobuf import symbol_database

    file_name = f"sb_dyn_{msg_name}.proto"
    package = "sb_dyn"
    full_name = f"{package}.{msg_name}"

    # Build the list of FieldDescriptorProtos.
    fields = []
    for field_name, field_def in schema.items():
        type_id = _PROTO_FIELD_TYPES.get(field_def.type)
        if type_id is None:
            raise ValueError(f"Unknown Protobuf field type {field_def.type!r} for field {field_name!r}")
        label = _LABEL_REPEATED if field_def.repeated else _LABEL_OPTIONAL
        fields.append(
            descriptor_pb2.FieldDescriptorProto(
                name=field_name,
                number=field_def.id,
                type=type_id,
                label=label,
            )
        )

    fdp = descriptor_pb2.FileDescriptorProto(
        name=file_name,
        package=package,
        syntax="proto3",
        message_type=[
            descriptor_pb2.DescriptorProto(name=msg_name, field=fields)
        ],
    )

    pool = _pool_mod.DescriptorPool()
    pool.Add(fdp)
    return pool.FindMessageTypeByName(full_name)


# Per-descriptor message class cache.
_class_cache: dict[int, Any] = {}
_class_cache_lock = threading.Lock()


def _get_message_class(descriptor: Any) -> Any:
    """Return (and cache) the generated message class for a descriptor."""
    key = id(descriptor)
    with _class_cache_lock:
        if key in _class_cache:
            return _class_cache[key]

    from google.protobuf import message_factory

    # ``MessageFactory`` API changed across protobuf versions.
    try:
        # protobuf >= 4.21 (4.x / 5.x)
        cls = message_factory.GetMessageClass(descriptor)
    except AttributeError:
        # Older protobuf
        factory = message_factory.MessageFactory()
        cls = factory.GetPrototype(descriptor)

    with _class_cache_lock:
        _class_cache[key] = cls
    return cls
