#!/usr/bin/env python3
"""Transformers generation entrypoint with additive model-owned NOOSPHERE.

This file is intentionally named ``custom_generate/generate.py``.  Modern
Transformers discovers it from the downloaded model directory and binds this
function as ``model.generate`` when ``trust_remote_code=True`` is used.

NOOSPHERE is not a document service and does not answer from an index.  It
operates on the model's own exact token stream:

1. cold pages receive deterministic one-bit semantic signatures;
2. selected pages are checked by VernadskySurface and VernadskyWeave;
3. ExactLedger rereads authoritative uint32 token IDs;
4. rare evidence tokens receive a bounded sparse decode-time bias;
5. the complete raw prompt still reaches the neural model and KV cache.

NOOSPHERE never substitutes its routed representation for the raw context.
Short prompts and explicit raw-certification mode use the normal generation
path without changing either inputs or logits.
"""

from __future__ import annotations

from collections.abc import Mapping, Sequence
import os
from typing import Any

from .noosphere.additive import build_additive_evidence_plan
from .noosphere.model_memory import (
    ModelMemoryConfig,
    ModelMemoryError,
    ModelOwnedNoosphere,
)


SCHEMA = "lomonosov_zenit_altay_custom_generate_noosphere_v1"
VALID_MODES = frozenset({"auto", "raw", "disabled"})


def _as_mapping(value: Any) -> dict[str, Any]:
    if isinstance(value, Mapping):
        return dict(value)
    if value is None:
        return {}
    if hasattr(value, "to_dict"):
        converted = value.to_dict()
        if isinstance(converted, Mapping):
            return dict(converted)
    return {
        key: getattr(value, key)
        for key in dir(value)
        if not key.startswith("_") and not callable(getattr(value, key, None))
    }


def noosphere_settings(model: Any) -> dict[str, Any]:
    """Return the quantization-independent model-owned memory settings."""

    config = getattr(model, "config", None)
    values = _as_mapping(getattr(config, "noosphere", None))
    if not values and isinstance(config, Mapping):
        values = _as_mapping(config.get("noosphere"))
    return values


def resolve_mode(model: Any, explicit_mode: str | None = None) -> str:
    settings = noosphere_settings(model)
    configured = explicit_mode or os.environ.get("ZENIT_NOOSPHERE_MODE")
    if configured is None:
        configured = settings.get("mode", "auto")
    mode = str(configured).strip().lower()
    if mode not in VALID_MODES:
        raise ModelMemoryError(f"unsupported NOOSPHERE mode: {mode!r}")
    if not bool(settings.get("enabled", True)):
        return "disabled"
    return mode


def memory_config(model: Any) -> ModelMemoryConfig:
    settings = noosphere_settings(model)
    allowed = {
        "context_tokens",
        "page_tokens",
        "hot_tokens",
        "query_tokens",
        "routed_pages",
        "widened_pages",
        "pinned_prefix_pages",
        "semantic_bits",
        "hashes_per_feature",
        "minimum_score",
        "boundary_margin",
    }
    return ModelMemoryConfig(
        **{key: settings[key] for key in allowed if key in settings}
    )


def protected_token_ids(model: Any) -> tuple[int, ...]:
    """Collect multimodal structural tokens that may never leave the prompt."""

    names = (
        "image_token_id",
        "video_token_id",
        "vision_start_token_id",
        "vision_end_token_id",
        "image_pad_token_id",
        "video_pad_token_id",
    )
    roots = [getattr(model, "config", None)]
    config = roots[0]
    if config is not None:
        roots.extend(
            [
                getattr(config, "text_config", None),
                getattr(config, "vision_config", None),
            ]
        )
    values: set[int] = set()
    for root in roots:
        if root is None:
            continue
        for name in names:
            value = (
                root.get(name)
                if isinstance(root, Mapping)
                else getattr(root, name, None)
            )
            if isinstance(value, int) and value >= 0:
                values.add(value)
    return tuple(sorted(values))


