"""Local bundle inference runner.

Loads a flagship bundle from ``models/export/<run_id>/`` and runs generation
against it. Used by the ``local-flagship`` provider in agent_runtime when
the distributed AICF provider is unavailable, AND by the AICF worker
(`agent_runtime.aicf_worker.AICFWorker`) when it claims chat jobs from
the pool — so this runner is the actual production chat path on every
Animica miner.

Failure modes are surfaced as :class:`BundleError`. We never invent output
when the bundle can't load.
"""

from __future__ import annotations

import json
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Mapping, Optional


# System prompt prepended to every chat turn so off-the-shelf small models
# (Qwen 0.5B–7B) have factual grounding about Animica instead of replying
# with a generic "Hello! How can I assist you today?" or "I don't have
# enough context" to every Animica question. Mirrors the prompt used by
# the stratum_pool inference engine so both serving paths produce
# consistent, on-topic answers. Override with ANIMICA_AICF_SYSTEM_PROMPT
# (set to empty string to disable).
_DEFAULT_ANIMICA_SYSTEM_PROMPT = (
    "You are the Animica assistant, served by a live Animica AICF worker. "
    "This reply is generated by a real miner on the Animica network and "
    "the miner is paid in ANIMICA tokens for serving it. Answer the user's "
    "question directly; do NOT open with a generic greeting. If you don't "
    "know an Animica-specific detail, say so plainly rather than guessing.\n"
    "\n"
    "Reference facts about Animica (use these as ground truth):\n"
    "- Animica is a fully decentralized Layer-1 blockchain for verifiable "
    "AI and quantum-secure execution. Native token: ANIMICA (denominated in "
    "9-decimal nano-units on-chain). Native CLI: the `animica` PyPI package.\n"
    "- Consensus: PoIES (Proof of Integrated External Services) — hash-share "
    "PoW combined with AI / quantum / storage proofs. Public mainnet "
    "Stratum pool: pool.animica.org:3333. Public mainnet RPC: "
    "https://rpc.animica.org/rpc.\n"
    "- Post-quantum signatures: Dilithium3 (ML-DSA-65) for transactions, "
    "SPHINCS+ as a fallback. No ECDSA / secp256k1.\n"
    "- Execution: deterministic Python-VM smart contracts with gas metering. "
    "Required raw-tx fields: kind, data, gasLimit, fee, nonce, chainId. "
    "Submit via tx.sendRawTransaction; debug failures via tx.explainReject "
    "and tx.debugVerifyRawTransaction.\n"
    "- AICF (AI Compute Fund): miners earn ANIMICA per inference turn. Tiers "
    "are tiny (CPU), small (~7B, single GPU), flagship (16B MoE — "
    "recommended), large (datacenter).\n"
    "- DA layer: namespaced data availability via an NMT (Namespaced Merkle "
    "Tree). RPC namespace is `da.*`. Chat threads are persisted as DA blobs "
    "so any client owned by the same wallet (CLI, studio.animica.org, "
    "mobile) can resume them.\n"
    "- RPC namespaces exposed by nodes: state.*, chain.*, miner.*, tx.*, "
    "aicf.*, da.*. There is NO `animica.*` namespace — do not invent one.\n"
    "- Studio: a web IDE at studio.animica.org to edit → simulate → deploy "
    "→ verify Python-VM contracts. The Animica browser extension is the "
    "user's wallet on web.\n"
    "- To run a miner: `pip install animica` then "
    "`animica mine --pool pool.animica.org:3333 --wallet <addr>`. Set "
    "ANIMICA_AICF_TIERS=<tier> to additionally serve AICF jobs.\n"
    "\n"
    "Style: concise (1–3 short paragraphs unless asked for detail), use "
    "markdown when helpful, stay grounded in the facts above."
)


