# CoderAI - OpenAI-compatible API server
# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.

# AI.PROMPT: Add Vulkan backend support for AMD GPUs using llama-cpp-python
# This backend handles GGUF models on AMD GPUs via Vulkan

import os
import json
from typing import AsyncIterator, Optional, Union, List, Dict, Any
from pathlib import Path

from codai.backends.base import ModelBackend
from codai.models.utils import (
    check_hf_chat_template,
    get_model_family,
    get_reasoning_stop_tokens
)
from codai.models.cache import get_cached_model_path
from codai.models.grammar import get_tool_call_grammar, is_grammar_available

# Import global flag from coderai (will be None if not running as server)
try:
    import coderai
    _grammar_guided_gen = getattr(coderai, 'grammar_guided_gen', False)
except (ImportError, AttributeError):
    _grammar_guided_gen = False


def _make_llama_thermal_criteria():
    """A llama.cpp StoppingCriteriaList that pauses generation while too hot.

    llama-cpp-python evaluates stopping criteria synchronously per token inside
    create_(chat_)completion, so blocking here pauses the GPU forward pass —
    mid-generation thermal protection for the GGUF/Vulkan/llama.cpp backend.
    The criterion never stops generation (returns False) and is throttled so it
    doesn't read sensors on every token. Returns None if unavailable.
    """
    try:
        from llama_cpp import StoppingCriteriaList
    except Exception:
        return None

    def _pause(input_ids, logits):
        try:
            from codai.models.thermal import checkpoint
            checkpoint(context="text-gen", throttle_seconds=2.0)
        except Exception:
            pass
        return False

    try:
        return StoppingCriteriaList([_pause])
    except Exception:
        return None

try:
    from llama_cpp import Llama
    from llama_cpp.llama_chat_format import ChatFormatterResponse
    import llama_cpp as _llama_cpp
    LLAMA_CPP_AVAILABLE = True
except ImportError:
    LLAMA_CPP_AVAILABLE = False
    Llama = None
    ChatFormatterResponse = None
    _llama_cpp = None


def _install_layer_log_callback():
    """Replace llama.cpp's log callback with one that prints load-time layer/buffer
    messages directly to stdout.  Returns the callback object — keep a reference
    alive for the duration of the load so ctypes doesn't garbage-collect it."""
    if _llama_cpp is None:
        return None

    # Keywords that identify interesting load-phase messages
    _KEEP = (
        'llm_load_tensors', 'llm_load_print_meta',
        'offload', 'layer', 'buffer size', 'buffer type',
        'GPU', 'CUDA', 'Vulkan', 'Metal', 'ROCm', 'SYCL',
        'CPU', 'VRAM', 'n_layer', 'n_gpu_layers',
    )

    @_llama_cpp.llama_log_callback
    def _cb(level, text, user_data):
        try:
            msg = (text.decode('utf-8', errors='replace') if isinstance(text, bytes) else str(text)).rstrip()
            if msg and any(k in msg for k in _KEEP):
                print(f"  [llama.cpp] {msg}", flush=True)
        except Exception:
            pass

    try:
        _llama_cpp.llama_log_set(_cb, None)
    except Exception:
        return None
    return _cb  # caller must hold this reference


