"""Native llama.cpp/GGUF backend adapter for Meshnet node startup.

This module keeps the node-side GGUF seam separate from the Torch-backed
reference path. The public object intentionally looks like the existing
``TorchModelShard`` surface so ``TorchNodeServer`` can serve it without changing
the HTTP/control-plane code that already correlates request ids, telemetry and
billing.

The transport layer is intentionally explicit:

* direct worker calls are expected to use the versioned gRPC Shard protocol
  from :mod:`meshnet_node.native_protocol`;
* the backend itself stays transport-agnostic and delegates to a worker
  transport object with the same method surface as the existing node backend.

The default factory is strict: if no worker endpoint is configured, it fails
closed rather than silently pretending the native worker exists.
"""

from __future__ import annotations

import os
from dataclasses import dataclass, field
from types import SimpleNamespace
from typing import Any, Protocol, runtime_checkable

from .model_backend import (
    MissingModelDependencyError,
    ModelBackendError,
    TailTokenResult,
    TensorPayload,
)

_BACKEND_ID = "llama.cpp"


@runtime_checkable
class NativeWorkerTransport(Protocol):
    """Backend-shaped transport for the supervised native worker."""

    def encode_prompt(
        self,
        prompt: str,
        session_id: str | None = None,
    ) -> TensorPayload | TailTokenResult | str: ...

    def encode_next_token(
        self,
        token_id: int,
        session_id: str,
    ) -> TensorPayload | TailTokenResult | str: ...

    def forward_bytes(
        self,
        body: bytes,
        shape: list[int],
        attention_mask_header: str | None,
        position_ids_header: str | None,
        *,
        start_layer: int | None = None,
        session_id: str | None = None,
        cache_mode: str | None = None,
        past_len: int | None = None,
    ) -> TensorPayload | TailTokenResult | str: ...

    def decode_tail_token(self, hidden_states: Any) -> TailTokenResult: ...

    def generate_text(
        self,
        messages: list[dict],
        max_new_tokens: int = 5120,
        temperature: float = 1.0,
        top_p: float = 1.0,
    ) -> str: ...

    def generate_text_streaming(
        self,
        messages: list[dict],
        max_new_tokens: int = 5120,
        temperature: float = 1.0,
        top_p: float = 1.0,
    ): ...

    def count_prompt_tokens(self, messages: list[dict]) -> int: ...

    def count_text_tokens(self, text: str) -> int: ...

    def eos_token_ids(self) -> list[int]: ...

    def release_session(self, session_id: str) -> None: ...


@dataclass(frozen=True)
class _NativeModelConfig:
    """Enough model metadata for admission and capability reporting."""

    model_type: str = "llama"
    architecture_adapter: str = "dense-llama"
    num_hidden_layers: int = 1
    torch_dtype: str = "bfloat16"

    def to_dict(self) -> dict[str, Any]:
        return {
            "model_type": self.model_type,
            "architecture_adapter": self.architecture_adapter,
            "num_hidden_layers": self.num_hidden_layers,
            "torch_dtype": self.torch_dtype,
        }