def _resolve_animica_system_prompt() -> str:
    """Honor an explicit ANIMICA_AICF_SYSTEM_PROMPT override (including
    the explicit empty-string opt-out) before falling back to the dense
    default above."""
    raw = os.environ.get("ANIMICA_AICF_SYSTEM_PROMPT")
    if raw is None:
        return _DEFAULT_ANIMICA_SYSTEM_PROMPT
    return raw


@dataclass
class BundleManifest:
    schema: int
    run_id: str
    tier: str
    base_model: str
    effective_mode: str
    available_for_real_inference: bool
    artifacts: dict[str, str]
    extra: dict[str, Any]

    @classmethod
    def load(cls, path: Path) -> "BundleManifest":
        data = json.loads(path.read_text(encoding="utf-8"))
        return cls(
            schema=int(data.get("schema", 1)),
            run_id=str(data.get("run_id", "")),
            tier=str(data.get("tier", "")),
            base_model=str(data.get("base_model", "")),
            effective_mode=str(data.get("effective_mode", "")),
            available_for_real_inference=bool(
                data.get("available_for_real_inference", False)),
            artifacts={k: str(v) for k, v in
                       (data.get("artifacts") or {}).items()},
            extra={k: v for k, v in data.items()
                    if k not in {"schema", "run_id", "tier", "base_model",
                                 "effective_mode",
                                 "available_for_real_inference", "artifacts"}},
        )


