"""PEFT-backed inference engine for non-LoRA adapters.

This module provides an inference engine using HuggingFace Transformers + PEFT
for adapter types that vLLM's --enable-lora doesn't support (e.g., PROMPT_TUNING).

The engine lazy-loads the base model on first use to minimize memory pressure
when multiple engines (vLLM + PEFT) coexist.
"""

from __future__ import annotations

from collections import OrderedDict
from collections.abc import Iterator
from threading import Lock, Thread
from typing import TYPE_CHECKING, Any

from appinfra.log import Logger

if TYPE_CHECKING:
    from ..context import RequestContext
    from ..serving.adapters import LoadedAdapter
    from ..serving.dispatch.config import PEFTEngineConfig


class _GenerationErrorHolder:
    """Thread-safe holder for exceptions from background generation thread."""

    def __init__(self) -> None:
        self.error: BaseException | None = None


class PEFTStreamingIterator:
    """Streaming iterator for PEFT generation using TextIteratorStreamer.

    Yields text chunks as they are generated by the model.
    """

    def __init__(
        self,
        streamer: Any,  # TextIteratorStreamer
        prompt_tokens: int,
        tokenizer: Any,
        error_holder: _GenerationErrorHolder | None = None,
    ) -> None:
        self._streamer = streamer
        self._tokenizer = tokenizer
        self._iter: Iterator[str] | None = None
        self._finished = False
        self._accumulated_text = ""
        self._error_holder = error_holder

        # Stats
        self.prompt_tokens: int = prompt_tokens
        self.completion_tokens: int = 0
        self.finish_reason: str = "stop"

    def __iter__(self) -> Iterator[str]:
        return self

    def __next__(self) -> str:
        if self._finished:
            raise StopIteration

        if self._iter is None:
            self._iter = iter(self._streamer)

        try:
            chunk = next(self._iter)
            self._accumulated_text += chunk
            return chunk
        except StopIteration:
            self._finished = True
            # Count actual tokens from accumulated text
            if self._accumulated_text:
                tokens = self._tokenizer.encode(
                    self._accumulated_text, add_special_tokens=False
                )
                self.completion_tokens = len(tokens)
            # Check if generation thread had an error
            if self._error_holder and self._error_holder.error:
                raise RuntimeError(
                    f"Generation failed: {self._error_holder.error}"
                ) from self._error_holder.error
            raise


class _StopOnSequences:
    """Stopping criteria for stop sequences in HuggingFace generate().

    Checks if any of the stop sequences appear at the end of generated tokens.
    """

    def __init__(self, stop_ids: list[list[int]]) -> None:
        self.stop_ids = stop_ids

    def __call__(self, input_ids: Any, scores: Any, **kwargs: Any) -> bool:
        for stop_id_seq in self.stop_ids:
            if len(stop_id_seq) <= input_ids.shape[1]:
                if input_ids[0, -len(stop_id_seq) :].tolist() == stop_id_seq:
                    return True
        return False


class AdapterRegistry:
    """Thread-safe registry tracking loaded adapter names.

    PEFT's multi-adapter support requires loading adapters onto a single PeftModel
    and switching between them with set_adapter(). This registry tracks which
    adapters have been loaded.
    """

    def __init__(self, maxsize: int = 4) -> None:
        # Maps adapter_path -> adapter_name (for PEFT's set_adapter)
        self._loaded: OrderedDict[str, str] = OrderedDict()
        self._maxsize = maxsize
        self._lock = Lock()

    def get_adapter_name(self, adapter_path: str) -> str | None:
        """Get adapter name if already loaded, moving to end (most recently used)."""
        with self._lock:
            if adapter_path in self._loaded:
                self._loaded.move_to_end(adapter_path)
                return self._loaded[adapter_path]
            return None

    def register(self, adapter_path: str, adapter_name: str) -> str | None:
        """Register adapter as loaded.

        Returns adapter_name to delete if at capacity, else None.
        """
        with self._lock:
            if adapter_path in self._loaded:
                self._loaded.move_to_end(adapter_path)
                return None

            to_delete = None
            if len(self._loaded) >= self._maxsize:
                _, to_delete = self._loaded.popitem(last=False)
            self._loaded[adapter_path] = adapter_name
            return to_delete

    def clear(self) -> list[str]:
        """Clear registry and return all adapter names."""
        with self._lock:
            names = list(self._loaded.values())
            self._loaded.clear()
            return names

    def __len__(self) -> int:
        with self._lock:
            return len(self._loaded)


