"""Dense-Llama GGUF ownership helpers.

This module keeps two related concerns together:

* selecting the tensors a dense-Llama GGUF shard is allowed to own; and
* inferring the authoritative loaded range / endpoint ownership from the
  tensors the model actually exposes.

The first is used by the range-aware loader seam. The second is used by the
doctor/admission/reporting path so the tracker sees what the model loaded, not
what a CLI flag claimed.
"""

from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Any, Iterable, Mapping

_BLOCK_RE = re.compile(r"^blk\.(\d+)\.")

_HEAD_TENSOR_NAMES = {
    "token_embd.weight",
    "token_embd.bias",
    "tok_embeddings.weight",
    "tok_embeddings.bias",
    "embed_tokens.weight",
    "embed_tokens.bias",
}

_TAIL_TENSOR_NAMES = {
    "output_norm.weight",
    "output_norm.bias",
    "output.weight",
    "output.bias",
    "lm_head.weight",
    "lm_head.bias",
}


@dataclass(frozen=True)
class DenseLlamaShardOwnership:
    """Authoritative ownership for one dense-Llama shard."""

    start_layer: int
    end_layer: int
    owns_embedding: bool
    owns_final_head: bool
    tensor_names: tuple[str, ...] = ()
    source_artifact_hash: str | None = None
    slice_artifact_hash: str | None = None
    derivative_slice: bool = False
    final_artifact_semantics: bool = True

    def __post_init__(self) -> None:
        if self.start_layer < 0:
            raise ValueError("start_layer must be non-negative")
        if self.end_layer < self.start_layer:
            raise ValueError("end_layer must be >= start_layer")
        if self.derivative_slice:
            if not self.source_artifact_hash or not self.slice_artifact_hash:
                raise ValueError(
                    "temporary derivative sub-GGUFs must carry source and slice hashes"
                )
            if self.final_artifact_semantics:
                raise ValueError(
                    "temporary derivative sub-GGUFs must not be claimed as final artifacts"
                )

    @property
    def range(self) -> tuple[int, int]:
        return self.start_layer, self.end_layer

    def to_dict(self) -> dict[str, Any]:
        return {
            "start_layer": self.start_layer,
            "end_layer": self.end_layer,
            "owns_embedding": self.owns_embedding,
            "owns_final_head": self.owns_final_head,
            "tensor_names": list(self.tensor_names),
            "source_artifact_hash": self.source_artifact_hash,
            "slice_artifact_hash": self.slice_artifact_hash,
            "derivative_slice": self.derivative_slice,
            "final_artifact_semantics": self.final_artifact_semantics,
        }


def select_dense_llama_tensor_names(
    tensor_names: Iterable[str],
    start_layer: int,
    end_layer: int,
    *,
    total_layers: int | None = None,
) -> set[str]:
    """Return the dense-Llama GGUF tensor names owned by an inclusive range."""
    if start_layer < 0:
        raise ValueError("start_layer must be non-negative")
    if end_layer < start_layer:
        raise ValueError("end_layer must be greater than or equal to start_layer")

    selected: set[str] = set()
    for tensor_name in tensor_names:
        if _tensor_belongs_to_range(tensor_name, start_layer, end_layer, total_layers):
            selected.add(tensor_name)
    return selected


def infer_dense_llama_ownership(
    tensor_names: Iterable[str],
    *,
    total_layers: int | None = None,
    source_artifact_hash: str | None = None,
    slice_artifact_hash: str | None = None,
    derivative_slice: bool = False,
    final_artifact_semantics: bool = True,
) -> DenseLlamaShardOwnership:
    """Infer authoritative loaded range and endpoint ownership from tensors."""
    names = tuple(str(name) for name in tensor_names if isinstance(name, str))
    if not names:
        raise ValueError("tensor inventory is empty")

    block_layers = sorted(
        {
            layer
            for name in names
            if (layer := _layer_index(name)) is not None
        }
    )
    if not block_layers:
        raise ValueError("tensor inventory does not contain any blk.N.* tensors")

    selected = tuple(sorted(names))
    return DenseLlamaShardOwnership(
        start_layer=block_layers[0],
        end_layer=block_layers[-1],
        owns_embedding=any(_is_head_tensor(name) for name in names),
        owns_final_head=any(
            _is_tail_tensor(name, total_layers=total_layers, loaded_end=block_layers[-1])
            for name in names
        ),
        tensor_names=selected,
        source_artifact_hash=source_artifact_hash,
        slice_artifact_hash=slice_artifact_hash,
        derivative_slice=derivative_slice,
        final_artifact_semantics=final_artifact_semantics,
    )