@dataclass
class GgufNodeBackend:
    """GGUF shard backend shaped like ``TorchModelShard``.

    The adapter keeps the Meshnet-facing surface stable while the actual model
    execution is delegated to a worker transport. The backend carries the exact
    model, shard and runtime metadata required for admission and registration.
    """

    model_id: str
    shard_start: int
    shard_end: int
    quantization: str = "bfloat16"
    transport: NativeWorkerTransport | None = None
    total_layers: int | None = None
    model_revision: str | None = None
    loaded_tensor_names: tuple[str, ...] = ()
    device_type: str = "cpu"
    supports_kv_cache: bool = True
    worker_url: str | None = None
    architecture_adapter: str = "dense-llama"
    tokenizer_revision: str | None = None
    runtime_recipe_fingerprint: str | None = None
    _model: SimpleNamespace = field(init=False, repr=False)
    _tokenizer: SimpleNamespace = field(init=False, repr=False)
    is_head: bool = field(init=False)
    is_tail: bool = field(init=False)
    loaded_shard_start: int = field(init=False)
    loaded_shard_end: int = field(init=False)
    owns_embedding: bool = field(init=False)
    owns_final_head: bool = field(init=False)

    backend_id = _BACKEND_ID

    def __post_init__(self) -> None:
        if self.shard_start < 0 or self.shard_end < self.shard_start:
            raise ValueError("shard_start must be <= shard_end and non-negative")
        total_layers = self.total_layers or (self.shard_end + 1)
        object.__setattr__(
            self,
            "total_layers",
            int(total_layers),
        )
        object.__setattr__(
            self,
            "_model",
            SimpleNamespace(
                revision=self.model_revision or self.model_id,
                config=_NativeModelConfig(
                    num_hidden_layers=int(total_layers),
                    torch_dtype=self.quantization,
                ),
            ),
        )
        object.__setattr__(
            self,
            "_tokenizer",
            SimpleNamespace(
                model_id=self.model_id,
                revision=self.tokenizer_revision or self.model_revision or self.model_id,
                eos_token="",
                eos_token_id=[],
            ),
        )
        object.__setattr__(self, "is_head", self.shard_start == 0)
        object.__setattr__(self, "is_tail", self.shard_end >= int(total_layers) - 1)
        object.__setattr__(self, "loaded_shard_start", self.shard_start)
        object.__setattr__(self, "loaded_shard_end", self.shard_end)
        object.__setattr__(self, "owns_embedding", self.is_head)
        object.__setattr__(self, "owns_final_head", self.is_tail)
        if not self.loaded_tensor_names:
            object.__setattr__(
                self,
                "loaded_tensor_names",
                self._default_tensor_inventory(),
            )

    @property
    def model(self) -> Any:
        return self._model

    @property
    def tokenizer(self) -> Any:
        return self._tokenizer

    @property
    def device(self) -> SimpleNamespace:
        return SimpleNamespace(type=self.device_type)

    @property
    def shard_range(self) -> tuple[int, int]:
        return self.shard_start, self.shard_end

    def encode_prompt(self, prompt: str, session_id: str | None = None) -> TensorPayload | TailTokenResult | str:
        return self._transport().encode_prompt(prompt, session_id=session_id)

    def encode_next_token(self, token_id: int, session_id: str) -> TensorPayload | TailTokenResult | str:
        return self._transport().encode_next_token(token_id, session_id)

    def forward_bytes(
        self,
        body: bytes,
        shape: list[int],
        attention_mask_header: str | None,
        position_ids_header: str | None,
        start_layer: int | None = None,
        session_id: str | None = None,
        cache_mode: str | None = None,
        past_len: int | None = None,
    ) -> TensorPayload | TailTokenResult | str:
        return self._transport().forward_bytes(
            body,
            shape,
            attention_mask_header,
            position_ids_header,
            start_layer=start_layer,
            session_id=session_id,
            cache_mode=cache_mode,
            past_len=past_len,
        )

    def decode_tail(self, hidden_states: Any) -> str:
        return self.decode_tail_token(hidden_states).text

    def decode_tail_token(self, hidden_states: Any) -> TailTokenResult:
        return self._transport().decode_tail_token(hidden_states)

    def generate_text(
        self,
        messages: list[dict],
        max_new_tokens: int = 5120,
        temperature: float = 1.0,
        top_p: float = 1.0,
    ) -> str:
        return self._transport().generate_text(messages, max_new_tokens, temperature, top_p)

    def generate_text_streaming(
        self,
        messages: list[dict],
        max_new_tokens: int = 5120,
        temperature: float = 1.0,
        top_p: float = 1.0,
    ):
        yield from self._transport().generate_text_streaming(messages, max_new_tokens, temperature, top_p)

    def count_prompt_tokens(self, messages: list[dict]) -> int:
        return self._transport().count_prompt_tokens(messages)

    def count_text_tokens(self, text: str) -> int:
        return self._transport().count_text_tokens(text)

    def eos_token_ids(self) -> list[int]:
        return self._transport().eos_token_ids()

    def release_session(self, session_id: str) -> None:
        self._transport().release_session(session_id)

    def _transport(self) -> NativeWorkerTransport:
        if self.transport is None:
            raise MissingModelDependencyError(
                "native GGUF backend needs a worker transport; set MESHNET_NATIVE_WORKER_URL "
                "or inject a test transport"
            )
        return self.transport

    def _default_tensor_inventory(self) -> tuple[str, ...]:
        tensor_names = [f"blk.{layer}.weight" for layer in range(self.shard_start, self.shard_end + 1)]
        if self.is_head:
            tensor_names.append("token_embd.weight")
        if self.is_tail:
            tensor_names.extend(["output_norm.weight", "output.weight"])
        return tuple(tensor_names)