class _SparseEvidenceLogitsProcessor:
    """Apply the bounded NOOSPHERE evidence bias without changing input IDs."""

    def __init__(
        self,
        token_ids: Sequence[int],
        token_biases: Sequence[float],
    ) -> None:
        if len(token_ids) != len(token_biases):
            raise ModelMemoryError("NOOSPHERE evidence bias shape mismatch")
        self._token_ids = tuple(int(value) for value in token_ids)
        self._token_biases = tuple(float(value) for value in token_biases)
        self._device_cache: dict[str, tuple[Any, Any]] = {}

    def __call__(self, input_ids: Any, scores: Any) -> Any:
        del input_ids
        if not self._token_ids:
            return scores
        import torch

        key = str(scores.device)
        cached = self._device_cache.get(key)
        if cached is None:
            indices = torch.tensor(
                self._token_ids,
                dtype=torch.long,
                device=scores.device,
            )
            biases = torch.tensor(
                self._token_biases,
                dtype=scores.dtype,
                device=scores.device,
            )
            cached = (indices, biases)
            self._device_cache[key] = cached
        indices, biases = cached
        scores[:, indices] = scores[:, indices] + biases
        return scores


def route_token_ids(
    model: Any,
    token_ids: Sequence[int],
    *,
    query_token_ids: Sequence[int] | None = None,
) -> tuple[tuple[int, ...], dict[str, Any]]:
    """Compile the exact routing view used to build additive evidence.

    This compatibility function does not define the release execution path:
    generation keeps the original raw token stream model-visible.
    """

    config = memory_config(model)
    horizon = ModelOwnedNoosphere(
        token_ids,
        config=config,
        protected_token_ids=protected_token_ids(model),
    )
    compiled = horizon.compile(query_token_ids)
    receipt = compiled.receipt.to_dict()
    receipt.update(
        {
            "generation_schema": SCHEMA,
            "architecture_owned": True,
            "quantization_independent": True,
            "logical_context_tokens": len(token_ids),
            "model_visible_tokens": len(compiled.token_ids),
        }
    )
    return compiled.token_ids, receipt


def _active_row(input_ids: Any, attention_mask: Any | None) -> Any:
    if getattr(input_ids, "ndim", None) != 2 or int(input_ids.shape[0]) != 1:
        raise ModelMemoryError(
            "model-owned NOOSPHERE currently requires batch size 1 for long prompts"
        )
    row = input_ids[0]
    if attention_mask is None:
        return row
    if getattr(attention_mask, "shape", None) != getattr(input_ids, "shape", None):
        raise ModelMemoryError("attention_mask shape does not match input_ids")
    return row[attention_mask[0].to(dtype=__import__("torch").bool)]


def _call_standard_generate(
    model: Any,
    *,
    inputs: Any,
    generation_config: Any,
    logits_processor: Any,
    stopping_criteria: Any,
    prefix_allowed_tokens_fn: Any,
    synced_gpus: Any,
    assistant_model: Any,
    streamer: Any,
    negative_prompt_ids: Any,
    negative_prompt_attention_mask: Any,
    kwargs: dict[str, Any],
) -> Any:
    from transformers.generation.utils import GenerationMixin

    return GenerationMixin.generate(
        model,
        inputs=inputs,
        generation_config=generation_config,
        logits_processor=logits_processor,
        stopping_criteria=stopping_criteria,
        prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
        synced_gpus=synced_gpus,
        assistant_model=assistant_model,
        streamer=streamer,
        negative_prompt_ids=negative_prompt_ids,
        negative_prompt_attention_mask=negative_prompt_attention_mask,
        **kwargs,
    )


def _restore_logical_prefix(
    output: Any,
    *,
    original_input_ids: Any,
    model_visible_tokens: int,
) -> Any:
    """Preserve the standard caller contract for completion slicing."""

    import torch

    sequences = output.sequences if hasattr(output, "sequences") else output
    if not isinstance(sequences, torch.Tensor) or sequences.ndim != 2:
        raise ModelMemoryError("standard generation returned no sequence tensor")
    if sequences.shape[-1] < model_visible_tokens:
        raise ModelMemoryError("generated sequence is shorter than model-visible prompt")
    suffix = sequences[:, model_visible_tokens:]
    prefix = original_input_ids
    if prefix.shape[0] != sequences.shape[0]:
        if prefix.shape[0] != 1 or sequences.shape[0] % prefix.shape[0] != 0:
            raise ModelMemoryError("cannot restore logical prefix for generated beams")
        prefix = prefix.repeat(sequences.shape[0], 1)
    restored = torch.cat((prefix.to(sequences.device), suffix), dim=-1)
    if hasattr(output, "sequences"):
        output.sequences = restored
        return output
    return restored