class PEFTEngine:
    """PEFT-backed inference engine for PROMPT_TUNING and similar adapters.

    Uses HuggingFace Transformers + PEFT library directly. Lazy-loads the
    base model on first adapter request to minimize memory when coexisting
    with vLLM.

    This is a secondary engine - it doesn't implement all InferenceEngineProtocol
    methods since it's only used for specific adapter types, not as a primary engine.
    """

    def __init__(
        self,
        lg: Logger,
        config: PEFTEngineConfig,
        model_path: str,
        adapters: list[LoadedAdapter] | None = None,
    ) -> None:
        """Initialize PEFT engine.

        Args:
            lg: Logger instance.
            config: PEFT engine configuration.
            model_path: Path to the base model weights.
            adapters: Pre-scanned adapters from AdapterManager (for metadata).
        """
        self._lg = lg
        self._config = config
        self._model_path = model_path

        # Lazy-loaded components
        self._base_model: Any = None  # AutoModelForCausalLM (before any adapters)
        self._peft_model: Any = None  # PeftModel (with adapters loaded)
        self._tokenizer: Any = None  # AutoTokenizer
        self._load_lock = Lock()
        self._model_loaded = False

        # Adapter registry - tracks loaded adapter names for set_adapter()
        self._adapter_registry = AdapterRegistry(maxsize=config.max_cached_adapters)
        self._active_adapter: str | None = None  # Currently active adapter name
        self._adapter_counter = 0  # For generating unique adapter names

        # Adapter metadata from AdapterManager (key -> {md5, mtime})
        self._adapter_paths: dict[str, str] = {}
        self._adapter_metadata: dict[str, dict[str, str]] = {}
        if adapters:
            self._init_adapters(adapters)

        # Compute dtype for 4-bit quantization (set during model load)
        self._compute_dtype: Any = None

    def _init_adapters(self, adapters: list[LoadedAdapter]) -> None:
        """Initialize adapter maps from pre-scanned LoadedAdapter list."""
        for adapter in adapters:
            # Resolve symlinks to match how handler resolves paths
            resolved_path = str(adapter.path.resolve())
            self._adapter_paths[adapter.key] = resolved_path
            self._adapter_metadata[adapter.key] = {
                "md5": adapter.md5 or "unknown",
                "mtime": adapter.mtime or "unknown",
            }
        if adapters:
            self._lg.info(
                "PEFT adapters initialized",
                extra={"count": len(adapters)},
            )

    def _get_adapter_key(self, adapter_path: str) -> str | None:
        """Look up adapter key from path."""
        for key, path in self._adapter_paths.items():
            if path == adapter_path:
                return key
        # Debug: log mismatch
        self._lg.warning(
            "adapter key lookup failed",
            extra={
                "lookup_path": adapter_path,
                "known_paths": list(self._adapter_paths.values()),
            },
        )
        return None

    def _build_adapter_info(self, adapter_path: str) -> dict[str, Any]:
        """Build adapter info dict with proper key and metadata."""
        key = self._get_adapter_key(adapter_path)
        if key:
            metadata = self._adapter_metadata.get(key, {})
            return {
                "requested": key,
                "actual": key,
                "fallback": False,
                "mtime": metadata.get("mtime"),
                "md5": metadata.get("md5"),
            }
        # Fallback: no metadata available, use path
        return {
            "requested": adapter_path,
            "actual": adapter_path,
            "fallback": False,
        }

    def _ensure_base_model_loaded(self) -> None:
        """Lazy-load the base model and tokenizer on first use."""
        if self._model_loaded:
            return

        with self._load_lock:
            if self._model_loaded:
                return

            from appinfra.time import since, start

            t0 = start()
            self._lg.info(
                "loading base model for PEFT engine",
                extra={"model_path": self._model_path},
            )
            self._load_tokenizer_and_model()
            self._model_loaded = True
            self._lg.info("PEFT base model loaded", extra={"after": since(t0)})

    def _load_tokenizer_and_model(self) -> None:
        """Load tokenizer and base model with configured options."""
        try:
            from transformers import AutoModelForCausalLM, AutoTokenizer
        except ImportError as e:
            raise ImportError(
                "PEFT engine requires transformers. "
                "Install with: pip install transformers peft"
            ) from e

        self._tokenizer = AutoTokenizer.from_pretrained(
            self._model_path,
            trust_remote_code=True,
        )

        # Store compute dtype for adapter embedding conversion
        self._compute_dtype = self._resolve_dtype()

        model_kwargs = self._build_model_kwargs()
        self._base_model = AutoModelForCausalLM.from_pretrained(
            self._model_path,
            **model_kwargs,
        )

    def _build_model_kwargs(self) -> dict[str, Any]:
        """Build kwargs for AutoModelForCausalLM.from_pretrained()."""
        dtype = self._resolve_dtype()
        kwargs: dict[str, Any] = {
            "trust_remote_code": True,
            "device_map": self._config.device,
        }

        if self._config.load_in_4bit:
            self._apply_4bit_quantization(kwargs, dtype)
        else:
            kwargs["torch_dtype"] = dtype

        return kwargs

    def _apply_4bit_quantization(self, kwargs: dict[str, Any], dtype: Any) -> None:
        """Apply 4-bit quantization config if bitsandbytes available."""
        try:
            from transformers import BitsAndBytesConfig

            kwargs["quantization_config"] = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype=dtype,
            )
        except ImportError:
            self._lg.warning(
                "bitsandbytes not available, loading without 4-bit quantization"
            )
            kwargs["torch_dtype"] = dtype

    def _resolve_dtype(self) -> Any:
        """Resolve dtype string to torch dtype.

        For 'auto', uses bfloat16 if supported by the device, otherwise float16.
        """
        import torch

        if self._config.dtype == "auto":
            # Prefer bfloat16 on supported hardware (Ampere+)
            if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
                return torch.bfloat16
            return torch.float16
        dtype_map = {
            "float16": torch.float16,
            "bfloat16": torch.bfloat16,
            "float32": torch.float32,
        }
        return dtype_map.get(self._config.dtype, torch.float16)

    def _cast_prompt_embeddings(self, peft_model: Any) -> None:
        """Cast prompt embeddings to compute dtype for 4-bit compatibility.

        When using 4-bit quantization, prompt embeddings (from PROMPT_TUNING)
        may be in a different dtype than the model's compute dtype, causing
        NaN during forward pass. This casts the embeddings to match.
        """
        if self._compute_dtype is None:
            return

        # Check if this is a prompt tuning model with prompt_encoder
        if not hasattr(peft_model, "prompt_encoder"):
            return

        try:
            for name, module in peft_model.prompt_encoder.named_modules():
                if hasattr(module, "embedding") and hasattr(module.embedding, "weight"):
                    original_dtype = module.embedding.weight.dtype
                    if original_dtype != self._compute_dtype:
                        module.embedding.weight.data = module.embedding.weight.data.to(
                            self._compute_dtype
                        )
                        self._lg.trace(
                            "cast prompt embeddings",
                            extra={
                                "embedding_module": name,
                                "from_dtype": str(original_dtype),
                                "to_dtype": str(self._compute_dtype),
                            },
                        )
        except Exception as e:
            self._lg.warning("failed to cast prompt embeddings", extra={"exception": e})

    def _generate_adapter_name(self, adapter_path: str) -> str:
        """Generate a unique adapter name from path."""
        import os

        # Use the last path component as the base name
        base_name = os.path.basename(adapter_path.rstrip("/"))
        self._adapter_counter += 1
        return f"{base_name}_{self._adapter_counter}"

    def _get_adapter_peft_type(self, adapter_path: str) -> str | None:
        """Read peft_type from adapter_config.json."""
        import json
        import os

        config_path = os.path.join(adapter_path, "adapter_config.json")
        try:
            with open(config_path) as f:
                config = json.load(f)
            peft_type: str | None = config.get("peft_type")
            return peft_type
        except Exception as e:
            self._lg.trace(
                "failed to read adapter config",
                extra={"config_path": config_path, "exception": e},
            )
            return None

    # Prompt-learning types supported by this engine
    _SUPPORTED_PEFT_TYPES = {"PROMPT_TUNING", "PREFIX_TUNING", "P_TUNING"}

    def _validate_adapter_type(self, adapter_path: str) -> None:
        """Validate adapter is a prompt-learning type (not LoRA).

        LoRA adapters should use vLLM which has optimized support.
        This engine is for PROMPT_TUNING and similar types that vLLM doesn't support.
        """
        peft_type = self._get_adapter_peft_type(adapter_path)
        if peft_type is None:
            self._lg.warning(
                "could not determine peft_type, proceeding anyway",
                extra={"adapter_path": adapter_path},
            )
            return

        if peft_type not in self._SUPPORTED_PEFT_TYPES:
            raise ValueError(
                f"PEFT engine only supports prompt-learning adapters "
                f"({', '.join(sorted(self._SUPPORTED_PEFT_TYPES))}). "
                f"Got '{peft_type}'. For LoRA adapters, use vLLM engine instead "
                f"(--engine vllm-server)."
            )

    def _do_peft_load(self, adapter_path: str, adapter_name: str) -> None:
        """Load adapter onto PEFT model."""
        try:
            from peft import PeftModel
        except ImportError as e:
            raise ImportError(
                "PEFT library required for PROMPT_TUNING adapters. "
                "Install with: pip install peft"
            ) from e

        if self._peft_model is None:
            self._peft_model = PeftModel.from_pretrained(
                self._base_model,
                adapter_path,
                adapter_name=adapter_name,
                is_trainable=False,
            )
        else:
            self._peft_model.load_adapter(
                adapter_path, adapter_name=adapter_name, is_trainable=False
            )
            self._peft_model.set_adapter(adapter_name)

    def _handle_adapter_eviction(self, adapter_path: str, adapter_name: str) -> None:
        """Register adapter and handle LRU eviction."""
        to_delete = self._adapter_registry.register(adapter_path, adapter_name)
        if to_delete is not None:
            self._lg.trace(
                "deleting evicted adapter", extra={"adapter_name": to_delete}
            )
            self._peft_model.delete_adapter(to_delete)

    def _load_adapter(self, adapter_path: str) -> None:
        """Load adapter and activate it.

        Uses PEFT's multi-adapter support:
        - First adapter: PeftModel.from_pretrained()
        - Subsequent: peft_model.load_adapter() + set_adapter()
        """
        self._ensure_base_model_loaded()

        # Check if already loaded
        adapter_name = self._adapter_registry.get_adapter_name(adapter_path)
        if adapter_name is not None:
            # Already loaded, just switch to it
            if self._active_adapter != adapter_name:
                self._peft_model.set_adapter(adapter_name)
                self._active_adapter = adapter_name
                self._lg.trace(
                    "switched to cached adapter",
                    extra={"adapter_path": adapter_path, "adapter_name": adapter_name},
                )
            return

        # Validate adapter type before loading
        self._validate_adapter_type(adapter_path)

        self._lg.info("loading PEFT adapter", extra={"adapter_path": adapter_path})
        adapter_name = self._generate_adapter_name(adapter_path)
        self._do_peft_load(adapter_path, adapter_name)
        self._cast_prompt_embeddings(self._peft_model)
        self._handle_adapter_eviction(adapter_path, adapter_name)
        self._active_adapter = adapter_name

    def _get_model_for_adapter(self, adapter_path: str | None) -> Any:
        """Get model with specified adapter active, or base model if None."""
        if adapter_path is None:
            self._ensure_base_model_loaded()
            return self._base_model
        self._load_adapter(adapter_path)
        return self._peft_model

    def _resolve_adapter_path(self, lora_request: Any) -> str | None:
        """Extract adapter path from lora_request parameter.

        Handles multiple input types:
        - str: direct path or adapter key
        - object with lora_path attribute: vLLM LoRARequest
        - object with lora_name attribute: warmup LoRARequest
        - None: no adapter
        """
        if lora_request is None:
            return None
        if isinstance(lora_request, str):
            return self._maybe_resolve_adapter_key(lora_request)
        if hasattr(lora_request, "lora_path"):
            return str(lora_request.lora_path)
        if hasattr(lora_request, "lora_name"):
            return self._maybe_resolve_adapter_key(str(lora_request.lora_name))
        return self._maybe_resolve_adapter_key(str(lora_request))

    def _maybe_resolve_adapter_key(self, path_or_key: str) -> str:
        """Resolve adapter key to full path if needed.

        First checks the pre-scanned adapter map, then falls back to path resolution.
        """
        import os

        # Check pre-scanned adapter map first (canonical resolved paths)
        if path_or_key in self._adapter_paths:
            return self._adapter_paths[path_or_key]
        # If it's already a path (has directory separator or exists), use it directly
        if os.path.sep in path_or_key or os.path.exists(path_or_key):
            return path_or_key
        # Otherwise, resolve relative to adapter_base_path
        if self._config.adapter_base_path:
            return os.path.join(self._config.adapter_base_path, path_or_key)
        return path_or_key

    def _prepare_prompt(
        self,
        prompt: str,
        messages: list[dict[str, Any]] | None,
        use_chat_template: bool | None,
    ) -> str:
        """Prepare prompt from messages or apply chat template."""
        if messages:
            return self._apply_chat_template(messages)
        if use_chat_template is not False:
            return self._apply_chat_template([{"role": "user", "content": prompt}])
        return prompt

    def _tokenize_input(self, prompt: str, device: Any) -> tuple[Any, Any, int]:
        """Tokenize prompt and move to device.

        Returns:
            Tuple of (input_ids, attention_mask, prompt_tokens).
        """
        inputs = self._tokenizer(prompt, return_tensors="pt")
        input_ids = inputs["input_ids"].to(device)
        attention_mask = inputs.get("attention_mask")
        if attention_mask is not None:
            attention_mask = attention_mask.to(device)
        return input_ids, attention_mask, input_ids.shape[1]

    def _run_generation(
        self,
        model: Any,
        input_ids: Any,
        attention_mask: Any,
        gen_kwargs: dict[str, Any],
        max_tokens: int,
    ) -> Any:
        """Run model.generate() with tracing."""
        import torch
        from appinfra.time import since, start

        self._lg.trace(
            "PEFT generate input",
            extra={"input_shape": list(input_ids.shape), "max_new_tokens": max_tokens},
        )
        t0 = start()
        with torch.no_grad():
            # IMPORTANT: Pass input_ids as keyword arg - PEFT ignores positional args
            # for prompt learning methods (PROMPT_TUNING, PREFIX_TUNING, etc.)
            outputs = model.generate(
                input_ids=input_ids, attention_mask=attention_mask, **gen_kwargs
            )
        self._lg.trace(
            "PEFT generate output",
            extra={
                "output_shape": list(outputs.shape),
                "output_len": len(outputs[0]),
                "after": since(t0),
            },
        )
        return outputs

    def _decode_and_build_result(
        self, outputs: Any, prompt_tokens: int, adapter_path: str | None
    ) -> dict[str, Any]:
        """Decode generated tokens and build result dict."""
        generated_tokens = outputs[0][prompt_tokens:]
        completion_tokens = len(generated_tokens)
        content = self._tokenizer.decode(generated_tokens, skip_special_tokens=True)
        self._lg.trace(
            "PEFT generation complete",
            extra={
                "prompt_tokens": prompt_tokens,
                "generated_tokens": completion_tokens,
                "first_gen_ids": generated_tokens[:5].tolist()
                if completion_tokens > 0
                else [],
            },
        )
        result: dict[str, Any] = {
            "content": content,
            "usage": {
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
            },
        }
        if adapter_path:
            result["adapter"] = self._build_adapter_info(adapter_path)
        return result

    def generate(
        self,
        prompt: str,
        max_tokens: int = 100,
        temperature: float = 1.0,
        top_p: float = 1.0,
        top_k: int = 0,
        repetition_penalty: float = 1.0,
        use_chat_template: bool | None = None,
        stop_sequences: list[str] | None = None,
        context: RequestContext | None = None,
        messages: list[dict[str, Any]] | None = None,
        lora_request: Any | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Generate text with specified adapter (blocking).

        Args:
            prompt: Input prompt text.
            max_tokens: Maximum tokens to generate.
            temperature: Sampling temperature.
            top_p: Nucleus sampling parameter.
            top_k: Top-k sampling parameter.
            repetition_penalty: Repetition penalty.
            use_chat_template: Whether to apply chat template (handled via messages).
            stop_sequences: Sequences that stop generation.
            context: Request context for logging.
            messages: Chat messages (alternative to prompt).
            lora_request: Adapter path (str) or LoRARequest object.
            **kwargs: Additional arguments (ignored for compatibility).

        Returns:
            Dict with "content", "usage", etc.
        """
        adapter_path = self._resolve_adapter_path(lora_request)
        self._lg.trace(
            "PEFT generate called",
            extra={"adapter_path": adapter_path, "has_messages": bool(messages)},
        )
        model = self._get_model_for_adapter(adapter_path)
        prompt = self._prepare_prompt(prompt, messages, use_chat_template)
        input_ids, attention_mask, prompt_tokens = self._tokenize_input(
            prompt, model.device
        )
        gen_kwargs = self._build_generation_kwargs(
            max_tokens, temperature, top_p, top_k, repetition_penalty, stop_sequences
        )
        outputs = self._run_generation(
            model, input_ids, attention_mask, gen_kwargs, max_tokens
        )

        return self._decode_and_build_result(outputs, prompt_tokens, adapter_path)

    def generate_stream_sync(
        self,
        prompt: str,
        max_tokens: int = 100,
        temperature: float = 1.0,
        top_p: float = 1.0,
        top_k: int = 0,
        repetition_penalty: float = 1.0,
        use_chat_template: bool | None = None,
        stop_sequences: list[str] | None = None,
        context: RequestContext | None = None,
        messages: list[dict[str, Any]] | None = None,
        lora_request: Any | None = None,
        **kwargs: Any,
    ) -> PEFTStreamingIterator:
        """Stream generated tokens with specified adapter.

        Args:
            prompt: Input prompt text.
            max_tokens: Maximum tokens to generate.
            temperature: Sampling temperature.
            top_p: Nucleus sampling parameter.
            top_k: Top-k sampling parameter.
            repetition_penalty: Repetition penalty.
            use_chat_template: Whether to apply chat template (handled via messages).
            stop_sequences: Sequences that stop generation.
            context: Request context for logging.
            messages: Chat messages (alternative to prompt).
            lora_request: Adapter path (str) or LoRARequest object.
            **kwargs: Additional arguments (ignored for compatibility).

        Returns:
            PEFTStreamingIterator that yields tokens.
        """
        from transformers import TextIteratorStreamer

        adapter_path = self._resolve_adapter_path(lora_request)
        model = self._get_model_for_adapter(adapter_path)
        prompt = self._prepare_prompt(prompt, messages, use_chat_template)
        input_ids, attention_mask, prompt_tokens = self._tokenize_input(
            prompt, model.device
        )
        gen_kwargs = self._build_generation_kwargs(
            max_tokens, temperature, top_p, top_k, repetition_penalty, stop_sequences
        )
        streamer = TextIteratorStreamer(
            self._tokenizer, skip_prompt=True, skip_special_tokens=True
        )
        gen_kwargs["streamer"] = streamer
        error_holder = _GenerationErrorHolder()
        self._start_generation_thread(
            model, input_ids, attention_mask, gen_kwargs, error_holder
        )
        return PEFTStreamingIterator(
            streamer, prompt_tokens, self._tokenizer, error_holder
        )

    def _start_generation_thread(
        self,
        model: Any,
        input_ids: Any,
        attention_mask: Any,
        gen_kwargs: dict[str, Any],
        error_holder: _GenerationErrorHolder,
    ) -> None:
        """Start background thread for streaming generation."""
        import torch

        streamer = gen_kwargs.get("streamer")

        def generate() -> None:
            try:
                with torch.no_grad():
                    model.generate(
                        input_ids=input_ids, attention_mask=attention_mask, **gen_kwargs
                    )
            except Exception as e:
                error_holder.error = e
                # Signal streamer to unblock iterator
                if streamer is not None:
                    streamer.end()

        Thread(target=generate, daemon=True).start()

    def _build_stopping_criteria(self, stop_sequences: list[str]) -> Any:
        """Build stopping criteria from stop sequences."""
        from transformers import StoppingCriteriaList

        stop_ids = [
            self._tokenizer.encode(seq, add_special_tokens=False)
            for seq in stop_sequences
        ]
        return StoppingCriteriaList([_StopOnSequences(stop_ids)])

    def _build_generation_kwargs(
        self,
        max_tokens: int,
        temperature: float,
        top_p: float,
        top_k: int,
        repetition_penalty: float,
        stop_sequences: list[str] | None,
    ) -> dict[str, Any]:
        """Build kwargs for model.generate()."""
        # Use pad_token_id if set, otherwise fall back to eos_token_id
        pad_token_id = self._tokenizer.pad_token_id or self._tokenizer.eos_token_id
        kwargs: dict[str, Any] = {
            "max_new_tokens": max_tokens,
            "do_sample": temperature > 0,
            "pad_token_id": pad_token_id,
            "use_cache": False,  # Required for PROMPT_TUNING compatibility
        }
        if temperature > 0:
            kwargs["temperature"] = temperature
            kwargs["top_p"] = top_p
            if top_k > 0:
                kwargs["top_k"] = top_k
        if repetition_penalty != 1.0:
            kwargs["repetition_penalty"] = repetition_penalty
        if stop_sequences:
            kwargs["stopping_criteria"] = self._build_stopping_criteria(stop_sequences)
        return kwargs

    def _apply_chat_template(self, messages: list[dict[str, Any]]) -> str:
        """Apply chat template to messages."""
        if hasattr(self._tokenizer, "apply_chat_template"):
            result: str = self._tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
            )
            return result
        # Fallback for tokenizers without chat template
        return "\n".join(
            f"{msg.get('role', 'user')}: {msg.get('content', '')}" for msg in messages
        )

    def count_tokens(self, text: str) -> int:
        """Count tokens in text."""
        self._ensure_base_model_loaded()
        return len(self._tokenizer.encode(text))

    def memory_stats(self) -> dict[str, int | float]:
        """Return memory statistics."""
        stats: dict[str, int | float] = {
            "allocated": 0,
            "reserved": 0,
            "peak": 0,
            "cached_adapters": len(self._adapter_registry),
        }

        if self._model_loaded:
            try:
                import torch

                if torch.cuda.is_available():
                    stats["allocated"] = torch.cuda.memory_allocated()
                    stats["reserved"] = torch.cuda.memory_reserved()
                    stats["peak"] = torch.cuda.max_memory_allocated()
            except Exception:
                pass

        return stats

    def _cleanup_adapters(self) -> None:
        """Clear adapter registry and delete loaded adapters."""
        adapter_names = self._adapter_registry.clear()
        if self._peft_model is not None:
            for name in adapter_names:
                try:
                    self._peft_model.delete_adapter(name)
                except Exception:
                    pass
            del self._peft_model
            self._peft_model = None

    def _cleanup_models(self) -> None:
        """Delete base model and tokenizer."""
        if self._base_model is not None:
            del self._base_model
            self._base_model = None
        if self._tokenizer is not None:
            del self._tokenizer
            self._tokenizer = None
        self._model_loaded = False
        self._active_adapter = None

    def _force_gc(self) -> None:
        """Force garbage collection and clear CUDA cache."""
        try:
            import gc

            import torch

            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
        except Exception:
            pass

    def shutdown(self) -> None:
        """Shutdown the engine and release resources."""
        self._lg.info("shutting down PEFT engine")
        self._cleanup_adapters()
        self._cleanup_models()
        self._force_gc()
        self._lg.info("PEFT engine shutdown complete")
