"""Loader and helpers for the versioned gRPC Shard protocol (ADR-0024, DGR-002).

The ``.proto`` schema at ``packages/node/native/proto/shard_runtime.proto`` is the
single source of truth. Rather than commit generated stubs (which pin a protobuf
runtime version and drift from the schema), this package generates the Python
stubs on demand into a gitignored build directory and imports them. Generation is
reproducible: it shells out to the pinned ``grpc_tools.protoc`` with the exact
same flags as ``packages/node/native/scripts/generate_python.py``.

Typical use::

    from meshnet_node import native_protocol as proto
    pb2 = proto.load()
    header = pb2.MessageHeader(work_id="w1", route_session_id="s1")

The checksum/fragment helpers encode the bounded-fragment tensor-bundle semantics
so callers (and DGR-008/DGR-009) do not re-derive them.
"""

from __future__ import annotations

import hashlib
import importlib
import importlib.util
import pathlib
import sys
import threading
import types
import zlib

# The wire schema version this build targets. Keep in sync with the
# ``SCHEMA_VERSION_1`` enum member in the .proto.
SCHEMA_VERSION = 1

_NATIVE_ROOT = pathlib.Path(__file__).resolve().parents[2] / "native"
PROTO_DIR = _NATIVE_ROOT / "proto"
PROTO_FILE = PROTO_DIR / "shard_runtime.proto"
# ``build/`` is globally gitignored, so generated stubs never enter version control.
GEN_DIR = _NATIVE_ROOT / "build" / "python"

_PB2_MODULE = "shard_runtime_pb2"
_GRPC_MODULE = "shard_runtime_pb2_grpc"

# Reentrant: load_grpc() holds the lock and calls load(), which re-acquires it.
_lock = threading.RLock()
_cached_pb2: types.ModuleType | None = None
_cached_grpc: types.ModuleType | None = None


class ProtocGenerationError(RuntimeError):
    """Raised when the protobuf stubs cannot be generated from the schema."""


def _needs_regen(target: pathlib.Path) -> bool:
    if not target.exists():
        return True
    try:
        return PROTO_FILE.stat().st_mtime > target.stat().st_mtime
    except OSError:
        return True


def generate(*, force: bool = False) -> pathlib.Path:
    """Generate ``shard_runtime_pb2{,_grpc}.py`` into :data:`GEN_DIR`.

    Returns the output directory. Reproducible and idempotent: regenerates only
    when the schema is newer than the stubs (or ``force`` is set). Requires the
    pinned ``grpc_tools`` (available in the project ``.venv``).
    """
    if not PROTO_FILE.exists():
        raise ProtocGenerationError(f"schema not found: {PROTO_FILE}")

    pb2_path = GEN_DIR / f"{_PB2_MODULE}.py"
    if not force and not _needs_regen(pb2_path):
        return GEN_DIR

    try:
        from grpc_tools import protoc
    except ImportError as exc:  # pragma: no cover - environment-dependent
        raise ProtocGenerationError(
            "grpc_tools is required to generate the Shard protocol stubs; "
            "install grpcio-tools (present in the project .venv)."
        ) from exc

    GEN_DIR.mkdir(parents=True, exist_ok=True)
    well_known = _well_known_include()
    args = [
        "grpc_tools.protoc",
        f"-I{PROTO_DIR}",
        *([f"-I{well_known}"] if well_known else []),
        f"--python_out={GEN_DIR}",
        f"--grpc_python_out={GEN_DIR}",
        str(PROTO_FILE.name),
    ]
    # protoc resolves the proto by name relative to -I, so run with PROTO_DIR
    # semantics by passing the bare filename plus the include path above.
    rc = protoc.main([a for a in args])
    if rc != 0:
        raise ProtocGenerationError(
            f"grpc_tools.protoc exited with status {rc} for {PROTO_FILE}"
        )
    if not pb2_path.exists():  # pragma: no cover - defensive
        raise ProtocGenerationError(f"protoc did not produce {pb2_path}")
    return GEN_DIR


def _well_known_include() -> str | None:
    """Bundled well-known .proto include dir shipped with grpc_tools, if any."""
    try:
        import grpc_tools

        candidate = pathlib.Path(grpc_tools.__file__).parent / "_proto"
        return str(candidate) if candidate.is_dir() else None
    except Exception:  # pragma: no cover - defensive
        return None


def _import_generated(module_name: str) -> types.ModuleType:
    gen_dir = str(GEN_DIR)
    if gen_dir not in sys.path:
        sys.path.insert(0, gen_dir)
    if module_name in sys.modules:
        return sys.modules[module_name]
    return importlib.import_module(module_name)


def load(*, force: bool = False) -> types.ModuleType:
    """Return the generated ``shard_runtime_pb2`` module (messages only).

    Generates the stubs on first use. Thread-safe and cached. Does not import
    grpc; message serialization/round-trip needs only this module.
    """
    global _cached_pb2
    with _lock:
        if _cached_pb2 is not None and not force:
            return _cached_pb2
        generate(force=force)
        _cached_pb2 = _import_generated(_PB2_MODULE)
        return _cached_pb2


def load_grpc(*, force: bool = False) -> types.ModuleType:
    """Return the generated ``shard_runtime_pb2_grpc`` module (service stubs).

    Requires the ``grpc`` runtime. Use for building the C++/Python worker; the
    round-trip/compat tests only need :func:`load`.
    """
    global _cached_grpc
    with _lock:
        if _cached_grpc is not None and not force:
            return _cached_grpc
        generate(force=force)
        load()  # ensure the _pb2 module the grpc stub imports is present
        _cached_grpc = _import_generated(_GRPC_MODULE)
        return _cached_grpc