def generate(
    model: Any,
    inputs: Any | None = None,
    generation_config: Any | None = None,
    logits_processor: Any | None = None,
    stopping_criteria: Any | None = None,
    prefix_allowed_tokens_fn: Any | None = None,
    synced_gpus: bool | None = None,
    assistant_model: Any | None = None,
    streamer: Any | None = None,
    negative_prompt_ids: Any | None = None,
    negative_prompt_attention_mask: Any | None = None,
    noosphere_mode: str | None = None,
    noosphere_query_token_ids: Sequence[int] | None = None,
    **kwargs: Any,
) -> Any:
    """Generate with automatic model-owned long-context retrieval."""

    import torch

    mode = resolve_mode(model, noosphere_mode)
    kwargs = dict(kwargs)
    kw_input_ids = kwargs.get("input_ids")
    if inputs is not None and kw_input_ids is not None:
        raise ModelMemoryError("input IDs were supplied twice")
    original_input_ids = inputs if inputs is not None else kw_input_ids

    should_route = (
        mode == "auto"
        and isinstance(original_input_ids, torch.Tensor)
        and original_input_ids.ndim == 2
    )
    config = memory_config(model)
    if should_route:
        active = _active_row(original_input_ids, kwargs.get("attention_mask"))
        should_route = int(active.numel()) > config.hot_tokens

    if not should_route:
        setattr(
            model,
            "_noosphere_last_receipt",
            {
                "generation_schema": SCHEMA,
                "status": "RAW_OR_SHORT_CONTEXT",
                "mode": mode,
                "architecture_owned": True,
            },
        )
        return _call_standard_generate(
            model,
            inputs=inputs,
            generation_config=generation_config,
            logits_processor=logits_processor,
            stopping_criteria=stopping_criteria,
            prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
            synced_gpus=synced_gpus,
            assistant_model=assistant_model,
            streamer=streamer,
            negative_prompt_ids=negative_prompt_ids,
            negative_prompt_attention_mask=negative_prompt_attention_mask,
            kwargs=kwargs,
        )

    active = _active_row(original_input_ids, kwargs.get("attention_mask"))
    active_ids = tuple(int(value) for value in active.detach().cpu().tolist())
    query_ids = (
        tuple(int(value) for value in noosphere_query_token_ids)
        if noosphere_query_token_ids is not None
        else None
    )
    plan = build_additive_evidence_plan(
        active_ids,
        config=config,
        protected_token_ids=protected_token_ids(model),
        query_token_ids=query_ids,
    )
    receipt = dict(plan.receipt)
    receipt.update(
        {
            "generation_schema": SCHEMA,
            "architecture_owned": True,
            "quantization_independent": True,
            "logical_context_tokens": len(active_ids),
            "model_visible_tokens": len(active_ids),
        }
    )
    setattr(model, "_noosphere_last_receipt", receipt)

    from transformers import LogitsProcessorList

    processors = LogitsProcessorList(
        list(logits_processor) if logits_processor is not None else []
    )
    processors.append(
        _SparseEvidenceLogitsProcessor(plan.token_ids, plan.token_biases)
    )
    return _call_standard_generate(
        model,
        inputs=inputs,
        generation_config=generation_config,
        logits_processor=processors,
        stopping_criteria=stopping_criteria,
        prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
        synced_gpus=synced_gpus,
        assistant_model=assistant_model,
        streamer=streamer,
        negative_prompt_ids=negative_prompt_ids,
        negative_prompt_attention_mask=negative_prompt_attention_mask,
        kwargs=kwargs,
    )


__all__ = [
    "SCHEMA",
    "generate",
    "memory_config",
    "noosphere_settings",
    "protected_token_ids",
    "resolve_mode",
    "route_token_ids",
]