class VulkanBackend(ModelBackend):
    """Backend for Vulkan (AMD GPUs) using llama-cpp-python with GGUF models."""
    
    def __init__(self, original_backend: str = None):
        self.model = None
        self.model_name = None
        self.n_gpu_layers = -1  # Offload all layers to GPU by default
        self.n_ctx = 2048
        self.verbose = True
        self.main_gpu = 0  # Default to first GPU
        self.chat_template = None  # Detected chat template name
        self.hf_tokenizer = None  # HuggingFace tokenizer for apply_chat_template
        self.force_cuda = original_backend in ("nvidia", "cuda")  # Force CUDA if original was nvidia
        if self.force_cuda:
            print("DEBUG: GGUF model will use CUDA backend (forced by --backend nvidia)")
        self._last_usage: dict = {}  # usage from the most recent completion call
        self._detect_chat_template()
    
    def _detect_chat_template(self):
        """Detect the chat template used by the model."""
        try:
            # Try to get the chat template from the model
            # llama.cpp models have a chat_template attribute
            from llama_cpp.llama_chat_format import ChatFormatterResponse
            # We'll detect it when the model is loaded
            self.chat_template = "unknown"
            print("DEBUG: Chat template detection will happen after model load")
        except Exception as e:
            print(f"DEBUG: Could not initialize chat template detection: {e}")
            self.chat_template = None
    
    def _load_huggingface_tokenizer(self, template_name: str = None):
        """Load HuggingFace tokenizer for apply_chat_template support.
        
        Args:
            template_name: Optional specific template to use (e.g., 'llama3', 'chatml'). 
                          If None, will auto-detect from tokenizer.
        """
        if self.hf_tokenizer is not None:
            return  # Already loaded
        
        model_path = getattr(self, 'model_name', None)
        if not model_path:
            print("DEBUG: No model name available for HuggingFace tokenizer")
            return
        
        # If model_path is a URL, try to get the cached local path first
        if model_path.startswith('http://') or model_path.startswith('https://'):
            cached_path = get_cached_model_path(model_path)
            if cached_path and os.path.exists(cached_path):
                model_path = cached_path
                print(f"DEBUG: Using cached model path for HF tokenizer: {model_path}")
        
        try:
            from transformers import AutoTokenizer
            
            # If a specific template is provided, we can use it directly without loading tokenizer
            if template_name:
                self.chat_template = template_name
                print(f"DEBUG: Using specified chat template: {template_name}")
                # Still need to load tokenizer to get the actual template
                # but we can use the specified template name
            
            # Try to determine the model identifier
            # If model_path is a GGUF file, try to find the corresponding HF model
            if model_path.endswith('.gguf'):
                # Try to extract model name from path
                # Common patterns: .../models/llama-3.1-8b-instruct-q4_k_m.gguf
                model_dir = os.path.dirname(model_path)
                model_file = os.path.basename(model_path)
                
                # Try to find a tokenizer config in the model directory
                tokenizer_config_path = os.path.join(model_dir, 'tokenizer_config.json')
                if os.path.exists(tokenizer_config_path):
                    # Load from local directory
                    self.hf_tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
                    print(f"DEBUG: Loaded HuggingFace tokenizer from local: {model_dir}")
                    if not template_name:
                        self.chat_template = "hf_local"
                    return
                
                # Try to infer model name from file name
                # Common patterns: llama-3.1-8b-instruct-q4_k_m.gguf -> llama-3.1-8b-instruct
                # Also handle cached files with hash prefix: hash_modelname.gguf -> modelname
                model_base = model_file.replace('.gguf', '')
                
                # Remove hash prefix (64 hex chars for SHA-256 followed by underscore)
                if len(model_base) > 64 and model_base[:64].isalnum():
                    model_base = model_base[65:]  # Skip hash + underscore
                
                # Remove common quantization suffixes (case-insensitive)
                for suffix in ['_q4_k_m', '_q4_k', '_q5_k', '_q5_k_m', '_q8_0', '_f16', '_q4_0', '_q3_k_m', '_q2_k', '_Q4_K_M', '_Q4_K', '_Q5_K', '_Q5_K_M', '_Q8_0', '_F16', '_Q4_0', '_Q3_K_M', '_Q2_K']:
                    model_base = model_base.replace(suffix, '')
                
                # Try to load from HuggingFace hub
                # First try the cleaned model_base
                model_names_to_try = [model_base]
                
                # Generate shorter versions of the model name for fallback
                # E.g., Qwen3.5-27B-Uncensored-HauhauCS-Aggressive -> try shorter variants
                parts = model_base.split('-')
                if len(parts) > 1:
                    # Try progressively shorter names by removing parts from the end
                    for i in range(len(parts) - 1, 0, -1):
                        shorter_name = '-'.join(parts[:i])
                        if shorter_name and shorter_name != model_base:
                            model_names_to_try.append(shorter_name)
                
                # Also try with just the first part (e.g., "Qwen" from "Qwen3.5-27B...")
                if len(parts) > 1:
                    model_names_to_try.append(parts[0])
                
                tokenizer_loaded = False
                last_error = None
                for model_id in model_names_to_try:
                    try:
                        self.hf_tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
                        print(f"DEBUG: Loaded HuggingFace tokenizer from hub: {model_id}")
                        if not template_name:
                            self.chat_template = "hf_hub"
                        tokenizer_loaded = True
                        break
                    except Exception as fallback_err:
                        # Silently skip failed tokenizer loads - we'll use manual formatting
                        pass
                
                if tokenizer_loaded:
                    return
                
                # If HF tokenizer loading failed, try to use known template names based on model name
                # This helps when we can't find the tokenizer but know the model family
                model_base_lower = model_base.lower()
                
                # Check if this looks like a Qwen model
                known_templates_to_try = []
                if 'qwen' in model_base_lower:
                    # Try known Qwen template names in order of specificity
                    if 'qwen3.5' in model_base_lower or 'qwen3' in model_base_lower:
                        known_templates_to_try = ['qwen3', 'qwen', None]  # None means use manual formatting
                    elif 'qwen2' in model_base_lower:
                        known_templates_to_try = ['qwen2', 'qwen', None]
                    else:
                        known_templates_to_try = ['qwen', None]
                elif 'llama' in model_base_lower:
                    known_templates_to_try = ['llama3', 'llama', None]
                elif 'phi' in model_base_lower:
                    known_templates_to_try = ['phi', None]
                elif 'mistral' in model_base_lower or 'mixtral' in model_base_lower:
                    known_templates_to_try = ['mistral', None]
                
                # Try each known template - directly use the template name without loading tokenizer
                # This is the key fix: instead of trying to load more non-existent tokenizers,
                # directly set the chat_template to the known template name
                for template_name in known_templates_to_try:
                    if template_name is None:
                        # No more templates to try, use manual formatting with generic format
                        self.chat_template = "chatml"  # Use ChatML as generic fallback
                        print(f"DEBUG: No known templates worked, using generic ChatML format")
                        break
                    # Directly use this known template name - no need to load tokenizer
                    # The manual formatting will use <|im_start|> tags which work for most models
                    self.chat_template = template_name
                    print(f"DEBUG: Using known template '{template_name}' for model family detection")
                    # Successfully set template - don't try to load tokenizer
                    break
                
                if self.chat_template:
                    return
                
                # All attempts failed - warn but continue without template
                print(f"Warning: Could not load HuggingFace tokenizer for any variant of '{model_base}'")
                print(f"Warning: Will not use apply_chat_template - model will use manual formatting")
                self.chat_template = None
            else:
                # Not a GGUF file, try to load directly
                self.hf_tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
                print(f"DEBUG: Loaded HuggingFace tokenizer from: {model_path}")
                if not template_name:
                    self.chat_template = "hf"
                return
                    
        except ImportError as e:
            print(f"DEBUG: transformers not installed, cannot use HuggingFace chat template: {e}")
            self.chat_template = None
        except Exception as e:
            print(f"DEBUG: Failed to load HuggingFace tokenizer: {e}")
            self.chat_template = None
    
    def _finalize_chat_template_detection(self):
        """Finalize chat template detection after model is loaded."""
        # Check if the loaded GGUF model has a built-in chat template
        if hasattr(self, 'model') and self.model is not None:
            try:
                # llama.cpp models have a chat_template attribute
                if hasattr(self.model, 'chat_template'):
                    template = self.model.chat_template
                    if template:
                        print(f"DEBUG: Using GGUF model's built-in chat template")
                        # Extract template name if possible
                        template_str = str(template)
                        if 'llama' in template_str.lower():
                            self.chat_template = 'llama3'
                        elif 'qwen' in template_str.lower():
                            self.chat_template = 'qwen3'
                        elif 'chatml' in template_str.lower():
                            self.chat_template = 'chatml'
                        elif 'mistral' in template_str.lower():
                            self.chat_template = 'mistral'
                        else:
                            self.chat_template = 'builtin'
                        print(f"DEBUG: chat_template set to: {self.chat_template}")
                        return  # Use built-in template
            except Exception as e:
                print(f"DEBUG: Could not get chat template from GGUF model: {e}")
        
        # Fallback: Try to load HuggingFace tokenizer for chat template
    
    def _apply_chat_template(self, messages: List[Dict[str, str]], add_generation_prompt: bool = True) -> str:
        """Apply chat template to messages.
        
        Tries multiple methods in order:
        1. GGUF model's built-in chat template (via llama.cpp)
        2. HuggingFace tokenizer's apply_chat_template
        3. Manual template application based on detected template name
        4. Generic ChatML format as fallback
        
        Args:
            messages: List of message dicts or ChatMessage objects
            add_generation_prompt: Whether to add generation prompt
        """
        # Convert ChatMessage objects to dicts if needed
        if messages and hasattr(messages[0], 'model_dump'):
            # It's a Pydantic model
            messages = [msg.model_dump() if hasattr(msg, 'model_dump') else {"role": msg.role, "content": msg.content} for msg in messages]
        # First try GGUF model's built-in template
        if hasattr(self, 'model') and self.model is not None:
            try:
                if hasattr(self.model, 'create_chat_completion'):
                    # llama.cpp can handle chat templates internally
                    # We'll use it to validate the format but return our own formatted string
                    pass  # Fall through to use template directly
            except Exception as e:
                print(f"DEBUG: GGUF model template check failed: {e}")
        
        # Try HuggingFace tokenizer
        if self.hf_tokenizer:
            try:
                # Check if tokenizer has apply_chat_template
                if hasattr(self.hf_tokenizer, 'apply_chat_template'):
                    # Use the tokenizer's chat template
                    # For ChatML/Others: add_generation_prompt=True adds the assistant token
                    # For Llama3: adds the correct prefix
                    rendered = self.hf_tokenizer.apply_chat_template(
                        messages,
                        tokenize=False,
                        add_generation_prompt=add_generation_prompt
                    )
                    print(f"DEBUG: Applied HuggingFace chat template")
                    return rendered
            except Exception as e:
                print(f"DEBUG: HuggingFace apply_chat_template failed: {e}")
        
        # Try llama.cpp's built-in chat template
        if hasattr(self, 'model') and self.model is not None:
            try:
                if hasattr(self.model, 'chat'):
                    # Get the chat template from llama.cpp
                    # This works for models that have chat templates defined
                    from llama_cpp.llama_chat_format import ChatCompletionMessage
                    
                    # Try to use the model's chat handler
                    # llama.cpp's create_chat_completion handles templates internally
                    # But we need raw text for streaming
                    pass  # Fall through to manual
            except Exception as e:
                print(f"DEBUG: llama.cpp chat handling failed: {e}")
        
        # Manual template application based on detected template
        template_name = self.chat_template or "chatml"
        
        # Format messages based on template name
        if template_name in ("llama3", "llama-3", "llama-3.1"):
            return self._format_llama3(messages, add_generation_prompt)
        elif template_name in ("qwen2", "qwen2.5", "qwen3"):
            return self._format_qwen(messages, add_generation_prompt)
        elif template_name == "chatml":
            return self._format_chatml(messages, add_generation_prompt)
        elif template_name in ("mistral", "mixtral"):
            return self._format_mistral(messages, add_generation_prompt)
        else:
            # Default to ChatML format
            return self._format_chatml(messages, add_generation_prompt)
    
    def _format_llama3(self, messages: List[Dict[str, str]], add_generation_prompt: bool = True) -> str:
        """Format messages for Llama 3.x models."""
        result = []
        for msg in messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            if role == "system":
                result.append(f"<|start_header_id|>system<|end_header_id|>\n\n{content}<|eot_id|>")
            elif role == "user":
                result.append(f"<|start_header_id|>user<|end_header_id|>\n\n{content}<|eot_id|>")
            elif role == "assistant":
                result.append(f"<|start_header_id|>assistant<|end_header_id|>\n\n{content}<|eot_id|>")
        
        if add_generation_prompt:
            result.append(f"<|start_header_id|>assistant<|end_header_id|>\n\n")
        
        return "".join(result)
    
    def _format_qwen(self, messages: List[Dict[str, str]], add_generation_prompt: bool = True) -> str:
        """Format messages for Qwen models."""
        result = []
        for msg in messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            if role == "system":
                result.append(f"<|im_start|>system\n{content}<|im_end|>")
            elif role == "user":
                result.append(f"<|im_start|>user\n{content}<|im_end|>")
            elif role == "assistant":
                result.append(f"<|im_start|>assistant\n{content}<|im_end|>")
        
        if add_generation_prompt:
            result.append(f"<|im_start|>assistant\n")
        
        return "".join(result)
    
    def _format_chatml(self, messages: List[Dict[str, str]], add_generation_prompt: bool = True) -> str:
        """Format messages using ChatML format."""
        result = []
        for msg in messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            if role == "system":
                result.append(f"<|im_start|>system\n{content}<|im_end|>")
            elif role == "user":
                result.append(f"<|im_start|>user\n{content}<|im_end|>")
            elif role == "assistant":
                result.append(f"<|im_start|>assistant\n{content}<|im_end|>")
        
        if add_generation_prompt:
            result.append(f"<|im_start|>assistant\n")
        
        return "".join(result)
    
    def _format_mistral(self, messages: List[Dict[str, str]], add_generation_prompt: bool = True) -> str:
        """Format messages for Mistral/Mixtral models."""
        result = []
        for msg in messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            if role == "system":
                result.append(f"[INST] <<SYS>>\n{content}\n<</SYS>>\n\n")
            elif role == "user":
                result.append(f"[INST]{content}[/INST]")
            elif role == "assistant":
                result.append(f"{content}")
        
        # Mistral uses a different format - the last user message has [/INST]
        # and we add the assistant prefix if generation prompt is needed
        if add_generation_prompt:
            # Find the last user message and add the assistant prefix after it
            # Actually, let's reformat more carefully
            pass  # Fall through to simpler format
        
        # Simpler approach: join with newlines
        formatted = ""
        for msg in messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            if role == "system":
                formatted += f"<<SYS>>\n{content}\n<</SYS>>\n\n"
            elif role == "user":
                formatted += f"[INST] {content} [/INST]"
            elif role == "assistant":
                formatted += f" {content}"
        
        if add_generation_prompt:
            formatted += " "
        
        return formatted
    
    def load_model(self, model_path: str, model_type: str = "text", **kwargs) -> None:
        """Load a GGUF model.
        
        Args:
            model_path: Path to the GGUF model file, HuggingFace model ID, or URL
            model_type: Type of model (text, image, audio)
            **kwargs: Additional parameters
        """
        if not LLAMA_CPP_AVAILABLE:
            raise ImportError("llama-cpp-python is required for GGUF models. Install with: pip install llama-cpp-python")
        
        # Check if this looks like a GGUF model repo (e.g., "unsloth/Qwen3-Coder-30B-A3B-Instruct-GGUF")
        is_gguf_repo = not model_path.endswith('.gguf') and not os.path.exists(model_path) and not model_path.startswith('http')
        
        if is_gguf_repo:
            # Try to find GGUF files in the repository
            try:
                from huggingface_hub import list_repo_files, hf_hub_download
                from codai.models.cache import get_hf_hub_cache_dir
                print(f"DEBUG: Searching for GGUF files in {model_path}...")
                files = list(list_repo_files(model_path, repo_type="model"))
                gguf_files = [f for f in files if f.lower().endswith('.gguf')]

                if gguf_files:
                    # Prefer Q4_K_M or Q4_K quantizations, otherwise use first available
                    preferred = [f for f in gguf_files if 'q4_k_m' in f.lower() or 'q4_k' in f.lower()]
                    selected = preferred[0] if preferred else gguf_files[0]
                    print(f"DEBUG: Found GGUF files: {gguf_files}")
                    print(f"DEBUG: Selected: {selected}")
                    model_path = hf_hub_download(repo_id=model_path, filename=selected, cache_dir=kwargs.get('cache_dir') or get_hf_hub_cache_dir())
                    print(f"DEBUG: Downloaded: {model_path}")
                else:
                    print(f"Warning: No GGUF files found in {model_path}, trying direct download...")
            except Exception as e:
                print(f"Warning: Could not search HuggingFace repo: {e}")
        # If it's a HuggingFace model ID (not ending in .gguf), try to download
        elif not model_path.endswith('.gguf') and not os.path.exists(model_path):
            # Try to get from HuggingFace
            try:
                from huggingface_hub import hf_hub_download
                from codai.models.cache import get_hf_hub_cache_dir
                # Download the GGUF file
                model_path = hf_hub_download(repo_id=model_path, filename="*.gguf", cache_dir=kwargs.get('cache_dir') or get_hf_hub_cache_dir())
            except Exception as e:
                print(f"Warning: Could not download from HuggingFace: {e}")
                # Try as-is
        
        # If it's a URL, download it
        if model_path.startswith('http://') or model_path.startswith('https://'):
            cached_path = get_cached_model_path(model_path)
            if cached_path:
                model_path = cached_path
            else:
                raise ValueError(f"Could not cache model from URL: {model_path}")
        
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"Model file not found: {model_path}")
        
        self.model_name = model_path
        
        # Determine model type
        is_image = model_type == "image" or model_path.startswith("image:")
        
        # Check for --no-ram mode from global args
        no_ram = kwargs.get('no_ram', False)
        if not no_ram:
            try:
                from codai.api.state import get_global_args
                _global_args = get_global_args()
                if _global_args and getattr(_global_args, 'no_ram', False):
                    no_ram = True
            except Exception:
                pass
        
        # Configure GPU layers
        n_gpu_layers = kwargs.get('n_gpu_layers', -1)
        if no_ram:
            # --no-ram: force all layers on GPU
            self.n_gpu_layers = -1
        elif n_gpu_layers != -1:
            self.n_gpu_layers = n_gpu_layers
        
        # Configure context size
        if no_ram:
            # --no-ram: ignore --n-ctx, let the model use its own default
            self.n_ctx = 0  # 0 means use model's built-in default in llama.cpp
            print("DEBUG: --no-ram mode: ignoring --n-ctx, using model default context size")
        else:
            n_ctx = kwargs.get('n_ctx', 2048)
            self.n_ctx = n_ctx
        
        # Set verbose
        self.verbose = kwargs.get('verbose', True)
        
        # Set main GPU
        self.main_gpu = kwargs.get('main_gpu', 0)
        
        # Build kwargs for Llama constructor
        llama_kwargs = {
            'model_path': model_path,
            'n_gpu_layers': self.n_gpu_layers,
            'n_ctx': self.n_ctx,
            'verbose': self.verbose,
            'main_gpu': self.main_gpu,
        }
        
        # --no-ram: disable mmap to prevent CPU RAM usage for memory-mapped files
        if no_ram:
            llama_kwargs['use_mmap'] = False
            print("DEBUG: --no-ram mode: use_mmap=False, n_gpu_layers=-1")
        
        # Add optional parameters
        if 'n_threads' in kwargs:
            llama_kwargs['n_threads'] = kwargs['n_threads']
        if 'n_threads_batch' in kwargs:
            llama_kwargs['n_threads_batch'] = kwargs['n_threads_batch']
        if 'rope_freq_base' in kwargs:
            llama_kwargs['rope_freq_base'] = kwargs['rope_freq_base']
        if 'rope_freq_scale' in kwargs:
            llama_kwargs['rope_freq_scale'] = kwargs['rope_freq_scale']
        
        # Force CUDA if requested
        if self.force_cuda:
            # Set environment variable to force CUDA
            os.environ['CUDA_VISIBLE_DEVICES'] = '0'
            # llama-cpp-python will use CUDA when available
        
        # Pre-load summary
        gpu_label = "all" if self.n_gpu_layers == -1 else str(self.n_gpu_layers)
        print(f"  n_gpu_layers : {gpu_label}  |  n_ctx : {self.n_ctx}  |  main_gpu : {self.main_gpu}")
        if _llama_cpp:
            gpu_supported = _llama_cpp.llama_supports_gpu_offload()
            print(f"  GPU offload  : {'supported' if gpu_supported else 'NOT supported by this build'}")

        _log_cb = _install_layer_log_callback()
        try:
            self.model = Llama(**llama_kwargs)
        except Exception as e:
            print(f"Error loading GGUF model: {e}")
            raise
        finally:
            # Restore llama.cpp's default (quiet) logging after load
            if _llama_cpp:
                try:
                    _llama_cpp.llama_log_set(None, None)
                except Exception:
                    pass
            _log_cb = None  # release callback

        # Post-load layer/buffer summary
        try:
            n_total = _llama_cpp.llama_model_n_layer(self.model.model)
            n_gpu_actual = n_total if self.n_gpu_layers == -1 else min(self.n_gpu_layers, n_total)
            n_cpu = n_total - n_gpu_actual
            print(f"  Layers total : {n_total}")
            print(f"  Layers → GPU : {n_gpu_actual}  |  Layers → CPU : {n_cpu}")
        except Exception:
            pass

        # Try to detect and set up chat template
        self._finalize_chat_template_detection()
        print(f"  chat_template: {self.chat_template}")
    
    def generate(
        self,
        prompt: str,
        max_tokens: Optional[int] = None,
        temperature: float = 0.7,
        top_p: float = 1.0,
        stop: Optional[List[str]] = None,
        grammar: Optional[str] = None,
        repeat_penalty: float = 1.0,
        presence_penalty: float = 0.0,
        frequency_penalty: float = 0.0
    ) -> str:
        """Generate text non-streaming.
        
        Args:
            prompt: Input prompt (or messages for chat)
            max_tokens: Maximum tokens to generate
            temperature: Sampling temperature
            top_p: Top-p sampling
            stop: Stop sequences
            grammar: Optional GBNF grammar string for constrained generation
            repeat_penalty: Repetition penalty (1.0 = no penalty)
            presence_penalty: Presence penalty
            frequency_penalty: Frequency penalty
            
        Returns:
            Generated text
        """
        kwargs = {}
        if max_tokens is not None:
            kwargs['max_tokens'] = max_tokens
        if temperature is not None:
            kwargs['temperature'] = temperature
        if top_p is not None:
            kwargs['top_p'] = top_p
        if stop is not None:
            kwargs['stop'] = stop
        if self.model is None:
            raise RuntimeError("Model not loaded")
        
        # Check if prompt looks like messages (list of dicts)
        if isinstance(prompt, list):
            prompt = self._apply_chat_template(prompt, add_generation_prompt=True)
        
        # Set defaults
        max_tokens = kwargs.get('max_tokens', 256)
        temperature = kwargs.get('temperature', 0.7)
        top_p = kwargs.get('top_p', 0.9)
        top_k = kwargs.get('top_k', 40)
        repeat_penalty = kwargs.get('repeat_penalty', 1.1)
        stream = kwargs.get('stream', False)
        
        # Get stop strings
        stop = kwargs.get('stop', None)
        if stop is None:
            # Get default stop tokens based on template
            stop = get_reasoning_stop_tokens(self.chat_template)
        
        # Check for grammar-guided generation
        use_grammar = grammar
        if use_grammar is None:
            # Check global flag
            try:
                import coderai
                if getattr(coderai, 'grammar_guided_gen', False):
                    use_grammar = get_tool_call_grammar()
            except (ImportError, AttributeError):
                pass
        
        try:
            result = self.model.create_completion(
                stopping_criteria=_make_llama_thermal_criteria(),
                prompt=prompt,
                max_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
                top_k=top_k,
                repeat_penalty=repeat_penalty,
                stop=stop,
                grammar=use_grammar,
            )
            usage = result.get('usage', {})
            self._store_usage(usage.get('prompt_tokens', 0), usage.get('completion_tokens', 0))
            return result['choices'][0]['text']
        except Exception as e:
            # If grammar generation fails, fall back to normal generation
            if use_grammar:
                print(f"Warning: Grammar-guided generation failed: {e}, falling back to normal generation")
                try:
                    result = self.model.create_completion(
                        stopping_criteria=_make_llama_thermal_criteria(),
                        prompt=prompt,
                        max_tokens=max_tokens,
                        temperature=temperature,
                        top_p=top_p,
                        top_k=top_k,
                        repeat_penalty=repeat_penalty,
                        stop=stop,
                    )
                    usage = result.get('usage', {})
                    self._store_usage(usage.get('prompt_tokens', 0), usage.get('completion_tokens', 0))
                    return result['choices'][0]['text']
                except Exception as e2:
                    print(f"Error during fallback generation: {e2}")
                    raise
            print(f"Error during generation: {e}")
            raise
    
    async def generate_stream(
        self,
        prompt: str,
        max_tokens: Optional[int] = None,
        temperature: float = 0.7,
        top_p: float = 1.0,
        stop: Optional[List[str]] = None,
        grammar: Optional[str] = None,
        repeat_penalty: float = 1.0,
        presence_penalty: float = 0.0,
        frequency_penalty: float = 0.0
    ) -> AsyncIterator[str]:
        """Generate text with streaming.
        
        Args:
            prompt: Input prompt (or messages for chat)
            max_tokens: Maximum tokens to generate
            temperature: Sampling temperature
            top_p: Top-p sampling
            stop: Stop sequences
            grammar: Optional GBNF grammar string
            repeat_penalty: Repetition penalty (1.0 = no penalty)
            presence_penalty: Presence penalty
            frequency_penalty: Frequency penalty
            
        Yields:
            Generated text chunks
        """
        if self.model is None:
            raise RuntimeError("Model not loaded")
        
        # Check if prompt looks like messages (list of dicts)
        if isinstance(prompt, list):
            prompt = self._apply_chat_template(prompt, add_generation_prompt=True)
        
        # Set defaults
        max_tokens = max_tokens if max_tokens is not None else 256
        temperature = temperature if temperature is not None else 0.7
        top_p = top_p if top_p is not None else 0.9
        top_k = 40
        repeat_penalty = repeat_penalty if repeat_penalty is not None else 1.1
        grammar = grammar
        
        # Get stop strings
        stop = stop if stop is not None else None
        if stop is None:
            # Get default stop tokens based on template
            stop = get_reasoning_stop_tokens(self.chat_template)
        
        # Check for grammar-guided generation
        use_grammar = grammar
        if use_grammar is None:
            # Check global flag
            try:
                import coderai
                if getattr(coderai, 'grammar_guided_gen', False):
                    use_grammar = get_tool_call_grammar()
            except (ImportError, AttributeError):
                pass
        
        try:
            # For chat, we need to extract just the new text from each chunk
            # The first chunk will have the full prompt + first token
            # Subsequent chunks only have new tokens
            
            first_chunk = True
            prompt_len = len(prompt) if isinstance(prompt, str) else 0
            
            for chunk in self.model.create_completion(
                stopping_criteria=_make_llama_thermal_criteria(),
                prompt=prompt,
                max_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
                top_k=top_k,
                repeat_penalty=repeat_penalty,
                stop=stop,
                stream=True,
                grammar=use_grammar,
            ):
                text = chunk['choices'][0].get('text', '')
                
                if first_chunk:
                    # Skip the prompt text on first chunk
                    # The first chunk includes the full prompt plus the first new token
                    if text and len(text) > prompt_len:
                        # Extract just the new part
                        new_text = text[prompt_len:]
                        if new_text:
                            yield new_text
                    first_chunk = False
                else:
                    # Subsequent chunks only have new tokens
                    if text:
                        yield text
                
                # Check for stop
                if chunk['choices'][0].get('finish_reason'):
                    break
                    
        except Exception as e:
            # If grammar generation fails, fall back to normal generation
            if use_grammar:
                print(f"Warning: Grammar-guided streaming generation failed: {e}, falling back to normal generation")
                try:
                    first_chunk = True
                    prompt_len = len(prompt) if isinstance(prompt, str) else 0
                    
                    for chunk in self.model.create_completion(
                        stopping_criteria=_make_llama_thermal_criteria(),
                        prompt=prompt,
                        max_tokens=max_tokens,
                        temperature=temperature,
                        top_p=top_p,
                        top_k=top_k,
                        repeat_penalty=repeat_penalty,
                        stop=stop,
                        stream=True,
                    ):
                        text = chunk['choices'][0].get('text', '')
                        
                        if first_chunk:
                            if text and len(text) > prompt_len:
                                new_text = text[prompt_len:]
                                if new_text:
                                    yield new_text
                            first_chunk = False
                        else:
                            if text:
                                yield text
                        
                        if chunk['choices'][0].get('finish_reason'):
                            break
                except Exception as e2:
                    print(f"Error during fallback streaming generation: {e2}")
                    raise
            else:
                print(f"Error during streaming generation: {e}")
                raise
    
    def chat(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> Dict[str, Any]:
        """Chat completion.
        
        Args:
            messages: List of message dicts with 'role' and 'content'
            **kwargs: Generation parameters
            
        Returns:
            Response dict with 'content' and metadata
        """
        if self.model is None:
            raise RuntimeError("Model not loaded")
        
        # Apply chat template
        prompt = self._apply_chat_template(messages, add_generation_prompt=True)
        
        # Generate
        max_tokens = kwargs.get('max_tokens', 256)
        temperature = kwargs.get('temperature', 0.7)
        top_p = kwargs.get('top_p', 0.9)
        repeat_penalty = kwargs.get('repeat_penalty', 1.1)
        
        stop = kwargs.get('stop', None)
        if stop is None:
            stop = get_reasoning_stop_tokens(self.chat_template)
        
        stream = kwargs.get('stream', False)
        
        if stream:
            # Return iterator for streaming
            async def generate_stream():
                first_chunk = True
                prompt_len = len(prompt)
                
                for chunk in self.model.create_completion(
                    stopping_criteria=_make_llama_thermal_criteria(),
                    prompt=prompt,
                    max_tokens=max_tokens,
                    temperature=temperature,
                    top_p=top_p,
                    repeat_penalty=repeat_penalty,
                    stop=stop,
                    stream=True,
                ):
                    text = chunk['choices'][0].get('text', '')
                    
                    if first_chunk:
                        if text and len(text) > prompt_len:
                            new_text = text[prompt_len:]
                            if new_text:
                                yield new_text
                        first_chunk = False
                    else:
                        if text:
                            yield text
                    
                    if chunk['choices'][0].get('finish_reason'):
                        break
            
            return {"stream": generate_stream(), "content": ""}
        else:
            result = self.model.create_completion(
                stopping_criteria=_make_llama_thermal_criteria(),
                prompt=prompt,
                max_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
                repeat_penalty=repeat_penalty,
                stop=stop,
            )
            
            content = result['choices'][0]['text']
            
            return {
                "content": content,
                "usage": result.get('usage', {}),
                "model": self.model_name,
            }
    
    async def embed(self, text: str) -> List[float]:
        """Generate embeddings.
        
        Args:
            text: Input text
            
        Returns:
            Embedding vector
        """
        if self.model is None:
            raise RuntimeError("Model not loaded")
        
        try:
            result = self.model.create_embedding(text)
            return result['data'][0]['embedding']
        except Exception as e:
            print(f"Error generating embeddings: {e}")
            raise
    
    def unload_model(self):
        """Unload the model from memory."""
        if self.model is not None:
            del self.model
            self.model = None
            self.model_name = None
            print("DEBUG: VulkanBackend unloaded model")
    
    @property
    def is_loaded(self) -> bool:
        """Check if model is loaded."""
        return self.model is not None
    
    def get_model_info(self) -> Dict[str, Any]:
        """Get information about the loaded model."""
        if self.model is None:
            return {"loaded": False}
        
        return {
            "loaded": True,
            "model_name": self.model_name,
            "chat_template": self.chat_template,
            "n_ctx": self.n_ctx,
            "n_gpu_layers": self.n_gpu_layers,
        }

    # ------------------------------------------------------------------
    # Usage / cache helpers
    # ------------------------------------------------------------------

    def _read_cached_tokens(self, prompt_tokens: int) -> int:
        """Extract cached token count from llama.cpp timings after a completion."""
        try:
            timings = getattr(self.model, 'timings', None)
            if timings is None:
                # Try the internal context if timings property not exposed
                ctx = getattr(self.model, '_ctx', None)
                if ctx and hasattr(ctx, 'timings'):
                    timings = ctx.timings()
            if timings is not None:
                n_p_eval = getattr(timings, 'n_p_eval', None)
                if n_p_eval is not None:
                    return max(0, prompt_tokens - int(n_p_eval))
        except Exception:
            pass
        return 0

    def _store_usage(self, prompt_tokens: int, completion_tokens: int) -> None:
        cached = self._read_cached_tokens(prompt_tokens)
        self._last_usage = {
            'prompt_tokens': prompt_tokens,
            'completion_tokens': completion_tokens,
            'total_tokens': prompt_tokens + completion_tokens,
            'cached_tokens': cached,
        }

    def get_last_usage(self) -> dict:
        """Return usage dict from the most recent completion (includes cached_tokens)."""
        return dict(self._last_usage)

    # ------------------------------------------------------------------
    # Chat-level generation (uses llama.cpp native chat template)
    # ------------------------------------------------------------------

    def generate_chat(self, messages, max_tokens=None, temperature=0.7, top_p=1.0,
                      stop=None, tools=None, response_format=None):
        """Non-streaming chat completion using llama.cpp's native chat handler."""
        if self.model is None:
            raise RuntimeError("Model not loaded")

        kwargs = dict(
            messages=messages,
            max_tokens=max_tokens or 512,
            temperature=temperature,
            top_p=top_p,
        )
        if stop:
            kwargs['stop'] = stop
        if response_format and response_format.get('type') == 'json_object':
            kwargs['response_format'] = {'type': 'json_object'}
        _tc = _make_llama_thermal_criteria()
        if _tc is not None:
            kwargs['stopping_criteria'] = _tc

        result = self.model.create_chat_completion(**kwargs)
        usage = result.get('usage', {})
        self._store_usage(
            prompt_tokens=usage.get('prompt_tokens', 0),
            completion_tokens=usage.get('completion_tokens', 0),
        )
        content = result['choices'][0]['message'].get('content') or ''
        return content

    async def generate_chat_stream(self, messages, max_tokens=None, temperature=0.7,
                                   top_p=1.0, stop=None, tools=None, response_format=None):
        """Streaming chat completion using llama.cpp's native chat handler."""
        if self.model is None:
            raise RuntimeError("Model not loaded")

        kwargs = dict(
            messages=messages,
            max_tokens=max_tokens or 512,
            temperature=temperature,
            top_p=top_p,
            stream=True,
        )
        if stop:
            kwargs['stop'] = stop
        _tc = _make_llama_thermal_criteria()
        if _tc is not None:
            kwargs['stopping_criteria'] = _tc

        prompt_tokens = 0
        completion_tokens = 0
        try:
            for chunk in self.model.create_chat_completion(**kwargs):
                delta = chunk['choices'][0].get('delta', {})
                text = delta.get('content') or ''
                if text:
                    completion_tokens += 1
                    yield text
                # Capture usage if present in final streaming chunk
                if chunk.get('usage'):
                    u = chunk['usage']
                    prompt_tokens = u.get('prompt_tokens', 0)
                    completion_tokens = u.get('completion_tokens', completion_tokens)
                if chunk['choices'][0].get('finish_reason'):
                    break
        finally:
            # Timings are available after the stream is exhausted
            if prompt_tokens == 0:
                # Estimate from word split if llama.cpp didn't report
                prompt_tokens = sum(
                    len(str(m.get('content', '')).split())
                    for m in messages
                )
            self._store_usage(prompt_tokens, completion_tokens)

    def get_model_name(self) -> str:
        """Return the loaded model name."""
        return self.model_name or "unknown"

    def format_messages(self, messages) -> str:
        """Format messages into a prompt string using chat template."""
        if isinstance(messages, list):
            return self._apply_chat_template(messages, add_generation_prompt=True)
        return str(messages)

    def cleanup(self) -> None:
        """Cleanup resources."""
        self.unload_model()
    
    def get_context_size(self) -> int:
        """Return the model's context window size."""
        return self.n_ctx