from .config import HF_TOKEN_ENV, get_env_var, set_env_var
import os
import logging
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from typing import Optional, Generator, Dict, Any, List, Union, Callable, AsyncGenerator
from fastapi import HTTPException
import time
from .config import (
    MODEL_REGISTRY, DEFAULT_MODEL, DEFAULT_MAX_LENGTH, DEFAULT_TEMPERATURE, DEFAULT_TOP_P,
    ENABLE_ATTENTION_SLICING, ENABLE_CPU_OFFLOADING, ENABLE_FLASH_ATTENTION,
    ENABLE_BETTERTRANSFORMER, ENABLE_QUANTIZATION, QUANTIZATION_TYPE, UNLOAD_UNUSED_MODELS, MODEL_TIMEOUT,
)
from .logger.logger import logger, log_model_loaded, log_model_unloaded
from .utils import check_resource_availability, get_device, format_model_size
import gc
from colorama import Fore, Style
import asyncio
import re
import zipfile
import tempfile
import json

QUANTIZATION_SETTINGS = {
    "fp16": {
        "load_in_8bit": False,
        "load_in_4bit": False,
        "torch_dtype": torch.float16,
        "device_map": "auto"
    },
    "int8": {
        "load_in_8bit": True,
        "load_in_4bit": False,
        "device_map": "auto"
    },
    "int4": {
        "load_in_8bit": False,
        "load_in_4bit": True,
        "device_map": "auto"
    }
}