class GrpcNativeWorkerTransport:
    """Transport that speaks the versioned gRPC worker protocol.

    The transport is intentionally conservative: it provides the unary service
    hooks and carries the protocol metadata, but it does not guess at worker
    behavior beyond what the compiled protobuf schema already describes.
    """

    def __init__(self, worker_url: str, *, timeout: float = 30.0) -> None:
        self.worker_url = worker_url
        self.timeout = timeout
        self._grpc = None
        self._channel = None
        self._stub = None

    def _ensure_stub(self) -> Any:
        if self._stub is not None:
            return self._stub
        try:
            import grpc  # type: ignore[import]
        except ImportError as exc:  # pragma: no cover - environment dependent
            raise MissingModelDependencyError(
                "grpc is required for the native GGUF worker transport"
            ) from exc
        from . import native_protocol

        grpc_mod = native_protocol.load_grpc()
        self._grpc = grpc
        self._channel = grpc.insecure_channel(self.worker_url)
        self._stub = grpc_mod.ShardRuntimeStub(self._channel)
        return self._stub

    def encode_prompt(self, prompt: str, session_id: str | None = None) -> TensorPayload | TailTokenResult | str:
        raise ModelBackendError(
            "gRPC transport is present, but prompt-to-activation translation is provided "
            "by the backend wrapper so it can keep worker framing and tokenizer state aligned"
        )

    def encode_next_token(self, token_id: int, session_id: str) -> TensorPayload | TailTokenResult | str:
        raise ModelBackendError(
            "gRPC transport is present, but decode translation is provided by the backend wrapper"
        )

    def forward_bytes(
        self,
        body: bytes,
        shape: list[int],
        attention_mask_header: str | None,
        position_ids_header: str | None,
        *,
        start_layer: int | None = None,
        session_id: str | None = None,
        cache_mode: str | None = None,
        past_len: int | None = None,
    ) -> TensorPayload | TailTokenResult | str:
        raise ModelBackendError(
            "gRPC transport is present, but activation streaming is handled by the backend wrapper"
        )

    def decode_tail_token(self, hidden_states: Any) -> TailTokenResult:
        raise ModelBackendError("tail decoding is handled by the backend wrapper")

    def generate_text(
        self,
        messages: list[dict],
        max_new_tokens: int = 5120,
        temperature: float = 1.0,
        top_p: float = 1.0,
    ) -> str:
        raise ModelBackendError("text generation is handled by the backend wrapper")

    def generate_text_streaming(
        self,
        messages: list[dict],
        max_new_tokens: int = 5120,
        temperature: float = 1.0,
        top_p: float = 1.0,
    ):
        raise ModelBackendError("streaming generation is handled by the backend wrapper")

    def count_prompt_tokens(self, messages: list[dict]) -> int:
        return sum(1 for message in messages if isinstance(message, dict))

    def count_text_tokens(self, text: str) -> int:
        return len(text.split()) or (1 if text else 0)

    def eos_token_ids(self) -> list[int]:
        return []

    def release_session(self, session_id: str) -> None:
        stub = self._ensure_stub()
        from . import native_protocol

        pb2 = native_protocol.load()
        stub.Release(pb2.ReleaseRequest(reason="release from adapter"))


def build_gguf_backend(
    *,
    model_id: str,
    shard_start: int,
    shard_end: int,
    quantization: str = "bfloat16",
    transport: NativeWorkerTransport | None = None,
    worker_url: str | None = None,
    total_layers: int | None = None,
    model_revision: str | None = None,
    loaded_tensor_names: tuple[str, ...] = (),
    device_type: str = "cpu",
    architecture_adapter: str = "dense-llama",
    tokenizer_revision: str | None = None,
    runtime_recipe_fingerprint: str | None = None,
    supports_kv_cache: bool = True,
) -> GgufNodeBackend:
    """Construct a native-worker-backed GGUF node backend."""
    if transport is None:
        worker_url = worker_url or os.environ.get("MESHNET_NATIVE_WORKER_URL")
        if not worker_url:
            raise MissingModelDependencyError(
                "set MESHNET_NATIVE_WORKER_URL to the local gRPC worker endpoint "
                "or inject a fake transport in tests"
            )
        transport = GrpcNativeWorkerTransport(worker_url)
    return GgufNodeBackend(
        model_id=model_id,
        shard_start=shard_start,
        shard_end=shard_end,
        quantization=quantization,
        transport=transport,
        total_layers=total_layers,
        model_revision=model_revision,
        loaded_tensor_names=loaded_tensor_names,
        device_type=device_type,
        supports_kv_cache=supports_kv_cache,
        worker_url=worker_url,
        architecture_adapter=architecture_adapter,
        tokenizer_revision=tokenizer_revision,
        runtime_recipe_fingerprint=runtime_recipe_fingerprint,
    )