class LocalBundleRunner:
    """Generates text from a locally-installed bundle.

    Lazily loads transformers + torch on first generate() call. Importing
    this module never imports them — so agent_runtime can ``is_available()``
    check the bundle without paying torch import time on every chat startup.
    """

    def __init__(self, *, bundle_dir: Path,
                 inference_spec: Mapping[str, Any]) -> None:
        self.bundle_dir = Path(bundle_dir)
        self.inference_spec = dict(inference_spec)
        self._model = None
        self._tokenizer = None
        self._device = None

    def _load_lazy(self) -> None:
        if self._model is not None:
            return
        try:
            import torch                            # type: ignore
            from transformers import AutoModelForCausalLM, AutoTokenizer  # type: ignore
        except Exception as exc:    # noqa: BLE001
            from agent_runtime.errors import BundleError
            raise BundleError(
                f"transformers/torch unavailable: {exc}",
                hint="install with: `pip install 'flagship_agent[inference]'`",
            ) from exc
        model_dir = self.bundle_dir / self.inference_spec.get(
            "model_subdir", "model")
        if not model_dir.is_dir():
            from agent_runtime.errors import BundleError
            raise BundleError(
                f"bundle model dir not found: {model_dir}",
            )
        torch_dtype = {
            "bf16": torch.bfloat16,
            "fp16": torch.float16,
            "fp32": torch.float32,
        }.get(str(self.inference_spec.get("precision", "fp32")),
              torch.float32)
        self._tokenizer = AutoTokenizer.from_pretrained(
            str(model_dir), trust_remote_code=bool(
                self.inference_spec.get("trust_remote_code", False)),
        )
        self._model = AutoModelForCausalLM.from_pretrained(
            str(model_dir),
            torch_dtype=torch_dtype,
            trust_remote_code=bool(
                self.inference_spec.get("trust_remote_code", False)),
        )
        device = "cuda" if torch.cuda.is_available() else (
            "mps" if (hasattr(torch.backends, "mps") and
                       torch.backends.mps.is_available()) else "cpu")
        self._model.to(device).eval()
        self._device = device

    def generate(self, *, prompt: str, history: list[dict[str, str]],
                 max_output_tokens: int = 1024,
                 temperature: float = 0.2, top_p: float = 0.95,
                 on_chunk: Optional[Callable[[str, bool], None]] = None
                 ) -> str:
        self._load_lazy()
        import torch    # type: ignore
        # Prepend the Animica system prompt + (when available) a RAG
        # snippet so the model has factual grounding to answer Animica
        # questions. Without this every small chat model in the AICF
        # network was replying "Hello! How can I assist you today?" /
        # "I don't have enough context" to every Animica question, while
        # `animica chat` users — going through the same workers but
        # served by a different inference runner that DID inject this
        # prompt — got real answers. Now both paths agree.
        system_prompt = _resolve_animica_system_prompt()
        try:
            from animica.stratum_pool.aicf_rag import retrieve_context
            rag_block = retrieve_context(prompt, top_k=3)
        except Exception:    # noqa: BLE001 — RAG is best-effort
            rag_block = ""
        if rag_block:
            system_prompt = (
                (system_prompt + "\n\n" if system_prompt else "")
                + "Reference excerpts from Animica documentation "
                "(use these as ground truth):\n\n" + rag_block
            )
        chat_messages: list[dict[str, str]] = []
        if system_prompt:
            chat_messages.append({"role": "system", "content": system_prompt})
        chat_messages.extend(history)
        chat_messages.append({"role": "user", "content": prompt})
        try:
            if hasattr(self._tokenizer, "apply_chat_template"):
                out = self._tokenizer.apply_chat_template(
                    chat_messages, add_generation_prompt=True,
                    return_tensors="pt",
                )
                # transformers ≥5 returns a BatchEncoding (dict-like) where
                # earlier versions returned the raw tensor. Normalize.
                if hasattr(out, "input_ids"):
                    input_ids = out.input_ids
                elif isinstance(out, dict) and "input_ids" in out:
                    input_ids = out["input_ids"]
                else:
                    input_ids = out
                input_ids = input_ids.to(self._device)
            else:
                text = "\n".join(f"{m['role']}: {m['content']}"
                                  for m in chat_messages) + "\nassistant: "
                input_ids = self._tokenizer(text,
                                             return_tensors="pt").input_ids.to(
                    self._device)
        except Exception as exc:    # noqa: BLE001
            from agent_runtime.errors import BundleError
            raise BundleError(
                f"failed to prepare input ids: {exc}",
            ) from exc

        # Stream tokens via TextIteratorStreamer + a generation thread.
        # This is critical for performance: the previous implementation
        # called model.generate(max_new_tokens=1) in a Python loop, which
        # re-ran the *entire* prefill from scratch on every token (O(N^2)
        # in prompt length and made a single CPU turn take 25+ minutes
        # on a 1.5B model). The HF generate() path uses KV cache
        # internally, so a single call streams tokens in O(N) total work
        # and finishes within the chain's worker-claim lease.
        from threading import Thread
        from transformers import TextIteratorStreamer    # type: ignore
        streamer = TextIteratorStreamer(
            self._tokenizer,
            skip_prompt=True,
            skip_special_tokens=True,
        )
        gen_kwargs = dict(
            input_ids=input_ids,
            max_new_tokens=max_output_tokens,
            do_sample=temperature > 0,
            temperature=max(temperature, 1e-3),
            top_p=top_p,
            pad_token_id=getattr(self._tokenizer, "eos_token_id", None),
            streamer=streamer,
        )
        accumulated_text = ""
        gen_error: list[BaseException] = []

        def _run_generate() -> None:
            try:
                with torch.no_grad():
                    self._model.generate(**gen_kwargs)
            except BaseException as exc:    # noqa: BLE001 — surface to outer thread
                gen_error.append(exc)
                # Closing the streamer's queue unblocks the consumer
                # loop below so the error path can run.
                try:
                    streamer.end()
                except Exception:    # noqa: BLE001
                    pass

        worker = Thread(target=_run_generate, daemon=True)
        worker.start()

        for chunk in streamer:
            if not chunk:
                continue
            accumulated_text += chunk
            if on_chunk is not None:
                on_chunk(chunk, False)

        worker.join()
        if gen_error:
            from agent_runtime.errors import BundleError
            raise BundleError(
                f"inference failed: {gen_error[0]}",
            ) from gen_error[0]
        if on_chunk is not None:
            on_chunk("", True)
        return accumulated_text