class ModelManager:
    def __init__(self):
        self.model = None
        self.tokenizer = None
        self.current_model = None
        self._loading = False
        self._last_use = time.time()  # Initialize _last_use
        self.response_cache = {}  # Add cache dictionary

    @property
    def last_used(self) -> float:
        """Get the timestamp of last model use"""
        return self._last_use

    @last_used.setter
    def last_used(self, value: float):
        """Set the timestamp of last model use"""
        self._last_use = value

    def _get_quantization_config(self) -> Optional[Dict[str, Any]]:
        """Get quantization configuration based on settings"""
        # Check if quantization is explicitly disabled (not just False but also '0', 'none', '')
        # Use the config system to get the value, which checks environment variables and config file
        from .cli.config import get_config_value

        # First check if CUDA is available - if not, we can't use bitsandbytes quantization
        if not torch.cuda.is_available():
            logger.warning("CUDA not available - quantization with bitsandbytes requires CUDA")
            logger.info("Disabling quantization and using CPU-compatible settings")
            return {
                "torch_dtype": torch.float32,
                "device_map": "auto"
            }

        enable_quantization = get_config_value('enable_quantization', ENABLE_QUANTIZATION)
        # Convert string values to boolean if needed
        if isinstance(enable_quantization, str):
            enable_quantization = enable_quantization.lower() not in ('false', '0', 'none', '')

        quantization_type = get_config_value('quantization_type', QUANTIZATION_TYPE)

        if not enable_quantization:
            logger.info("Quantization is disabled, using default precision")
            return {
                "torch_dtype": torch.float16 if torch.cuda.is_available() else torch.float32,
                "device_map": "auto"
            }

        try:
            import bitsandbytes as bnb
            from packaging import version

            if version.parse(bnb.__version__) < version.parse("0.41.1"):
                logger.warning(
                    f"bitsandbytes version {bnb.__version__} may not support all quantization features. "
                    "Please upgrade to version 0.41.1 or higher."
                )
                return {
                    "torch_dtype": torch.float16 if torch.cuda.is_available() else torch.float32,
                    "device_map": "auto"
                }

            # Check for empty quantization type
            if not quantization_type or quantization_type.lower() in ('none', ''):
                logger.info(
                    "No quantization type specified, defaulting to fp16")
                return {
                    "torch_dtype": torch.float16 if torch.cuda.is_available() else torch.float32,
                    "device_map": "auto"
                }

            if quantization_type == "int8":
                logger.info("Using INT8 quantization")
                return {
                    "device_map": "auto",
                    "quantization_config": BitsAndBytesConfig(
                        load_in_8bit=True,
                        llm_int8_threshold=6.0,
                        bnb_8bit_compute_dtype=torch.float16,
                        bnb_8bit_use_double_quant=True
                    )
                }
            elif quantization_type == "int4":
                logger.info("Using INT4 quantization")
                return {
                    "device_map": "auto",
                    "quantization_config": BitsAndBytesConfig(
                        load_in_4bit=True,
                        bnb_4bit_compute_dtype=torch.float16,
                        bnb_4bit_use_double_quant=True,
                        bnb_4bit_quant_type="nf4"
                    )
                }
            else:
                logger.info(f"Unrecognized quantization type '{quantization_type}', defaulting to fp16")
                return {
                    "torch_dtype": torch.float16 if torch.cuda.is_available() else torch.float32,
                    "device_map": "auto"
                }

        except ImportError:
            logger.warning(
                "bitsandbytes package not found or incompatible. "
                "Falling back to fp16. Please install bitsandbytes>=0.41.1 for quantization support."
            )
            return {
                "torch_dtype": torch.float16 if torch.cuda.is_available() else torch.float32,
                "device_map": "auto"
            }

    def _apply_optimizations(self, model: AutoModelForCausalLM) -> AutoModelForCausalLM:
        """Apply various optimizations to the model"""
        try:
            # Import the config system
            from .cli.config import get_config_value

            # Only apply attention slicing if explicitly enabled and not empty
            enable_attention_slicing = get_config_value('enable_attention_slicing', ENABLE_ATTENTION_SLICING)
            if isinstance(enable_attention_slicing, str):
                enable_attention_slicing = enable_attention_slicing.lower() not in ('false', '0', 'none', '')

            if enable_attention_slicing:
                if hasattr(model, 'enable_attention_slicing'):
                    # Use more aggressive slicing for faster inference
                    model.enable_attention_slicing("max")
                    logger.info("Attention slicing enabled with max setting")
                else:
                    logger.info(
                        "Attention slicing not available for this model")

            # Only apply CPU offloading if explicitly enabled and not empty
            enable_cpu_offloading = get_config_value('enable_cpu_offloading', ENABLE_CPU_OFFLOADING)
            if isinstance(enable_cpu_offloading, str):
                enable_cpu_offloading = enable_cpu_offloading.lower() not in ('false', '0', 'none', '')

            if enable_cpu_offloading:
                if hasattr(model, "enable_cpu_offload"):
                    model.enable_cpu_offload()
                    logger.info("CPU offloading enabled")
                else:
                    logger.info("CPU offloading not available for this model")

            # Only apply BetterTransformer if explicitly enabled and not empty
            enable_bettertransformer = get_config_value('enable_better_transformer', ENABLE_BETTERTRANSFORMER)
            if isinstance(enable_bettertransformer, str):
                enable_bettertransformer = enable_bettertransformer.lower() not in ('false', '0', 'none', '')

            if enable_bettertransformer:
                try:
                    from optimum.bettertransformer import BetterTransformer
                    model = BetterTransformer.transform(model)
                    logger.info("BetterTransformer optimization applied")
                except ImportError:
                    logger.warning(
                        "BetterTransformer not available - install 'optimum' for this feature")
                except Exception as e:
                    logger.warning(
                        f"BetterTransformer optimization failed: {str(e)}")

            # Only apply Flash Attention if explicitly enabled and not empty
            enable_flash_attention = get_config_value('enable_flash_attention', ENABLE_FLASH_ATTENTION)
            if isinstance(enable_flash_attention, str):
                enable_flash_attention = enable_flash_attention.lower() not in ('false', '0', 'none', '')

            if enable_flash_attention:
                try:
                    # Try to enable flash attention directly on the model config
                    if hasattr(model.config, "attn_implementation"):
                        model.config.attn_implementation = "flash_attention_2"
                        logger.info("Flash Attention 2 enabled via config")
                    # For older models, try the flash_attn module
                    else:
                        import flash_attn
                        logger.info("Flash Attention enabled via module")
                except ImportError:
                    logger.warning(
                        "Flash Attention not available - install 'flash-attn' for this feature")
                except Exception as e:
                    logger.warning(
                        f"Flash Attention optimization failed: {str(e)}")

            # Enable memory efficient attention if available
            try:
                if hasattr(model, "enable_xformers_memory_efficient_attention"):
                    model.enable_xformers_memory_efficient_attention()
                    logger.info("XFormers memory efficient attention enabled")
            except Exception as e:
                logger.info(f"XFormers memory efficient attention not available: {str(e)}")

            # Enable gradient checkpointing for memory efficiency if available
            try:
                if hasattr(model, "gradient_checkpointing_enable"):
                    model.gradient_checkpointing_enable()
                    logger.info("Gradient checkpointing enabled for memory efficiency")
            except Exception as e:
                logger.info(f"Gradient checkpointing not available: {str(e)}")

            # Set model to evaluation mode for faster inference
            model.eval()
            logger.info("Model set to evaluation mode for faster inference")

            return model
        except Exception as e:
            logger.warning(f"Some optimizations could not be applied: {str(e)}")
            return model

    async def _load_model_with_optimizations(self, model_id: str):
        """Load and optimize a model with all configured optimizations"""
        try:
            # Get HF token
            from .config import get_hf_token
            hf_token = get_hf_token(interactive=False)

            # Apply quantization settings
            quant_config = self._get_quantization_config()

            # Load tokenizer first
            self.tokenizer = AutoTokenizer.from_pretrained(
                model_id,
                token=hf_token if hf_token else None
            )

            # Load model with optimizations
            self.model = AutoModelForCausalLM.from_pretrained(
                model_id,
                token=hf_token if hf_token else None,
                **quant_config
            )

            # Apply additional optimizations
            self.model = self._apply_optimizations(self.model)

            # Set model to evaluation mode
            self.model.eval()

            return self.model

        except Exception as e:
            logger.error(f"Error loading model: {str(e)}")
            raise

    async def load_model(self, model_id: str) -> None:
        """Load a model but don't persist it to config"""
        if self._loading:
            raise RuntimeError("Another model is currently loading")

        self._loading = True

        try:
            # Unload current model if any
            if self.model:
                self.unload_model()

            # Load the new model
            logger.info(f"Loading model: {model_id}")
            start_time = time.time()

            # Apply optimizations based on config
            self.model = await self._load_model_with_optimizations(model_id)
            self.current_model = model_id

            load_time = time.time() - start_time
            logger.info(f"Model loaded in {load_time:.2f} seconds")

            # Log the model load but don't persist to config
            log_model_loaded(model_id, load_time)

        except Exception as e:
            logger.error(f"Error loading model: {str(e)}")
            raise
        finally:
            self._loading = False

    def check_model_timeout(self):
        """Check if model should be unloaded due to inactivity"""
        if not UNLOAD_UNUSED_MODELS or not self.model:
            return

        if time.time() - self.last_used > MODEL_TIMEOUT:
            logger.info(f"Unloading model {self.current_model} due to inactivity")
            model_id = self.current_model
            del self.model
            self.model = None
            self.current_model = None
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            log_model_unloaded(model_id)

    async def generate(
        self,
        prompt: str,
        stream: bool = False,
        max_length: Optional[int] = None,
        max_new_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        top_k: Optional[int] = None,
        repetition_penalty: Optional[float] = None,
        system_instructions: Optional[str] = None
    ) -> str:
        """Generate text from the model"""
        # Check model timeout
        self.check_model_timeout()

        if not self.model or not self.tokenizer:
            raise HTTPException(
                status_code=400, 
                detail="No model is currently loaded. Please load a model first using the /models/load endpoint."
            )

        self.last_used = time.time()

        try:
            # Get appropriate system instructions
            from .config import system_instructions
            instructions = str(system_instructions.get_instructions(
                self.current_model)) if not system_instructions else str(system_instructions)

            # Format prompt with system instructions
            formatted_prompt = f"""<|system|>{instructions}</|system|>\n<|user|>{prompt}</|user|>\n<|assistant|>"""

            # Check cache for non-streaming requests with default parameters
            cache_key = None
            if not stream and not any([max_length, max_new_tokens, temperature, top_p, top_k, repetition_penalty]):
                cache_key = f"{self.current_model}:{hash(formatted_prompt)}"
                if cache_key in self.response_cache:
                    logger.info(f"Cache hit for prompt: {prompt[:30]}...")
                    return self.response_cache[cache_key]

            # Get model-specific generation parameters
            from .config import get_model_generation_params
            gen_params = get_model_generation_params(self.current_model)

            # Set optimized defaults for faster generation
            if not max_length and not max_new_tokens:
                # Use a smaller default max_length for faster generation
                gen_params["max_length"] = min(gen_params.get("max_length", DEFAULT_MAX_LENGTH), 512)

            if not temperature:
                # Slightly lower temperature for faster and more focused generation
                gen_params["temperature"] = min(gen_params.get("temperature", DEFAULT_TEMPERATURE), 0.7)

            if not top_k:
                # Add top_k for faster sampling
                gen_params["top_k"] = 40

            # Handle max_new_tokens parameter (map to max_length)
            if max_new_tokens is not None:
                max_length = max_new_tokens

            # Override with user-provided parameters if specified
            if max_length is not None:
                try:
                    gen_params["max_length"] = int(max_length)
                except (ValueError, TypeError):
                    logger.warning(f"Invalid max_length value: {max_length}. Using model default.")
            if temperature is not None:
                try:
                    gen_params["temperature"] = float(temperature)
                except (ValueError, TypeError):
                    logger.warning(f"Invalid temperature value: {temperature}. Using model default.")
            if top_p is not None:
                try:
                    gen_params["top_p"] = float(top_p)
                except (ValueError, TypeError):
                    logger.warning(f"Invalid top_p value: {top_p}. Using model default.")
            if top_k is not None:
                try:
                    gen_params["top_k"] = int(top_k)
                except (ValueError, TypeError):
                    logger.warning(f"Invalid top_k value: {top_k}. Using model default.")
            if repetition_penalty is not None:
                try:
                    gen_params["repetition_penalty"] = float(
                        repetition_penalty)
                except (ValueError, TypeError):
                    logger.warning(f"Invalid repetition_penalty value: {repetition_penalty}. Using model default.")

            inputs = self.tokenizer(formatted_prompt, return_tensors="pt")

            # Get the actual device of the model
            model_device = next(self.model.parameters()).device

            # Move inputs to the same device as the model
            for key in inputs:
                inputs[key] = inputs[key].to(model_device)

            if stream:
                return self.async_stream_generate(inputs, gen_params)

            # Check if we need to clear CUDA cache before generation
            if torch.cuda.is_available():
                current_mem = torch.cuda.memory_allocated() / (1024 * 1024 * 1024)  # GB
                total_mem = torch.cuda.get_device_properties(0).total_memory / (1024 * 1024 * 1024)  # GB
                if current_mem > 0.8 * total_mem:  # If using >80% of GPU memory
                    # Clear cache to avoid OOM
                    torch.cuda.empty_cache()
                    logger.info("Cleared CUDA cache before generation to avoid out of memory error")

            with torch.no_grad():
                try:
                    generate_params = {
                        **inputs,
                        "max_new_tokens": gen_params["max_length"],
                        "temperature": gen_params["temperature"],
                        "top_p": gen_params["top_p"],
                        "do_sample": True,
                        "pad_token_id": self.tokenizer.eos_token_id,
                        # Fix the early stopping warning by setting num_beams explicitly
                        "num_beams": 1
                    }

                    # Add optional parameters if present in gen_params
                    if "top_k" in gen_params:
                        generate_params["top_k"] = gen_params["top_k"]
                    if "repetition_penalty" in gen_params:
                        generate_params["repetition_penalty"] = gen_params["repetition_penalty"]

                    # Set a reasonable max time for generation to prevent hanging
                    if "max_time" not in generate_params and not stream:
                        generate_params["max_time"] = 30.0  # 30 seconds max for generation

                    # Use efficient attention implementation if available
                    if hasattr(self.model.config, "attn_implementation"):
                        generate_params["attn_implementation"] = "flash_attention_2"

                    # Generate text
                    start_time = time.time()
                    outputs = self.model.generate(**generate_params)
                    generation_time = time.time() - start_time
                    logger.info(f"Generation completed in {generation_time:.2f} seconds")

                except RuntimeError as e:
                    if "CUDA out of memory" in str(e):
                        # If we run out of memory, clear cache and try again with smaller parameters
                        torch.cuda.empty_cache()
                        logger.warning("CUDA out of memory during generation. Cleared cache and reducing parameters.")

                        # Reduce parameters for memory efficiency
                        generate_params["max_new_tokens"] = min(generate_params.get("max_new_tokens", 512), 256)

                        # Try again with reduced parameters
                        outputs = self.model.generate(**generate_params)
                    else:
                        # For other errors, re-raise
                        raise

            # Decode the raw response without any formatting
            response = self.tokenizer.decode(
                outputs[0][len(inputs["input_ids"][0]):], skip_special_tokens=True)

            # Cache the response if we have a cache key
            if cache_key:
                self.response_cache[cache_key] = response
                # Limit cache size to prevent memory issues
                if len(self.response_cache) > 100:
                    # Remove oldest entries
                    for _ in range(10):
                        self.response_cache.pop(next(iter(self.response_cache)), None)

            return response

        except Exception as e:
            logger.error(f"Generation failed: {str(e)}")
            raise HTTPException(
                status_code=500, detail=f"Generation failed: {str(e)}")


    def _stream_generate(
        self,
        inputs: Dict[str, torch.Tensor],
        max_length: int = None,
        temperature: float = None,
        top_p: float = None,
        gen_params: Optional[Dict[str, Any]] = None
    ) -> Generator[str, None, None]:
        """Stream generate text from the model with optimizations for low-resource computers"""
        try:
            # If gen_params is provided, use it instead of individual parameters
            if gen_params is not None:
                max_length = gen_params.get("max_length", DEFAULT_MAX_LENGTH)
                temperature = gen_params.get(
                    "temperature", DEFAULT_TEMPERATURE)
                top_p = gen_params.get("top_p", DEFAULT_TOP_P)
                top_k = gen_params.get("top_k", 40)  # Default to 40 for better quality
                repetition_penalty = gen_params.get("repetition_penalty", 1.1)
            else:
                # Use provided individual parameters or defaults
                max_length = max_length or min(DEFAULT_MAX_LENGTH, 256)  # More conservative limit for low-resource computers
                temperature = temperature or 0.7  # Use same temperature as non-streaming
                top_p = top_p or DEFAULT_TOP_P
                top_k = 40  # Default to 40 for better quality
                repetition_penalty = 1.1

            # For low-resource computers, limit max_length to prevent OOM
            if max_length > 256:
                logger.info(f"Limiting max_length from {max_length} to 256 for better performance on low-resource computers")
                max_length = 256

            # Get the actual device of the model
            model_device = next(self.model.parameters()).device

            # Ensure inputs are on the same device as the model
            for key in inputs:
                if inputs[key].device != model_device:
                    inputs[key] = inputs[key].to(model_device)

            # For low-resource computers, generate one token at a time
            # This is more memory-efficient and prevents OOM errors
            tokens_to_generate_per_step = 1

            # Track generated text for quality control
            generated_text = ""

            # Define stop sequences for proper termination
            stop_sequences = ["</s>", "<|endoftext|>", "<|im_end|>", "<|assistant|>", "<|user|>"]

            # Memory monitoring variables
            last_memory_check = time.time()
            memory_check_interval = 2.0  # Check memory every 2 seconds

            # Get input IDs and attention mask
            input_ids = inputs["input_ids"]
            attention_mask = inputs["attention_mask"]

            with torch.no_grad():
                for step in range(0, max_length, tokens_to_generate_per_step):
                    # Periodically check memory usage and clear cache if needed
                    current_time = time.time()
                    if current_time - last_memory_check > memory_check_interval:
                        last_memory_check = current_time
                        if torch.cuda.is_available():
                            current_mem = torch.cuda.memory_allocated() / (1024 * 1024 * 1024)  # GB
                            total_mem = torch.cuda.get_device_properties(0).total_memory / (1024 * 1024 * 1024)  # GB
                            if current_mem > 0.7 * total_mem:  # If using >70% of GPU memory
                                torch.cuda.empty_cache()
                                logger.info("Cleared CUDA cache during streaming to prevent OOM")

                    # Generate parameters - optimized for low-resource computers
                    generate_params = {
                        "input_ids": input_ids,
                        "attention_mask": attention_mask,
                        "max_new_tokens": tokens_to_generate_per_step,
                        "temperature": temperature,
                        "top_p": top_p,
                        "top_k": top_k,
                        "do_sample": True,
                        "pad_token_id": self.tokenizer.eos_token_id,
                        "repetition_penalty": repetition_penalty,
                        "num_beams": 1  # Explicitly set to 1 to avoid warnings
                    }

                    # Use efficient attention if available, but only on CUDA devices
                    if torch.cuda.is_available() and hasattr(self.model.config, "attn_implementation"):
                        generate_params["attn_implementation"] = "flash_attention_2"

                    try:
                        # Generate tokens
                        outputs = self.model.generate(**generate_params)

                        # Get the new tokens (skip the input tokens)
                        new_tokens = outputs[0][len(input_ids[0]):]

                        # Decode the new token
                        new_text = self.tokenizer.decode(new_tokens, skip_special_tokens=True)

                        # If no new text was generated or it's just whitespace, skip this token but continue generation
                        if not new_text or new_text.isspace():
                            continue

                        # Add to generated text for quality control
                        generated_text += new_text

                        # Check for stop sequences
                        should_stop = False
                        for stop_seq in stop_sequences:
                            if stop_seq in generated_text:
                                # We've reached a stop sequence, stop generation
                                should_stop = True
                                break

                        # Check for repetition (a sign of poor quality)
                        if len(generated_text) > 50:
                            # Check for repeating patterns of 5+ characters
                            last_50_chars = generated_text[-50:]
                            for pattern_len in range(5, 15):
                                if pattern_len < len(last_50_chars) // 2:
                                    pattern = last_50_chars[-pattern_len:]
                                    if pattern in last_50_chars[:-pattern_len]:
                                        # Detected repetition, stop generation
                                        logger.warning("Detected repetition in streaming generation, stopping")
                                        should_stop = True
                                        break

                        # Yield the new text
                        yield new_text

                        # Stop if needed
                        if should_stop:
                            break

                        # Update input_ids for next iteration
                        input_ids = outputs
                        attention_mask = torch.ones_like(input_ids)

                        # Ensure the updated inputs are on the correct device
                        if input_ids.device != model_device:
                            input_ids = input_ids.to(model_device)
                        if attention_mask.device != model_device:
                            attention_mask = attention_mask.to(model_device)

                    except RuntimeError as e:
                        error_msg = str(e)
                        if "CUDA out of memory" in error_msg or "DefaultCPUAllocator: can't allocate memory" in error_msg:
                            # If we run out of memory, clear cache and try to recover
                            if torch.cuda.is_available():
                                torch.cuda.empty_cache()

                            logger.warning(f"Memory error during streaming: {error_msg}. Attempting to recover...")

                            # Wait a moment to let memory clear
                            time.sleep(0.5)

                            try:
                                # Try again with minimal parameters
                                minimal_params = {
                                    "input_ids": input_ids,
                                    "attention_mask": attention_mask,
                                    "max_new_tokens": 1,
                                    "do_sample": False,  # Turn off sampling to save memory
                                    "pad_token_id": self.tokenizer.eos_token_id,
                                    "num_beams": 1
                                }
                                outputs = self.model.generate(**minimal_params)

                                # Continue as before
                                new_tokens = outputs[0][len(input_ids[0]):]
                                new_text = self.tokenizer.decode(new_tokens, skip_special_tokens=True)

                                if new_text and not new_text.isspace():
                                    generated_text += new_text
                                    yield new_text

                                    input_ids = outputs
                                    attention_mask = torch.ones_like(input_ids)
                                else:
                                    # If recovery didn't produce valid text, stop generation
                                    logger.warning("Recovery attempt didn't produce valid text, stopping generation")
                                    break
                            except Exception as recovery_error:
                                # If recovery fails, yield an error message and stop
                                logger.error(f"Recovery failed: {str(recovery_error)}")
                                yield "\nError: Out of memory. Try a shorter response or reduce model parameters."
                                break
                        else:
                            # For other errors, log and yield error message
                            logger.error(f"Error during streaming: {error_msg}")
                            yield f"\nError: {error_msg}"
                            break

        except Exception as e:
            logger.error(f"Streaming generation failed: {str(e)}")
            # Instead of raising an exception, yield the error as text
            # This allows the client to display the error message
            yield f"\nError: {str(e)}"

    async def async_stream_generate(self, inputs: Dict[str, torch.Tensor] = None, gen_params: Dict[str, Any] = None, prompt: str = None, system_prompt: Optional[str] = None, **kwargs):
        """Convert the synchronous stream generator to an async generator with optimizations for low-resource computers.

        This can be called either with:
        1. inputs and gen_params directly (internal use)
        2. prompt, system_prompt and other kwargs (from generate_stream adapter)
        """
        # If called with prompt, prepare inputs and parameters
        if prompt is not None:
            # Get appropriate system instructions
            from .config import system_instructions
            instructions = str(system_instructions.get_instructions(
                self.current_model)) if not system_prompt else str(system_prompt)

            # Format prompt with system instructions
            formatted_prompt = f"""<|system|>{instructions}</|system|>\n<|user|>{prompt}</|user|>\n<|assistant|>"""

            # Get model-specific generation parameters
            from .config import get_model_generation_params
            gen_params = get_model_generation_params(self.current_model)

            # Set optimized defaults for streaming on low-resource computers
            if not kwargs.get("max_length") and not kwargs.get("max_new_tokens"):
                # Use a conservative default max_length for low-resource computers
                gen_params["max_length"] = min(gen_params.get("max_length", DEFAULT_MAX_LENGTH), 256)

            if not kwargs.get("temperature"):
                # Use a slightly lower temperature for more predictable generation
                gen_params["temperature"] = min(gen_params.get("temperature", DEFAULT_TEMPERATURE), 0.7)

            if not kwargs.get("top_k"):
                # Add top_k for better quality
                gen_params["top_k"] = 40

            if not kwargs.get("repetition_penalty"):
                # Add repetition penalty to avoid loops
                gen_params["repetition_penalty"] = 1.1

            # Update with provided kwargs
            for key, value in kwargs.items():
                if key in ["max_length", "temperature", "top_p", "top_k", "repetition_penalty"]:
                    gen_params[key] = value
                elif key == "max_new_tokens":
                    # Handle the max_new_tokens parameter by mapping to max_length
                    gen_params["max_length"] = value

            # Tokenize the prompt
            inputs = self.tokenizer(formatted_prompt, return_tensors="pt")

            # Get the actual device of the model
            model_device = next(self.model.parameters()).device

            # Move inputs to the same device as the model
            for key in inputs:
                inputs[key] = inputs[key].to(model_device)

        # Check if we need to clear memory before generation
        if torch.cuda.is_available():
            # Clear CUDA cache before streaming to ensure maximum available memory
            torch.cuda.empty_cache()
            logger.info("Cleared CUDA cache before streaming generation to ensure optimal performance")

        # Create a token-level streaming generator optimized for low-resource computers
        async def token_level_stream_generator():
            # Define stop sequences for proper termination
            stop_sequences = ["</s>", "<|endoftext|>", "<|im_end|>", "<|assistant|>", "<|user|>"]
            accumulated_text = ""
            token_buffer = ""
            last_yield_time = time.time()
            error_reported = False

            try:
                # Use the optimized _stream_generate method
                for token in self._stream_generate(inputs, gen_params=gen_params):
                    # Check if this is an error message
                    if token.startswith("\nError:"):
                        if not error_reported:
                            yield token
                            error_reported = True
                        break

                    # Skip empty tokens
                    if not token or token.isspace():
                        continue

                    # Add token to buffer
                    token_buffer += token
                    accumulated_text += token

                    # Check for stop sequences
                    should_stop = False
                    for stop_seq in stop_sequences:
                        if stop_seq in accumulated_text:
                            # Truncate at stop sequence
                            accumulated_text = accumulated_text.split(stop_seq)[0]
                            token_buffer = ""
                            should_stop = True
                            break

                    # Determine if we should yield the buffer now
                    current_time = time.time()
                    time_since_last_yield = current_time - last_yield_time

                    # Yield in these cases:
                    # 1. Buffer contains complete word (ends with space)
                    # 2. Buffer contains punctuation
                    # 3. It's been more than 100ms since last yield
                    # 4. Buffer is getting large (>10 chars)
                    should_yield = (
                        token.endswith(" ") or
                        any(p in token for p in ".,!?;:") or
                        time_since_last_yield > 0.1 or
                        len(token_buffer) > 10
                    )

                    if should_yield and token_buffer:
                        # Yield the raw token buffer without any formatting
                        yield token_buffer
                        last_yield_time = current_time
                        token_buffer = ""

                    # Stop if we've reached a stop sequence
                    if should_stop:
                        break

                    # Also stop if we've generated too much text
                    if len(accumulated_text) > gen_params.get("max_length", 256) * 4:
                        logger.warning("Stream generation exceeded maximum length - stopping")
                        break

                    # Yield control to allow other async tasks to run
                    await asyncio.sleep(0)

                # Yield any remaining text in the buffer
                if token_buffer and not error_reported:
                    # Yield the raw token buffer without any formatting
                    yield token_buffer

            except Exception as e:
                logger.error(f"Error in token-level stream generation: {str(e)}")
                # Send error message to client if not already reported
                if not error_reported:
                    yield f"\nError: {str(e)}"

        # Use the token-level streaming generator
        async for token in token_level_stream_generator():
            yield token

    def get_model_info(self) -> Dict[str, Any]:
        """Get information about the currently loaded model"""
        if not self.current_model:
            return {"status": "No model loaded"}

        memory_used = 0
        if self.model:
            memory_used = sum(p.numel() * p.element_size()
                              for p in self.model.parameters())
            num_parameters = sum(p.numel() for p in self.model.parameters())

        model_name = self.model_config.get("name", self.current_model) if isinstance(
            self.model_config, dict) else self.current_model
        max_length = self.model_config.get("max_length", DEFAULT_MAX_LENGTH) if isinstance(
            self.model_config, dict) else DEFAULT_MAX_LENGTH
        ram_required = self.model_config.get("ram", "Unknown") if isinstance(
            self.model_config, dict) else "Unknown"
        vram_required = self.model_config.get("vram", "Unknown") if isinstance(
            self.model_config, dict) else "Unknown"

        # Check environment variables for optimization settings
        enable_quantization = os.environ.get('LOCALLAB_ENABLE_QUANTIZATION', '').lower() not in ('false', '0', 'none', '')
        quantization_type = os.environ.get('LOCALLAB_QUANTIZATION_TYPE', '') if enable_quantization else "None"

        # If quantization is disabled or type is empty, set to "None"
        if not enable_quantization or not quantization_type or quantization_type.lower() in ('none', ''):
            quantization_type = "None"

        model_info = {
            "model_id": self.current_model,
            "model_name": model_name,
            "parameters": f"{num_parameters/1e6:.1f}M",
            "architecture": self.model.__class__.__name__ if self.model else "Unknown",
            "device": self.device,
            "max_length": max_length,
            "ram_required": ram_required,
            "vram_required": vram_required,
            "memory_used": f"{memory_used / (1024 * 1024):.2f} MB",
            "quantization": quantization_type,
            "optimizations": {
                "attention_slicing": os.environ.get('LOCALLAB_ENABLE_ATTENTION_SLICING', '').lower() not in ('false', '0', 'none', ''),
                "flash_attention": os.environ.get('LOCALLAB_ENABLE_FLASH_ATTENTION', '').lower() not in ('false', '0', 'none', ''),
                "better_transformer": os.environ.get('LOCALLAB_ENABLE_BETTERTRANSFORMER', '').lower() not in ('false', '0', 'none', '')
            }
        }

        # Log detailed model information
        logger.info(f"""
{Fore.CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━{Style.RESET_ALL}
{Fore.GREEN}📊 Model Information{Style.RESET_ALL}
{Fore.CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━{Style.RESET_ALL}

• Model: {Fore.YELLOW}{model_name}{Style.RESET_ALL}
• Parameters: {Fore.YELLOW}{model_info['parameters']}{Style.RESET_ALL}
• Architecture: {Fore.YELLOW}{model_info['architecture']}{Style.RESET_ALL}
• Device: {Fore.YELLOW}{self.device}{Style.RESET_ALL}
• Memory Used: {Fore.YELLOW}{model_info['memory_used']}{Style.RESET_ALL}
• Quantization: {Fore.YELLOW}{model_info['quantization']}{Style.RESET_ALL}

{Fore.GREEN}🔧 Optimizations{Style.RESET_ALL}
• Attention Slicing: {Fore.YELLOW}{str(model_info['optimizations']['attention_slicing'])}{Style.RESET_ALL}
• Flash Attention: {Fore.YELLOW}{str(model_info['optimizations']['flash_attention'])}{Style.RESET_ALL}
• Better Transformer: {Fore.YELLOW}{str(model_info['optimizations']['better_transformer'])}{Style.RESET_ALL}
""")

        return model_info

    async def load_custom_model(self, model_name: str, fallback_model: Optional[str] = "qwen-0.5b") -> bool:
        """Load a custom model from Hugging Face Hub with resource checks"""
        try:
            from huggingface_hub import model_info
            info = model_info(model_name)

            estimated_ram = info.siblings[0].size / (1024 * 1024)
            estimated_vram = estimated_ram * 1.5

            temp_config = {
                "name": model_name,
                "ram": estimated_ram,
                "vram": estimated_vram,
                "max_length": 2048,
                "fallback": fallback_model,
                "description": f"Custom model: {info.description}",
                "quantization": "int8",
                "tags": info.tags
            }

            if not check_resource_availability(temp_config["ram"]):
                if fallback_model:
                    logger.warning(
                        f"Insufficient resources for {model_name} "
                        f"(Requires ~{format_model_size(temp_config['ram'])} RAM), "
                        f"falling back to {fallback_model}"
                    )
                    raise HTTPException(
                        status_code=400,
                        detail=f"Insufficient resources. Model requires ~{format_model_size(temp_config['ram'])} RAM"
                    )
                raise HTTPException(
                    status_code=400,
                    detail=f"Insufficient resources. Model requires ~{format_model_size(temp_config['ram'])} RAM"
                )

            if self.model:
                del self.model
                torch.cuda.empty_cache()

            logger.info(f"Loading custom model: {model_name}")

            quant_config = BitsAndBytesConfig(
                load_in_8bit=True,
                llm_int8_threshold=6.0
            )

            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
            self.model = AutoModelForCausalLM.from_pretrained(
                model_name,
                torch_dtype=torch.float16,
                device_map="auto",
                quantization_config=quant_config
            )

            self.model = self._apply_optimizations(self.model)

            self.current_model = f"custom/{model_name}"
            self.model_config = temp_config
            self.last_used = time.time()

            model_size = sum(p.numel() * p.element_size()
                             for p in self.model.parameters())
            logger.info(f"Custom model loaded successfully. Size: {format_model_size(model_size)}")

            return True

        except Exception as e:
            logger.error(f"Failed to load custom model {model_name}: {str(e)}")
            if fallback_model:
                logger.warning(f"Attempting to load fallback model: {fallback_model}")
                return await self.load_model(fallback_model)
            raise HTTPException(
                status_code=500,
                detail=f"Failed to load model: {str(e)}"
            )

    # Add adapter methods to match the interface expected by the routes
    async def generate_text(self, prompt: str, system_prompt: Optional[str] = None, **kwargs) -> str:
        """
        Adapter method that calls the generate method.
        This is used to maintain compatibility with routes that call generate_text.
        """
        # Make sure we're not streaming when generating text
        kwargs["stream"] = False

        # Handle max_new_tokens parameter by mapping to max_length if needed
        if "max_new_tokens" in kwargs and "max_length" not in kwargs:
            kwargs["max_length"] = kwargs.pop("max_new_tokens")

        # Directly await the generate method to return the string result
        return await self.generate(prompt=prompt, system_instructions=system_prompt, **kwargs)

    async def generate_stream(self, prompt: str, system_prompt: Optional[str] = None, **kwargs) -> AsyncGenerator[str, None]:
        """Adapter method for streaming text generation.
        Calls the async_stream_generate method with proper parameters."""
        # Ensure streaming is enabled
        kwargs["stream"] = True

        # Handle max_new_tokens parameter by mapping to max_length
        if "max_new_tokens" in kwargs and "max_length" not in kwargs:
            kwargs["max_length"] = kwargs.pop("max_new_tokens")

        # Call async_stream_generate with the prompt and parameters
        async for token in self.async_stream_generate(prompt=prompt, system_prompt=system_prompt, **kwargs):
            yield token

    def is_model_loaded(self, model_id: str) -> bool:
        """Check if a specific model is loaded.

        Args:
            model_id: The ID of the model to check

        Returns:
            True if the model is loaded, False otherwise
        """
        return (self.model is not None) and (self.current_model == model_id)

    def unload_model(self) -> None:
        """Unload the current model to free memory resources.

        This method removes the current model from memory and clears
        the tokenizer and model references.
        """
        if self.model is not None:
            # Log which model is being unloaded
            model_id = self.current_model

            # Clear model and tokenizer
            self.model = None
            self.tokenizer = None
            self.current_model = None

            # Clean up memory
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            gc.collect()

            # Log model unloading
            log_model_unloaded(model_id)

            logger.info(f"Model {model_id} unloaded successfully")