def authoritative_dense_llama_ownership(
    backend: Any,
    selection: Any | None = None,
) -> DenseLlamaShardOwnership:
    """Return the most authoritative dense-Llama ownership the backend exposes."""
    tensor_names = _tensor_names_from_backend(backend)
    if tensor_names:
        try:
            return infer_dense_llama_ownership(
                tensor_names,
                total_layers=_backend_total_layers(backend, selection),
            )
        except ValueError:
            pass

    start, end = _backend_loaded_bounds(backend, selection)
    return DenseLlamaShardOwnership(
        start_layer=start,
        end_layer=end,
        owns_embedding=_backend_owns_embedding(backend, start),
        owns_final_head=_backend_owns_final_head(backend, end),
    )


def _backend_loaded_bounds(backend: Any, selection: Any | None) -> tuple[int, int]:
    start = getattr(backend, "loaded_shard_start", None)
    end = getattr(backend, "loaded_shard_end", None)
    if start is None:
        start = getattr(backend, "shard_start", None)
    if end is None:
        end = getattr(backend, "shard_end", None)
    if start is None or end is None:
        if selection is None:
            raise ValueError("backend does not expose a loaded shard range")
        start = getattr(selection, "shard_start")
        end = getattr(selection, "shard_end")
    return int(start), int(end)


def _backend_owns_embedding(backend: Any, start: int) -> bool:
    value = getattr(backend, "owns_embedding", None)
    if value is None:
        value = getattr(backend, "is_head", start == 0)
    return bool(value)


def _backend_owns_final_head(backend: Any, end: int) -> bool:
    value = getattr(backend, "owns_final_head", None)
    if value is None:
        value = getattr(backend, "is_tail", False)
    return bool(value)


def _backend_total_layers(backend: Any, selection: Any | None) -> int | None:
    value = getattr(backend, "total_layers", None)
    if isinstance(value, int) and value > 0:
        return value
    if selection is None:
        return None
    total = getattr(selection, "total_layers", None)
    if isinstance(total, int) and total > 0:
        return total
    return None


def _tensor_names_from_backend(backend: Any) -> tuple[str, ...]:
    for attr in ("loaded_tensor_names", "tensor_names", "tensor_inventory"):
        value = getattr(backend, attr, None)
        names = _normalise_tensor_names(value)
        if names:
            return names
    return ()


def _normalise_tensor_names(value: Any) -> tuple[str, ...]:
    if value is None:
        return ()
    if isinstance(value, Mapping):
        items = value.keys()
    else:
        try:
            items = list(value)
        except TypeError:
            return ()
    names = [str(item) for item in items if isinstance(item, str) and item.strip()]
    return tuple(names)


def _tensor_belongs_to_range(
    tensor_name: str,
    start_layer: int,
    end_layer: int,
    total_layers: int | None,
) -> bool:
    layer = _layer_index(tensor_name)
    if layer is not None:
        return start_layer <= layer <= end_layer

    if start_layer == 0 and _is_head_tensor(tensor_name):
        return True

    if total_layers is not None and end_layer >= total_layers - 1 and _is_tail_tensor(
        tensor_name, total_layers=total_layers, loaded_end=end_layer
    ):
        return True

    return False


def _layer_index(tensor_name: str) -> int | None:
    match = _BLOCK_RE.match(tensor_name)
    if match is None:
        return None
    return int(match.group(1))


def _is_head_tensor(tensor_name: str) -> bool:
    lowered = tensor_name.lower()
    return lowered in _HEAD_TENSOR_NAMES or any(
        lowered.startswith(prefix)
        for prefix in ("token_embd.", "tok_embeddings.", "embed_tokens.")
    )


def _is_tail_tensor(
    tensor_name: str,
    *,
    total_layers: int | None,
    loaded_end: int,
) -> bool:
    lowered = tensor_name.lower()
    if lowered in _TAIL_TENSOR_NAMES:
        return True
    if total_layers is not None and loaded_end >= total_layers - 1:
        return any(
            lowered.startswith(prefix)
            for prefix in ("output_norm.", "final_norm.", "norm.")
        )
    return False