"""Hugging Face backend with microbatching, KV cache, and bounded activation cards."""

from __future__ import annotations

import asyncio
import hashlib
import threading
import time
from contextlib import ExitStack, contextmanager, nullcontext
from pathlib import Path, PurePosixPath, PureWindowsPath
from typing import Any, Iterator, Literal, Mapping

from ..contracts import BackendCapabilities
from ..errors import UnsupportedCapability
from .base import Backend, BackendInput, BackendResult


def adapter_directory_sha256(path: str | Path) -> str:
    """Commit the PEFT config and weight file without publishing their path."""

    root = Path(path).resolve()
    config = root / "adapter_config.json"
    weights = next(
        (
            candidate
            for candidate in (
                root / "adapter_model.safetensors",
                root / "adapter_model.bin",
            )
            if candidate.is_file()
        ),
        None,
    )
    if not config.is_file() or weights is None:
        raise ValueError(
            "adapter directory must contain adapter_config.json and one weight file"
        )
    digest = hashlib.sha256()
    for source in (config, weights):
        with source.open("rb") as handle:
            for chunk in iter(lambda: handle.read(1024 * 1024), b""):
                digest.update(chunk)
    return digest.hexdigest()


def _peft_adapter_name(card_id: str) -> str:
    """Map a public card id onto a PyTorch-safe internal module name."""

    return "bp_" + hashlib.sha256(card_id.encode("utf-8")).hexdigest()


def _receipt_model_label(model_id: str, model_label: str | None) -> str:
    """Require a path-free receipt label whenever the model source is local."""

    source = str(model_id)
    local_source = (
        Path(source).expanduser().is_dir()
        or PureWindowsPath(source).is_absolute()
        or PurePosixPath(source).is_absolute()
    )
    if local_source and not model_label:
        raise ValueError("local model paths require an explicit path-free model_label")
    label = str(model_label or source).strip()
    if not label:
        raise ValueError("model_label cannot be blank")
    if PureWindowsPath(label).is_absolute() or PurePosixPath(label).is_absolute():
        raise ValueError("model_label must not be an absolute filesystem path")
    return label