# ---------------------------------------------------------------------------
# Checksum + bounded-fragment helpers (shared bundle semantics)
# ---------------------------------------------------------------------------

# Algorithm-name strings mirror the ChecksumAlgorithm enum members without
# importing the generated module (so this table is usable before load()).
_CHECKSUM_CRC32C = "CHECKSUM_CRC32C"
_CHECKSUM_CRC32 = "CHECKSUM_CRC32"
_CHECKSUM_SHA256 = "CHECKSUM_SHA256"
_CHECKSUM_NONE = "CHECKSUM_NONE"


def _crc32c(data: bytes) -> int:
    """Castagnoli CRC32C (software table). Deterministic, no external deps."""
    crc = 0xFFFFFFFF
    for byte in data:
        crc ^= byte
        for _ in range(8):
            crc = (crc >> 1) ^ (0x82F63B78 & -(crc & 1))
    return crc ^ 0xFFFFFFFF


def compute_checksum(algorithm: int, data: bytes):
    """Build a ``Checksum`` message for ``data`` under the given enum value.

    ``algorithm`` is a ``ChecksumAlgorithm`` enum int from the generated module.
    Uses only the standard library (crc32c software table, zlib.crc32, hashlib).
    """
    pb2 = load()
    name = pb2.ChecksumAlgorithm.Name(algorithm)
    if name == _CHECKSUM_SHA256:
        value = hashlib.sha256(data).digest()
    elif name == _CHECKSUM_CRC32C:
        value = _crc32c(data).to_bytes(4, "big")
    elif name == _CHECKSUM_CRC32:
        value = (zlib.crc32(data) & 0xFFFFFFFF).to_bytes(4, "big")
    elif name == _CHECKSUM_NONE:
        value = b""
    else:
        raise ValueError(f"unsupported checksum algorithm: {name}")
    return pb2.Checksum(algorithm=algorithm, value=value)


def verify_checksum(checksum, data: bytes) -> bool:
    """True if ``checksum`` matches ``data`` (CHECKSUM_NONE always verifies)."""
    pb2 = load()
    if checksum.algorithm in (0, pb2.CHECKSUM_NONE):
        return True
    return compute_checksum(checksum.algorithm, data).value == checksum.value


def fragment_tensor(
    *,
    name: str,
    shape,
    dtype: int,
    payload: bytes,
    byte_order: int | None = None,
    max_fragment_bytes: int = 1 << 20,
    compression: int | None = None,
    checksum_algorithm: int | None = None,
):
    """Build a :class:`NamedTensor` splitting ``payload`` into bounded fragments.

    Fragments are ordered by ``byte_offset`` and each carries an optional
    per-fragment checksum. ``payload`` is treated as already compressed if
    ``compression`` is set; this helper does not compress (that is the seam's
    policy in ``activation_compression``), it only frames.
    """
    if max_fragment_bytes <= 0:
        raise ValueError("max_fragment_bytes must be positive")
    pb2 = load()
    if byte_order is None:
        byte_order = pb2.BYTE_ORDER_LITTLE_ENDIAN
    if compression is None:
        compression = pb2.COMPRESSION_NONE

    chunks = [
        payload[i : i + max_fragment_bytes]
        for i in range(0, len(payload), max_fragment_bytes)
    ] or [b""]
    fragments = []
    offset = 0
    for index, chunk in enumerate(chunks):
        frag = pb2.TensorFragment(
            fragment_index=index,
            fragment_count=len(chunks),
            byte_offset=offset,
            data=chunk,
        )
        if checksum_algorithm is not None:
            frag.checksum.CopyFrom(compute_checksum(checksum_algorithm, chunk))
        fragments.append(frag)
        offset += len(chunk)
    return pb2.NamedTensor(
        name=name,
        shape=list(shape),
        dtype=dtype,
        byte_order=byte_order,
        total_byte_length=len(payload),
        compression=compression,
        fragments=fragments,
    )


def reassemble_tensor(named_tensor) -> bytes:
    """Concatenate a :class:`NamedTensor`'s fragments back into the full payload.

    Validates fragment ordering, total length, and any per-fragment checksums.
    """
    fragments = sorted(named_tensor.fragments, key=lambda f: f.byte_offset)
    out = bytearray()
    for frag in fragments:
        if frag.byte_offset != len(out):
            raise ValueError(
                f"non-contiguous fragment at offset {frag.byte_offset} "
                f"(expected {len(out)})"
            )
        if frag.HasField("checksum") and not verify_checksum(frag.checksum, frag.data):
            raise ValueError(f"fragment {frag.fragment_index} checksum mismatch")
        out.extend(frag.data)
    if named_tensor.total_byte_length and len(out) != named_tensor.total_byte_length:
        raise ValueError(
            f"reassembled length {len(out)} != declared "
            f"{named_tensor.total_byte_length}"
        )
    return bytes(out)


__all__ = [
    "SCHEMA_VERSION",
    "PROTO_FILE",
    "PROTO_DIR",
    "GEN_DIR",
    "ProtocGenerationError",
    "generate",
    "load",
    "load_grpc",
    "compute_checksum",
    "verify_checksum",
    "fragment_tensor",
    "reassemble_tensor",
]