class TransformersBackend(Backend):
    """Execute trusted built-in cards against a local Hugging Face model.

    The backend owns model mutation under a lock and removes every hook before
    releasing it. The scheduler only batches requests with the same card stack.
    """

    def __init__(
        self,
        model_id: str = "Qwen/Qwen3-1.7B",
        *,
        device: str = "cuda",
        dtype: str = "float16",
        local_files_only: bool = False,
        quantization: Literal["none", "4bit"] = "none",
        adapter_registry: Mapping[str, str | Path] | None = None,
        model_label: str | None = None,
        system_prompt: str | None = None,
    ):
        try:
            import torch
            from transformers import (
                AutoModelForCausalLM,
                AutoTokenizer,
                BitsAndBytesConfig,
            )
        except ImportError as exc:
            raise RuntimeError(
                "install backplane-runtime[transformers] for this backend"
            ) from exc
        if quantization not in {"none", "4bit"}:
            raise ValueError("quantization must be 'none' or '4bit'")
        if quantization == "4bit" and not device.startswith("cuda"):
            raise ValueError("4-bit Backplane execution currently requires CUDA")

        self.torch = torch
        self.model_id = model_id
        self.model_label = _receipt_model_label(model_id, model_label)
        self.device = device
        self.quantization = quantization
        self.system_prompt = system_prompt.strip() if system_prompt else None
        self.system_prompt_sha256 = (
            hashlib.sha256(self.system_prompt.encode("utf-8")).hexdigest()
            if self.system_prompt
            else None
        )
        self._adapter_paths = {
            str(card_id): Path(path).resolve()
            for card_id, path in (adapter_registry or {}).items()
        }
        for card_id, path in self._adapter_paths.items():
            if not card_id.strip():
                raise ValueError("adapter registry card ids cannot be blank")
            if not path.is_dir():
                raise ValueError(f"adapter registry path not found for {card_id}")
        self._adapter_digests = {
            card_id: adapter_directory_sha256(path)
            for card_id, path in self._adapter_paths.items()
        }
        self._peft_adapter_names = {
            card_id: _peft_adapter_name(card_id)
            for card_id in self._adapter_paths
        }
        dtype_value = getattr(torch, dtype)
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_id,
            local_files_only=local_files_only,
            trust_remote_code=False,
        )
        if self.tokenizer.pad_token_id is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        self.tokenizer.padding_side = "left"
        if quantization == "4bit":
            quantization_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=dtype_value,
                bnb_4bit_use_double_quant=True,
            )
            self.model = AutoModelForCausalLM.from_pretrained(
                model_id,
                local_files_only=local_files_only,
                trust_remote_code=False,
                quantization_config=quantization_config,
                device_map={"": 0},
            )
        else:
            try:
                self.model = AutoModelForCausalLM.from_pretrained(
                    model_id,
                    dtype=dtype_value,
                    local_files_only=local_files_only,
                    trust_remote_code=False,
                ).to(device)
            except TypeError:
                self.model = AutoModelForCausalLM.from_pretrained(
                    model_id,
                    torch_dtype=dtype_value,
                    local_files_only=local_files_only,
                    trust_remote_code=False,
                ).to(device)
        if self._adapter_paths:
            try:
                from peft import PeftModel
            except ImportError as exc:
                raise RuntimeError(
                    "install backplane-runtime[transformers] for adapter execution"
                ) from exc
            registrations = iter(self._adapter_paths.items())
            first_card, first_path = next(registrations)
            self.model = PeftModel.from_pretrained(
                self.model,
                str(first_path),
                adapter_name=self._peft_adapter_names[first_card],
                is_trainable=False,
            )
            for card_id, path in registrations:
                self.model.load_adapter(
                    str(path),
                    adapter_name=self._peft_adapter_names[card_id],
                    is_trainable=False,
                )
        self.model.eval()
        for parameter in self.model.parameters():
            parameter.requires_grad_(False)
        self.layers = self._find_layers()
        self.output_embeddings = self.model.get_output_embeddings().weight
        self._lock = threading.RLock()
        self._capabilities = BackendCapabilities(
            backend="transformers",
            model=self.model_label,
            batching=True,
            kv_cache=True,
            activation_cards=True,
            per_request_adapters=bool(self._adapter_paths),
            metrics=(
                "latency_ms",
                "batch_size",
                "input_tokens",
                "output_tokens",
                "gpu_peak_memory_bytes",
            ),
            limitations=(
                "Trusted built-in activation cards only.",
                "Adapter paths are operator-owned and never accepted from card packages.",
                "Optimized serving engines require their own card execution plugin.",
            ),
        )

    @property
    def capabilities(self) -> BackendCapabilities:
        return self._capabilities

    def _find_layers(self):
        for path in (
            "model.layers",
            "model.model.layers",
            "model.language_model.layers",
            "transformer.h",
            "gpt_neox.layers",
        ):
            current: Any = self.model
            for part in path.split("."):
                if not hasattr(current, part):
                    break
                current = getattr(current, part)
            else:
                return current
        raise UnsupportedCapability("could not locate decoder layers")

    async def generate_batch(self, requests: list[BackendInput]) -> list[BackendResult]:
        return await asyncio.to_thread(self._generate_sync, requests)

    def _generate_sync(self, requests: list[BackendInput]) -> list[BackendResult]:
        if not requests:
            return []
        if len({request.stack_key for request in requests}) != 1:
            raise RuntimeError("mixed card stacks reached TransformersBackend")
        activation_specs = requests[0].activation_specs
        if any(request.activation_specs != activation_specs for request in requests):
            raise RuntimeError("mixed activation specs reached TransformersBackend")

        torch = self.torch
        with self._lock:
            encoded = self.tokenizer(
                [self._format_prompt(request.prompt) for request in requests],
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=8_192,
            ).to(self.device)
            input_width = int(encoded["input_ids"].shape[1])
            input_counts = encoded["attention_mask"].sum(dim=1).tolist()
            max_new_tokens = max(request.max_new_tokens for request in requests)
            if self.device.startswith("cuda"):
                torch.cuda.reset_peak_memory_stats()
            started = time.perf_counter()
            with self._adapter_context(activation_specs) as adapter_card, ExitStack() as hooks:
                for spec in (
                    item for item in activation_specs if item.get("kind") == "steering"
                ):
                    handle = self._install_activation(spec)
                    hooks.callback(handle.remove)
                with torch.inference_mode():
                    generated = self.model.generate(
                        **encoded,
                        max_new_tokens=max_new_tokens,
                        do_sample=False,
                        use_cache=True,
                        pad_token_id=self.tokenizer.pad_token_id,
                        eos_token_id=self.tokenizer.eos_token_id,
                    )
            elapsed_ms = (time.perf_counter() - started) * 1000
            peak = (
                int(torch.cuda.max_memory_allocated())
                if self.device.startswith("cuda")
                else None
            )
            results: list[BackendResult] = []
            for index, request in enumerate(requests):
                continuation = generated[index, input_width : input_width + request.max_new_tokens]
                text = self.tokenizer.decode(continuation, skip_special_tokens=True)
                results.append(
                    BackendResult(
                        text=text,
                        input_tokens=int(input_counts[index]),
                        output_tokens=int(continuation.numel()),
                        latency_ms=round(elapsed_ms, 3),
                        gpu_memory_bytes=peak,
                        metadata={
                            "batch_size": len(requests),
                            "kv_cache": True,
                            "activation_cards": len(activation_specs),
                            "adapter_card": adapter_card,
                        },
                    )
                )
            return results

    def _format_prompt(self, prompt: str) -> str:
        if self.system_prompt is None:
            return prompt
        if not getattr(self.tokenizer, "chat_template", None):
            raise UnsupportedCapability(
                "configured system prompt requires a tokenizer chat template"
            )
        messages = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": prompt},
        ]
        try:
            return self.tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
                enable_thinking=False,
            )
        except TypeError:
            return self.tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
            )

    @contextmanager
    def _adapter_context(
        self,
        activation_specs: tuple[dict[str, Any], ...],
    ) -> Iterator[str | None]:
        adapters = [
            spec for spec in activation_specs if spec.get("kind") == "adapter"
        ]
        if len(adapters) > 1:
            raise UnsupportedCapability(
                "TransformersBackend allows at most one adapter card per request"
            )
        if not adapters:
            context = (
                self.model.disable_adapter()
                if self._adapter_paths
                else nullcontext()
            )
            with context:
                yield None
            return
        spec = adapters[0]
        card_id = str(spec.get("card_id", ""))
        if card_id not in self._adapter_paths:
            raise UnsupportedCapability(
                f"adapter card is not registered with this backend: {card_id}"
            )
        expected = str(spec.get("artifact_sha256", ""))
        if expected != self._adapter_digests[card_id]:
            raise UnsupportedCapability(
                f"adapter artifact hash mismatch for card: {card_id}"
            )
        expected_system_prompt = spec.get("system_prompt_sha256")
        if (
            expected_system_prompt is not None
            and expected_system_prompt != self.system_prompt_sha256
        ):
            raise UnsupportedCapability(
                f"adapter system prompt hash mismatch for card: {card_id}"
            )
        self.model.set_adapter(self._peft_adapter_names[card_id])
        yield card_id

    def _install_activation(self, spec: dict[str, Any]):
        if spec.get("kind") != "steering":
            raise UnsupportedCapability(
                f"TransformersBackend cannot execute activation card {spec.get('kind')}"
            )
        token = str(spec.get("token", "")).strip()
        if not token:
            raise UnsupportedCapability("steering card requires a token")
        token_ids = self.tokenizer(
            token, add_special_tokens=False, return_tensors="pt"
        )["input_ids"][0].to(self.device)
        rows = self.output_embeddings[token_ids].float()
        direction = rows.mean(0)
        direction = direction / direction.norm().clamp_min(1e-6)
        fraction = float(spec.get("depth", 0.78))
        slot = max(0, min(len(self.layers) - 1, round(fraction * (len(self.layers) - 1))))
        magnitude = min(abs(float(spec.get("magnitude", 0.08))), 0.35)
        sign = -1.0 if spec.get("mode") == "subtract" else 1.0

        def hook(_, __, output):
            is_tuple = isinstance(output, tuple)
            hidden = output[0] if is_tuple else output
            delta = direction.to(hidden.dtype) * hidden[:, -1, :].float().norm(
                dim=-1, keepdim=True
            ).to(hidden.dtype)
            edited = hidden.clone()
            edited[:, -1, :] = edited[:, -1, :] + sign * magnitude * delta
            return (edited,) + tuple(output[1:]) if is_tuple else edited

        return self.layers[slot].register_forward_hook(hook)

