import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer, BitsAndBytesConfig
from threading import Thread
import time
import logging
import gc
from pathlib import Path
import re
from huggingface_hub import HfApi, list_models
import os
import queue
import threading
from collections import deque

# Set PyTorch memory management environment variables
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('gradio-chat-ui.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# Log memory management settings
logger.info(f"PyTorch CUDA allocation config: {os.environ.get('PYTORCH_CUDA_ALLOC_CONF')}")
logger.info(f"CUDA device count: {torch.cuda.device_count() if torch.cuda.is_available() else 'N/A'}")

# Model parameters
MODEL_NAME = "No Model Loaded"
MAX_LENGTH = 16384
DEFAULT_TEMPERATURE = 0.15
DEFAULT_TOP_P = 0.93
DEFAULT_TOP_K = 50
DEFAULT_REP_PENALTY = 1.15

# Base location for local models
LOCAL_MODELS_BASE = "/home/llm-models/"

# Global variables
model = None
tokenizer = None
hf_api = HfApi()

# Generation metadata storage with automatic cleanup
generation_metadata = deque(maxlen=100)  # Fixed size deque to prevent unlimited growth

class RAMSavingIteratorStreamer:
    """
    Custom streamer that saves VRAM by moving tokens to CPU and provides iteration interface for Gradio.
    Combines the benefits of TextStreamer (RAM saving) with TextIteratorStreamer (iteration).
    """
    def __init__(self, tokenizer, skip_special_tokens=True, skip_prompt=True, timeout=None):
        self.tokenizer = tokenizer
        self.skip_special_tokens = skip_special_tokens
        self.skip_prompt = skip_prompt
        self.timeout = timeout
        
        # Token and text storage (CPU-based)
        self.generated_tokens = []
        self.generated_text = ""
        self.token_cache = ""
        
        # Queue for streaming interface
        self.text_queue = queue.Queue()
        self.stop_signal = threading.Event()
        
        # Track prompt tokens to skip them
        self.prompt_length = 0
        self.tokens_processed = 0
        
        # Decoding state
        self.print_len = 0
        
    def put(self, value):
        """
        Receive new token(s) and process them for streaming.
        This method is called by the model during generation.
        """
        try:
            # Handle different input types
            if isinstance(value, torch.Tensor):
                if value.dim() > 1:
                    value = value[0]  # Remove batch dimension if present
                token_ids = value.tolist()
                
                # Store CPU version to save VRAM
                self.generated_tokens.append(value.detach().cpu())
            else:
                token_ids = value if isinstance(value, list) else [value]
                self.generated_tokens.append(torch.tensor(token_ids, dtype=torch.long))
            
            # Track tokens processed
            if isinstance(token_ids, list):
                self.tokens_processed += len(token_ids)
            else:
                self.tokens_processed += 1
            
            # Skip prompt tokens if requested
            if self.skip_prompt and self.tokens_processed <= self.prompt_length:
                return
            
            # Decode incrementally for real-time streaming
            try:
                # Get all generated tokens so far
                if self.generated_tokens:
                    all_tokens = []
                    for tokens in self.generated_tokens:
                        if isinstance(tokens, torch.Tensor):
                            if tokens.dim() == 0:
                                all_tokens.append(tokens.item())
                            else:
                                all_tokens.extend(tokens.tolist())
                        elif isinstance(tokens, list):
                            all_tokens.extend(tokens)
                        else:
                            all_tokens.append(tokens)
                    
                    # Decode the full sequence
                    full_text = self.tokenizer.decode(
                        all_tokens, 
                        skip_special_tokens=self.skip_special_tokens
                    )
                    
                    # Get new text since last update
                    if len(full_text) > self.print_len:
                        new_text = full_text[self.print_len:]
                        self.print_len = len(full_text)
                        self.generated_text = full_text
                        
                        # Put new text in queue for iteration
                        if new_text:
                            self.text_queue.put(new_text)
                            
            except Exception as decode_error:
                logger.warning(f"Decoding error in streamer: {decode_error}")
                
        except Exception as e:
            logger.error(f"Error in RAMSavingIteratorStreamer.put: {e}")
            
    def end(self):
        """Signal end of generation."""
        self.text_queue.put(None)  # Sentinel value
        
    def __iter__(self):
        """Make this streamer iterable for Gradio compatibility."""
        return self
        
    def __next__(self):
        """Get next chunk of text for streaming."""
        try:
            value = self.text_queue.get(timeout=self.timeout)
            if value is None:  # End signal
                raise StopIteration
            return value
        except queue.Empty:
            raise StopIteration
            
    def set_prompt_length(self, prompt_length):
        """Set the length of prompt tokens to skip."""
        self.prompt_length = prompt_length
        
    def get_generated_text(self):
        """Get the complete generated text."""
        return self.generated_text
        
    def get_generated_tokens(self):
        """Get all generated tokens as a single tensor."""
        if not self.generated_tokens:
            return torch.tensor([])
            
        # Combine all tokens
        all_tokens = []
        for tokens in self.generated_tokens:
            if isinstance(tokens, torch.Tensor):
                if tokens.dim() == 0:
                    all_tokens.append(tokens.item())
                else:
                    all_tokens.extend(tokens.tolist())
            elif isinstance(tokens, list):
                all_tokens.extend(tokens)
            else:
                all_tokens.append(tokens)
                
        return torch.tensor(all_tokens, dtype=torch.long)
        
    def cleanup(self):
        """Clean up resources."""
        self.generated_tokens.clear()
        self.generated_text = ""
        self.token_cache = ""
        
        # Clear queue
        while not self.text_queue.empty():
            try:
                self.text_queue.get_nowait()
            except queue.Empty:
                break
                
        self.stop_signal.set()

def scan_local_models(base_path=LOCAL_MODELS_BASE):
    """Scan for valid models in the local models directory"""
    try:
        base_path = Path(base_path)
        if not base_path.exists():
            logger.warning(f"Base path does not exist: {base_path}")
            return []
        
        valid_models = []
        
        # Scan subdirectories (depth 1 only)
        for item in base_path.iterdir():
            if item.is_dir():
                # Check if directory contains required model files
                config_file = item / "config.json"
                
                # Look for model weight files (safetensors or bin)
                safetensors_files = list(item.glob("*.safetensors"))
                bin_files = list(item.glob("*.bin"))
                
                # Check if it's a valid model directory
                if config_file.exists() and (safetensors_files or bin_files):
                    valid_models.append(str(item))
                    logger.info(f"Found valid model: {item}")
        
        # Sort models for consistent ordering
        valid_models.sort()
        logger.info(f"Found {len(valid_models)} valid models in {base_path}")
        
        return valid_models
        
    except Exception as e:
        logger.error(f"Error scanning local models: {e}")
        return []

def update_local_models_dropdown(base_path):
    """Update the local models dropdown based on base path"""
    if not base_path or not base_path.strip():
        return gr.Dropdown(choices=[], value=None, interactive=True)
    
    models = scan_local_models(base_path)
    model_choices = [Path(model).name for model in models]  # Show just the model name
    model_paths = models  # Keep full paths for internal use
    
    # Create a mapping for display name to full path
    if model_choices:
        return gr.Dropdown(
            choices=list(zip(model_choices, model_paths)), 
            value=model_paths[0] if model_paths else None,
            label="📋 Available Local Models",
            interactive=True,
            allow_custom_value=False,  # Don't allow custom for local models
            filterable=True
        )
    else:
        return gr.Dropdown(
            choices=[], 
            value=None,
            label="📋 Available Local Models (None found)",
            interactive=True,
            allow_custom_value=False,
            filterable=True
        )

def search_hf_models(query, limit=20):
    """Enhanced search for models on Hugging Face Hub with better coverage"""
    if not query or len(query.strip()) < 2:
        return []
    
    try:
        query = query.strip()
        model_choices = []
        
        # Strategy 1: Direct model ID search (if query looks like a model ID)
        if '/' in query:
            try:
                # Try to get the specific model
                model_info = hf_api.model_info(query)
                if model_info and hasattr(model_info, 'id'):
                    model_choices.append(model_info.id)
                    logger.info(f"Found direct model: {model_info.id}")
            except Exception as direct_error:
                logger.debug(f"Direct model search failed: {direct_error}")
        
        # Strategy 2: Search with different parameters
        search_strategies = [
            # Exact search
            {"search": query, "sort": "downloads", "direction": -1, "limit": limit//2},
            # Author search (if query contains /)
            {"author": query.split('/')[0] if '/' in query else query, "sort": "downloads", "direction": -1, "limit": limit//4} if '/' in query else None,
            # Broader search
            {"search": query, "sort": "trending", "direction": -1, "limit": limit//4},
        ]
        
        for strategy in search_strategies:
            if strategy is None:
                continue
                
            try:
                models = list_models(
                    task="text-generation",
                    **strategy
                )
                
                for model in models:
                    if model.id not in model_choices:
                        model_choices.append(model.id)
                        
            except Exception as strategy_error:
                logger.debug(f"Search strategy failed: {strategy_error}")
        
        # Remove duplicates while preserving order
        seen = set()
        unique_choices = []
        for choice in model_choices:
            if choice not in seen:
                seen.add(choice)
                unique_choices.append(choice)
        
        # Limit results
        final_choices = unique_choices[:limit]
        logger.info(f"HF search for '{query}' returned {len(final_choices)} models")
        
        return final_choices
        
    except Exception as e:
        logger.error(f"Error searching models: {str(e)}")
        return []

def update_model_dropdown(query):
    """Update dropdown with enhanced search results"""
    if not query or len(query.strip()) < 2:
        return gr.Dropdown(choices=[], value=None, interactive=True)
    
    choices = search_hf_models(query, limit=20)
    return gr.Dropdown(
        choices=choices, 
        value=choices[0] if choices else None,
        interactive=True,
        allow_custom_value=True,  # Allow manual typing
        filterable=True
    )

def load_model_with_progress(model_source, hf_model, local_path, local_model_selection, quantization, memory_optimization):
    """Load model with progress tracking and memory optimization"""
    global model, tokenizer, MODEL_NAME
    
    # Determine model path based on source
    if model_source == "Hugging Face Model":
        if not hf_model:
            return "❌ Error: Please select a model from the dropdown"
        model_path = hf_model
    else:
        # Use selected local model if available, otherwise use manual path
        if local_model_selection:
            model_path = local_model_selection
        else:
            model_path = local_path
            if not Path(model_path).exists():
                logger.error(f"Local path does not exist: {model_path}")
                return f"❌ Error: Local path does not exist: {model_path}"
    
    MODEL_NAME = model_path.split("/")[-1] if "/" in model_path else model_path
    logger.info(f"Loading model from {model_path} with memory optimization: {memory_optimization}")
    
    try:
        # Yield progress updates
        yield "🔄 Initializing model loading..."
        
        # Setup memory configuration (GPU-only, generous allocation)
        if torch.cuda.is_available():
            device_properties = torch.cuda.get_device_properties(0)
            total_memory_gb = device_properties.total_memory / (1024**3)
            
            # Set max memory to 11GB as requested (GPU-bound)
            max_memory_val = 11.5  # Fixed 11GB allocation
            max_memory = f"{max_memory_val}GB"
            logger.info(f"Setting max GPU memory to {max_memory} (Total available: {total_memory_gb:.2f}GB)")
        else:
            max_memory = "11GB"
            logger.info("CUDA not available. Using CPU fallback.")
        
        yield "🔄 Configuring quantization settings..."
        
        # Configure quantization (removed CPU offloading)
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=quantization == "4bit", 
            load_in_8bit=quantization == "8bit",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_quant_type="nf4",
        )
        
        yield "🔄 Loading tokenizer..."
        
        # Load tokenizer
        if model_source == "Local Path":
            tokenizer = AutoTokenizer.from_pretrained(
                model_path, 
                trust_remote_code=True, 
                local_files_only=True
            )
        else:
            tokenizer = AutoTokenizer.from_pretrained(
                model_path, 
                trust_remote_code=True
            )
        
        yield "🔄 Cleaning memory cache..."
        
        # Clean memory
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        
        # Determine torch dtype
        if quantization in ["4bit", "8bit"]:
            torch_dtype = torch.bfloat16
        elif quantization == "f16":
            torch_dtype = torch.float16
        else:  # bf16
            torch_dtype = torch.bfloat16
        
        yield "🔄 Loading model weights (this may take a while)..."
        
        # Simple GPU-only model loading parameters
        model_kwargs = {
            "device_map": "auto",
            "max_memory": {0: max_memory} if torch.cuda.is_available() else None,
            "torch_dtype": torch_dtype,
            "quantization_config": bnb_config if quantization in ["4bit", "8bit"] else None,
            "trust_remote_code": True,
        }
        
        # Memory optimization specific settings (GPU-only)
        if memory_optimization:
            model_kwargs.update({
                "attn_implementation": "flash_attention_2" if torch.cuda.is_available() else "sdpa",
                "use_cache": False,  # Disable cache by default for memory optimization
            })
        else:
            model_kwargs.update({
                "attn_implementation": "flash_attention_2" if torch.cuda.is_available() else "sdpa",
                #"use_cache": True,  # Enable cache for performance
            })
        
        # Add local files only for local models
        if model_source == "Local Path":
            model_kwargs["local_files_only"] = True
        
        # Load model
        model = AutoModelForCausalLM.from_pretrained(model_path, **model_kwargs)
        
        # Post-loading memory optimization
        if memory_optimization:
            yield "🔄 Applying memory optimizations..."
            
            # Additional memory cleanup after loading
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                torch.cuda.synchronize()
        
        logger.info("Model loaded successfully with memory optimization")
        yield "✅ Model loaded successfully with memory optimization!" if memory_optimization else "✅ Model loaded successfully!"
        
    except Exception as e:
        logger.error(f"Error loading model: {str(e)}", exc_info=True)
        yield f"❌ Error loading model: {str(e)}"

def unload_model():
    """Unload the model and free memory with aggressive cleanup"""
    global model, tokenizer, MODEL_NAME
    
    if model is None:
        return "No model loaded"
    
    try:
        logger.info("Unloading model with aggressive memory cleanup...")
        
        # Step 1: Move model to CPU first (if it was on GPU)
        if torch.cuda.is_available() and hasattr(model, 'device'):
            try:
                model.cpu()
                logger.info("Model moved to CPU")
            except Exception as cpu_error:
                logger.warning(f"Could not move model to CPU: {cpu_error}")
        
        # Step 2: Clear model cache if available
        if hasattr(model, 'clear_cache'):
            model.clear_cache()
        
        # Step 3: Delete model and tokenizer references
        del model
        del tokenizer
        model = None
        tokenizer = None
        
        # Step 4: Reset model name
        MODEL_NAME = "No Model Loaded"
        
        # Step 5: Clear metadata deque
        generation_metadata.clear()
        
        # Step 6: Aggressive garbage collection (multiple rounds)
        for i in range(5):  # More aggressive - 5 rounds
            gc.collect()
            time.sleep(0.1)  # Small delay between rounds
        
        # Step 7: Aggressive CUDA cleanup
        if torch.cuda.is_available():
            logger.info("Performing aggressive CUDA cleanup...")
            
            # Multiple rounds of cache clearing
            for i in range(5):
                torch.cuda.empty_cache()
                torch.cuda.synchronize()
                
                # Additional PyTorch CUDA cleanup
                if hasattr(torch.cuda, 'ipc_collect'):
                    torch.cuda.ipc_collect()
                
                # Reset memory stats
                if hasattr(torch.cuda, 'reset_peak_memory_stats'):
                    torch.cuda.reset_peak_memory_stats()
                if hasattr(torch.cuda, 'reset_accumulated_memory_stats'):
                    torch.cuda.reset_accumulated_memory_stats()
                
                time.sleep(0.1)
        
        # Step 8: Force PyTorch to release all unused memory
        if torch.cuda.is_available():
            try:
                # Try to trigger the memory pool cleanup
                torch.cuda.empty_cache()
                
                # Force a small allocation and deallocation to trigger cleanup
                dummy_tensor = torch.zeros(1, device='cuda')
                del dummy_tensor
                torch.cuda.empty_cache()
                
                logger.info("Forced memory pool cleanup")
            except Exception as cleanup_error:
                logger.warning(f"Advanced cleanup failed: {cleanup_error}")
        
        # Step 9: Final garbage collection
        gc.collect()
        
        logger.info("Model unloaded successfully with aggressive cleanup")
        return "✅ Model unloaded with aggressive memory cleanup"
        
    except Exception as e:
        logger.error(f"Error unloading model: {str(e)}", exc_info=True)
        # Emergency cleanup even if unload fails
        model = None
        tokenizer = None
        MODEL_NAME = "No Model Loaded"
        generation_metadata.clear()
        
        # Emergency memory cleanup
        for _ in range(3):
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
        
        return f"❌ Error unloading model: {str(e)} (Emergency cleanup performed)"

def cleanup_memory():
    """Enhanced memory cleanup function with PyTorch optimizations"""
    try:
        # Clear Python garbage
        gc.collect()
        
        # Clear CUDA cache if available
        if torch.cuda.is_available():
            # Multiple aggressive cleanup rounds
            for i in range(3):
                torch.cuda.empty_cache()
                torch.cuda.synchronize()
                if hasattr(torch.cuda, 'ipc_collect'):
                    torch.cuda.ipc_collect()
                
                # PyTorch specific memory management
                if hasattr(torch.cuda, 'reset_peak_memory_stats'):
                    torch.cuda.reset_peak_memory_stats()
                if hasattr(torch.cuda, 'reset_accumulated_memory_stats'):
                    torch.cuda.reset_accumulated_memory_stats()
                
                # Brief pause between cleanup rounds
                time.sleep(0.1)
        
        # Clear metadata deque
        generation_metadata.clear()
        
        # Force garbage collection again
        gc.collect()
        
        logger.info("Enhanced memory cleanup completed")
        return "🧹 Enhanced memory cleanup completed"
    except Exception as e:
        logger.error(f"Memory cleanup error: {e}")
        return f"Memory cleanup error: {e}"

def nuclear_memory_cleanup():
    """Nuclear option: Complete VRAM reset (use if normal unload doesn't work)"""
    global model, tokenizer, MODEL_NAME
    
    try:
        logger.info("Performing nuclear memory cleanup...")
        
        # Force unload everything
        model = None
        tokenizer = None
        MODEL_NAME = "No Model Loaded"
        generation_metadata.clear()
        
        # Import PyTorch again to reset some internal states
        import torch
        
        # Multiple aggressive cleanup rounds
        for round_num in range(10):  # Very aggressive - 10 rounds
            gc.collect()
            
            if torch.cuda.is_available():
                # Multiple types of CUDA cleanup
                torch.cuda.empty_cache()
                torch.cuda.synchronize()
                
                # Try to reset CUDA context
                try:
                    if hasattr(torch.cuda, 'ipc_collect'):
                        torch.cuda.ipc_collect()
                    if hasattr(torch.cuda, 'memory_summary'):
                        logger.info(f"Round {round_num + 1}: {torch.cuda.memory_summary()}")
                except Exception:
                    pass
                
                # Reset memory stats
                try:
                    if hasattr(torch.cuda, 'reset_peak_memory_stats'):
                        torch.cuda.reset_peak_memory_stats()
                    if hasattr(torch.cuda, 'reset_accumulated_memory_stats'):
                        torch.cuda.reset_accumulated_memory_stats()
                except Exception:
                    pass
            
            time.sleep(0.1)
        
        # Final attempt: allocate and free a small tensor to trigger cleanup
        if torch.cuda.is_available():
            try:
                for _ in range(5):
                    dummy = torch.zeros(1024, 1024, device='cuda')  # 4MB tensor
                    del dummy
                    torch.cuda.empty_cache()
                    torch.cuda.synchronize()
            except Exception as nuclear_error:
                logger.warning(f"Nuclear tensor cleanup failed: {nuclear_error}")
        
        logger.info("Nuclear memory cleanup completed")
        return "☢️ Nuclear memory cleanup completed! VRAM should be minimal now."
        
    except Exception as e:
        logger.error(f"Nuclear cleanup error: {e}")
        return f"☢️ Nuclear cleanup error: {e}"

def get_memory_stats():
    """Get comprehensive VRAM usage information"""
    if not torch.cuda.is_available():
        return """
        <div style="text-align: center; padding: 15px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 10px; color: white;">
            <h3 style="margin: 0; font-size: 16px;">💻 CPU Mode</h3>
            <p style="margin: 5px 0; opacity: 0.9;">GPU not available</p>
        </div>
        """
    
    try:
        torch.cuda.synchronize()
        total = torch.cuda.get_device_properties(0).total_memory / (1024**3)
        allocated = torch.cuda.memory_allocated(0) / (1024**3)
        reserved = torch.cuda.memory_reserved(0) / (1024**3)
        free = total - reserved
        usage_percent = (reserved/total)*100
        
        # Get peak memory if available
        peak_allocated = 0
        if hasattr(torch.cuda, 'max_memory_allocated'):
            peak_allocated = torch.cuda.max_memory_allocated(0) / (1024**3)
        
        # Dynamic color based on usage
        if usage_percent < 50:
            color = "#10b981"  # Green
        elif usage_percent < 80:
            color = "#f59e0b"  # Orange
        else:
            color = "#ef4444"  # Red
        
        return f"""
        <div style="text-align: center; padding: 15px; background: linear-gradient(135deg, {color}22 0%, {color}44 100%); border: 2px solid {color}; border-radius: 10px;">
            <h3 style="margin: 0; font-size: 16px; color: {color};">🎮 VRAM Usage</h3>
            <div style="margin: 10px 0;">
                <div style="background: #f3f4f6; border-radius: 8px; height: 8px; overflow: hidden;">
                    <div style="width: {usage_percent}%; height: 100%; background: {color}; transition: width 0.3s ease;"></div>
                </div>
            </div>
            <p style="margin: 5px 0; font-weight: 600;">Total: {total:.2f} GB</p>
            <p style="margin: 5px 0;">Allocated: {allocated:.2f} GB ({usage_percent:.1f}%)</p>
            <p style="margin: 5px 0;">Reserved: {reserved:.2f} GB</p>
            <p style="margin: 5px 0;">Free: {free:.2f} GB</p>
            <p style="margin: 5px 0; font-size: 12px; opacity: 0.8;">Peak: {peak_allocated:.2f} GB</p>
            <p style="margin: 5px 0; font-size: 10px; opacity: 0.6;">RAM-Saving Streamer Active</p>
        </div>
        """
    except Exception as e:
        logger.error(f"Error getting memory stats: {str(e)}")
        return f"""
        <div style="text-align: center; padding: 15px; background: #fee2e2; border: 2px solid #ef4444; border-radius: 10px;">
            <h3 style="margin: 0; color: #ef4444;">❌ Error</h3>
            <p style="margin: 5px 0;">{str(e)}</p>
        </div>
        """

def process_latex_content(text):
    """Enhanced LaTeX processing for streaming without UI glitches"""
    # Don't process LaTeX here - let Gradio handle it natively
    # Just return the text as-is for now
    return text

def process_think_tags(text):
    """Process thinking tags with progressive streaming support"""
    # Check if we're in the middle of generating a think section
    if '<think>' in text and '</think>' not in text:
        # We're currently generating inside a think section
        parts = text.split('<think>')
        if len(parts) == 2:
            before_think = parts[0]
            thinking_content = parts[1]
            
            # Create a progressive thinking display
            formatted_thinking = f"""
<div style="background: linear-gradient(135deg, #e0e7ff 0%, #c7d2fe 100%); border-left: 4px solid #6366f1; padding: 12px; margin: 8px 0; border-radius: 8px;">
    <div style="display: flex; align-items: center; margin-bottom: 8px;">
        <span style="font-size: 16px; margin-right: 8px;">🤔</span>
        <strong style="color: #4338ca;">Thinking...</strong>
    </div>
    <div style="color: #475569; font-style: italic;">{thinking_content}</div>
</div>

"""
            return before_think + formatted_thinking
    
    # Handle completed think sections
    think_pattern = re.compile(r'<think>(.*?)</think>', re.DOTALL)
    
    def replace_think(match):
        think_content = match.group(1).strip()
        return f"""
<div style="background: linear-gradient(135deg, #e0e7ff 0%, #c7d2fe 100%); border-left: 4px solid #6366f1; padding: 12px; margin: 8px 0; border-radius: 8px;">
    <div style="display: flex; align-items: center; margin-bottom: 8px;">
        <span style="font-size: 16px; margin-right: 8px;">🤔</span>
        <strong style="color: #4338ca;">Thinking...</strong>
    </div>
    <div style="color: #475569; font-style: italic;">{think_content}</div>
</div>

"""
    
    # Replace completed <think> tags with formatted version
    processed_text = think_pattern.sub(replace_think, text)
    
    return processed_text

def calculate_generation_metrics(start_time, total_tokens):
    """Calculate generation metrics"""
    end_time = time.time()
    generation_time = end_time - start_time
    tokens_per_second = total_tokens / generation_time if generation_time > 0 else 0
    
    return {
        "generation_time": generation_time,
        "total_tokens": total_tokens,
        "tokens_per_second": tokens_per_second,
        "model_name": MODEL_NAME
    }

def format_metadata_tooltip(metadata):
    """Format metadata for tooltip display"""
    return f"""Model: {metadata['model_name']}
Tokens: {metadata['total_tokens']}
Speed: {metadata['tokens_per_second']:.2f} tok/s
Time: {metadata['generation_time']:.2f}s"""

def add_metadata_to_response(response_text, metadata):
    """Add metadata icon with tooltip to the response"""
    tooltip_content = format_metadata_tooltip(metadata)
    
    # Create a metadata icon with tooltip using HTML
    metadata_html = f"""
<div style="position: relative; display: inline-block; margin-left: 8px;">
    <span class="metadata-icon" style="cursor: help; opacity: 0.6; font-size: 14px;" title="{tooltip_content}">ℹ️</span>
</div>
"""
    
    # Add metadata icon at the end of the response
    return response_text + "\n\n" + metadata_html

def chat_with_model(message, history, system_prompt, temp, top_p_val, top_k_val, rep_penalty_val, memory_opt):
    """
    Enhanced chat function with RAM-saving streamer and improved memory management.
    Uses direct generation approach for better memory control and VRAM efficiency.
    """
    global model, tokenizer, generation_metadata
    
    # Check if model is loaded
    if model is None or tokenizer is None:
        return "❌ Model not loaded. Please load the model first."
    
    # Initialize variables for cleanup
    input_ids = None
    streamer = None
    
    try:
        # Record start time for metrics
        start_time = time.time()
        token_count = 0
        
        # Format conversation for model
        messages = [{"role": "system", "content": system_prompt}]
        
        # Add chat history - HANDLE BOTH FORMATS (tuples from original and dicts from new)
        for h in history:
            if isinstance(h, dict):
                # New dict format
                if h.get("role") == "user":
                    messages.append({"role": "user", "content": h["content"]})
                elif h.get("role") == "assistant":
                    messages.append({"role": "assistant", "content": h["content"]})
            else:
                # Original tuple format (user_msg, bot_msg)
                if len(h) >= 2:
                    messages.append({"role": "user", "content": h[0]})
                    if h[1] is not None:
                        messages.append({"role": "assistant", "content": h[1]})
            
        # Add the current message
        messages.append({"role": "user", "content": message})
        
        # Wrap generation in torch.no_grad() to prevent gradient accumulation
        with torch.no_grad():
            # Create model input with memory-efficient approach
            input_ids = tokenizer.apply_chat_template(
                messages,
                tokenize=True,
                add_generation_prompt=True,
                return_tensors="pt"
            )
            
            # Handle edge case
            if input_ids.ndim == 1:
                input_ids = input_ids.unsqueeze(0)
            
            # Move to device
            input_ids = input_ids.to(model.device)
            
            # Setup RAM-saving streamer
            streamer = RAMSavingIteratorStreamer(
                tokenizer,
                skip_special_tokens=True,
                skip_prompt=True,
                timeout=1.0
            )
            
            # Set prompt length for the streamer
            streamer.set_prompt_length(input_ids.shape[1])
            
            # Pre-generation memory cleanup (only if memory optimization is on)
            if memory_opt:
                gc.collect()
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            
            # Conditional generation parameters based on memory optimization
            gen_kwargs = {
                "input_ids": input_ids,
                "max_new_tokens": MAX_LENGTH,
                "temperature": temp,
                "top_p": top_p_val,
                "top_k": top_k_val,
                "repetition_penalty": rep_penalty_val,
                "do_sample": temp > 0,
                "streamer": streamer,
                "use_cache": not memory_opt,  # Disable cache only if memory optimization is on
            }
            
            # Generate in a thread for real-time streaming
            thread = Thread(
                target=model.generate, 
                kwargs=gen_kwargs,
                daemon=True
            )
            thread.start()
            
            # Stream the response with conditional memory management
            partial_text = ""
            try:
                for new_text in streamer:
                    partial_text += new_text
                    token_count += 1
                    
                    # Process the text to handle think tags while preserving LaTeX
                    processed_text = process_think_tags(partial_text)
                    
                    yield processed_text
                    
                    # Conditional cleanup based on memory optimization setting (less frequent)
                    if memory_opt and token_count % 150 == 0:  # Reduced frequency for performance
                        gc.collect()  # Only light cleanup if memory optimization is on
                        
            except StopIteration:
                # Normal end of generation
                pass
            except Exception as stream_error:
                logger.error(f"Streaming error: {stream_error}")
                yield f"❌ Streaming error: {stream_error}"
                return
                
            finally:
                # Add metadata to final response
                try:
                    metrics = calculate_generation_metrics(start_time, token_count)
                    partial_text = add_metadata_to_response(partial_text, metrics)
                except Exception as e:
                    logger.warning(f"Couldn't add metadata: {str(e)}")
                
                yield partial_text
                
                # Ensure thread completion
                if thread.is_alive():
                    thread.join(timeout=5.0)
                    if thread.is_alive():
                        logger.warning("Generation thread did not complete in time")
                
                # Calculate generation metrics
                try:
                    metrics = calculate_generation_metrics(start_time, token_count)
                    
                    # Store metadata (using deque with max size to prevent memory leaks)
                    generation_metadata.append(metrics)
                    
                    # Log the metrics
                    logger.info(f"Generation metrics - Tokens: {metrics['total_tokens']}, Speed: {metrics['tokens_per_second']:.2f} tok/s, Time: {metrics['generation_time']:.2f}s")
                except Exception as metrics_error:
                    logger.warning(f"Error calculating metrics: {metrics_error}")
        
        # Final cleanup
        try:
            # Clean up streamer
            if streamer:
                streamer.cleanup()
                del streamer
                streamer = None
            
            # Clean up input tensors
            if input_ids is not None:
                del input_ids
                input_ids = None
            
            # Conditional cleanup based on memory optimization setting
            if memory_opt:
                # Aggressive cleanup only if memory optimization is enabled
                if torch.cuda.is_available():
                    for _ in range(2):  # Reduced rounds for performance
                        torch.cuda.empty_cache()
                        torch.cuda.synchronize()
                # Force garbage collection
                for _ in range(2):
                    gc.collect()
            else:
                # Light cleanup for performance mode
                gc.collect()
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            
            logger.info(f"Generation completed, {token_count} tokens, memory_opt: {memory_opt}, VRAM saved with RAM-saving streamer")
            
        except Exception as cleanup_error:
            logger.warning(f"Final cleanup warning: {cleanup_error}")
        
    except Exception as e:
        logger.error(f"Error in chat_with_model: {str(e)}", exc_info=True)
        
        # Emergency cleanup
        try:
            if streamer:
                streamer.cleanup()
                del streamer
            if input_ids is not None:
                del input_ids
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
        except Exception as emergency_cleanup_error:
            logger.error(f"Emergency cleanup failed: {emergency_cleanup_error}")
        
        yield f"❌ Error: {str(e)}"

def update_model_name():
    """Update the displayed model name"""
    return f"🔮 AI Chat Assistant ({MODEL_NAME})"

def add_page_refresh_warning():
    """Add JavaScript to warn about page refresh when model is loaded"""
    return """
    <script>
    window.addEventListener('beforeunload', function (e) {
        // Check if model is loaded by looking for specific text in the page
        const statusElements = document.querySelectorAll('input[type="text"], textarea');
        let modelLoaded = false;
        
        statusElements.forEach(element => {
            if (element.value && element.value.includes('Model loaded successfully')) {
                modelLoaded = true;
            }
        });
        
        if (modelLoaded) {
            e.preventDefault();
            e.returnValue = 'A model is currently loaded. Are you sure you want to leave?';
            return 'A model is currently loaded. Are you sure you want to leave?';
        }
    });
    </script>
    """

# Custom CSS for elegant styling with fixed dropdown behavior
custom_css = """
/* Main container styling */
.gradio-container {
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif !important;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    min-height: 100vh;
}

/* Header styling */
.header-text {
    background: rgba(255, 255, 255, 0.95);
    backdrop-filter: blur(10px);
    border-radius: 15px;
    padding: 20px;
    margin: 20px 0;
    text-align: center;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.1);
    border: 1px solid rgba(255, 255, 255, 0.2);
}

/* Chat interface styling */
.chat-container {
    background: rgba(255, 255, 255, 0.95) !important;
    border-radius: 20px !important;
    box-shadow: 0 20px 40px rgba(0, 0, 0, 0.1) !important;
    border: 1px solid rgba(255, 255, 255, 0.2) !important;
    backdrop-filter: blur(10px) !important;
}

/* Control panel styling */
.control-panel {
    background: rgba(255, 255, 255, 0.9) !important;
    border-radius: 15px !important;
    padding: 20px !important;
    box-shadow: 0 10px 30px rgba(0, 0, 0, 0.1) !important;
    border: 1px solid rgba(255, 255, 255, 0.3) !important;
    backdrop-filter: blur(10px) !important;
    overflow: visible !important; /* Allow dropdowns to overflow */
}

/* Button styling */
.btn-primary {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    border-radius: 10px !important;
    color: white !important;
    font-weight: 600 !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 15px rgba(102, 126, 234, 0.4) !important;
}

.btn-primary:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 8px 25px rgba(102, 126, 234, 0.6) !important;
}

.btn-secondary {
    background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%) !important;
    border: none !important;
    border-radius: 10px !important;
    color: white !important;
    font-weight: 600 !important;
    transition: all 0.3s ease !important;
}

/* Input field styling */
.input-field {
    border-radius: 10px !important;
    border: 2px solid rgba(102, 126, 234, 0.2) !important;
    transition: all 0.3s ease !important;
}

.input-field:focus {
    border-color: #667eea !important;
    box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
}

/* Dropdown fixes */
.dropdown-container {
    position: relative !important;
    z-index: 1000 !important;
    overflow: visible !important;
}

/* Fix dropdown menu positioning and styling */
.dropdown select, 
.dropdown-menu,
.svelte-select,
.svelte-select-list {
    position: relative !important;
    z-index: 1001 !important;
    background: white !important;
    border: 2px solid rgba(102, 126, 234, 0.2) !important;
    border-radius: 10px !important;
    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.15) !important;
    max-height: 200px !important;
    overflow-y: auto !important;
}

/* Fix dropdown option styling */
.dropdown option,
.svelte-select-option {
    padding: 8px 12px !important;
    background: white !important;
    color: #333 !important;
    border: none !important;
}

.dropdown option:hover,
.svelte-select-option:hover {
    background: #f0f0f0 !important;
    color: #667eea !important;
}

/* Ensure dropdown arrow is clickable */
.dropdown::after,
.dropdown-arrow {
    pointer-events: none !important;
    z-index: 1002 !important;
}

/* Fix any overflow issues in parent containers */
.gradio-group,
.gradio-column {
    overflow: visible !important;
}

/* Accordion styling */
.accordion {
    border-radius: 10px !important;
    border: 1px solid rgba(102, 126, 234, 0.2) !important;
    overflow: visible !important; /* Allow dropdowns to overflow accordion */
}

/* Status indicators */
.status-success {
    color: #10b981 !important;
    font-weight: 600 !important;
}

.status-error {
    color: #ef4444 !important;
    font-weight: 600 !important;
}

/* Reduced transition frequency to avoid conflicts */
.gradio-container * {
    transition: background-color 0.3s ease, border-color 0.3s ease !important;
}

/* Chat bubble styling */
.message {
    border-radius: 18px !important;
    padding: 12px 16px !important;
    margin: 8px 0 !important;
    max-width: 80% !important;
}

.user-message {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    color: white !important;
    margin-left: auto !important;
}

.bot-message {
    background: #f8fafc !important;
    border: 1px solid #e2e8f0 !important;
}

/* Metadata tooltip styling - Enhanced */
.metadata-icon {
    display: inline-block;
    margin-left: 8px;
    cursor: help;
    opacity: 0.6;
    transition: opacity 0.3s ease, transform 0.2s ease;
    font-size: 14px;
    user-select: none;
    vertical-align: middle;
}

.metadata-icon:hover {
    opacity: 1;
    transform: scale(1.1);
}

/* Enhanced tooltip styling */
.metadata-icon[title]:hover::after {
    content: attr(title);
    position: absolute;
    bottom: 100%;
    left: 50%;
    transform: translateX(-50%);
    background: rgba(0, 0, 0, 0.9);
    color: white;
    padding: 8px 12px;
    border-radius: 6px;
    font-size: 12px;
    white-space: pre-line;
    z-index: 1000;
    box-shadow: 0 4px 12px rgba(0, 0, 0, 0.3);
    margin-bottom: 5px;
    min-width: 200px;
    text-align: left;
}

.metadata-icon[title]:hover::before {
    content: '';
    position: absolute;
    bottom: 100%;
    left: 50%;
    transform: translateX(-50%);
    border: 5px solid transparent;
    border-top-color: rgba(0, 0, 0, 0.9);
    z-index: 1001;
}

/* Compact system prompt */
.compact-prompt {
    min-height: 40px !important;
    transition: min-height 0.3s ease !important;
}

.compact-prompt:focus {
    min-height: 80px !important;
}
"""

# Main application
with gr.Blocks(css=custom_css, title="🔮 AI Chat Assistant") as demo:
    # Add page refresh warning script
    gr.HTML(add_page_refresh_warning())
    
    # Header
    with gr.Row():
        title = gr.Markdown("# 🔮 AI Chat Assistant (No Model Loaded)", elem_classes="header-text")
    
    with gr.Row(equal_height=True):
        # Main chat area (left side - 70% width)
        with gr.Column(scale=7, elem_classes="chat-container"):
            # Compact system prompt (changed from 4 lines to 1)
            system_prompt = gr.Textbox(
                label="🎯 System Prompt",
                value="You are a helpful AI assistant.",
                lines=1,  # Changed from 4 to 1
                elem_classes="input-field compact-prompt"
            )
            
            # Generation settings in accordion
            with gr.Accordion("⚙️ Generation Settings", open=False, elem_classes="accordion"):
                with gr.Row():
                    temperature = gr.Slider(0.0, 2.0, DEFAULT_TEMPERATURE, step=0.05, label="🌡️ Temperature")
                    top_p = gr.Slider(0.0, 1.0, DEFAULT_TOP_P, step=0.01, label="🎯 Top-p")
                with gr.Row():
                    top_k = gr.Slider(1, 200, DEFAULT_TOP_K, step=1, label="🔝 Top-k")
                    rep_penalty = gr.Slider(1.0, 2.0, DEFAULT_REP_PENALTY, step=0.01, label="🔄 Repetition Penalty")
                
                # Memory optimization for chat (moved here to be defined before use)
                memory_opt_chat = gr.Checkbox(
                    label="🧠 Memory Optimization for Chat", 
                    value=True,
                    info="Use memory optimization during chat generation (disables KV cache)"
                )
            
            # Chat interface using original gr.ChatInterface for fast streaming and stop button
            chatbot = gr.Chatbot(
                height=500,
                latex_delimiters=[
                    {"left": "$", "right": "$", "display": True},
                    {"left": "$", "right": "$", "display": False},
                    {"left": "\\(", "right": "\\)", "display": False},
                    {"left": "\\[", "right": "\\]", "display": True}
                ],
                show_copy_button=True,
                avatar_images=("👤", "🤖"),
                type="messages",
                render_markdown=True
            )
            
            chat_interface = gr.ChatInterface(
                fn=chat_with_model,
                chatbot=chatbot,
                additional_inputs=[system_prompt, temperature, top_p, top_k, rep_penalty, memory_opt_chat],
                type="messages",
                submit_btn="Send 📤",
                stop_btn="⏹️ Stop"
            )
        
        # Control panel (right side - 30% width)
        with gr.Column(scale=3, elem_classes="control-panel"):
            # Model status and controls
            with gr.Group():
                gr.Markdown("### 🚀 Model Controls")
                
                with gr.Row():
                    load_btn = gr.Button("🚀 Load Model", variant="primary", elem_classes="btn-primary")
                    unload_btn = gr.Button("🗑️ Unload", variant="secondary", elem_classes="btn-secondary")
                
                model_status = gr.Textbox(
                    label="📊 Status", 
                    value="Model not loaded",
                    interactive=False,
                    elem_classes="input-field"
                )
                
                progress_display = gr.Textbox(
                    label="📈 Progress",
                    value="Ready to load model",
                    interactive=False,
                    elem_classes="input-field"
                )
            
            # Model selection
            with gr.Group():
                gr.Markdown("### 🎛️ Model Selection")
                
                model_source = gr.Radio(
                    choices=["Hugging Face Model", "Local Path"],
                    value="Local Path",  # Changed default to Local Path
                    label="📍 Model Source"
                )
                
                # HF Model search and selection (initially hidden)
                with gr.Group(visible=False) as hf_group:
                    model_search = gr.Textbox(
                        label="🔍 Search Models",
                        placeholder="e.g., microsoft/Phi-3, meta-llama/Llama-3, ykarout/your-model",
                        elem_classes="input-field"
                    )
                    
                    hf_model = gr.Dropdown(
                        label="📋 Select Model",
                        choices=[],
                        interactive=True,
                        elem_classes="input-field dropdown-container",
                        allow_custom_value=True,  # Allow typing custom model names
                        filterable=True  # Enable filtering
                    )
                
                # Local path group (visible by default)
                with gr.Group(visible=True) as local_group:
                    local_path = gr.Textbox(
                        value=LOCAL_MODELS_BASE,  # Changed default to new base location
                        label="📁 Local Models Base Path",
                        elem_classes="input-field"
                    )
                    
                    # Button to refresh local models
                    refresh_local_btn = gr.Button("🔄 Scan Local Models", elem_classes="btn-secondary")
                    
                    # Dropdown for local models with better configuration
                    local_models_dropdown = gr.Dropdown(
                        label="📋 Available Local Models",
                        choices=[],
                        interactive=True,
                        elem_classes="input-field dropdown-container",
                        allow_custom_value=False,  # Don't allow custom for local models
                        filterable=True  # Enable filtering
                    )
                
                quantization = gr.Radio(
                    choices=["4bit", "8bit", "bf16", "f16"],
                    value="4bit",
                    label="⚡ Quantization"
                )
                
                # Advanced memory optimization toggle
                memory_optimization = gr.Checkbox(
                    label="🧠 Advanced Memory Optimization", 
                    value=True,
                    info="Reduces VRAM usage but may slightly impact speed"
                )
                
                # Note: Memory optimization for chat is now in Generation Settings
            
            # Memory stats with cleanup buttons
            with gr.Group():
                gr.Markdown("### 💾 System Status")
                memory_info = gr.HTML()
                with gr.Row():
                    refresh_btn = gr.Button("↻ Refresh Stats", elem_classes="btn-secondary")
                    cleanup_btn = gr.Button("🧹 Clean Memory", elem_classes="btn-secondary")
                with gr.Row():
                    nuclear_btn = gr.Button("☢️ Nuclear Cleanup", elem_classes="btn-secondary", variant="stop")
    
    # Event handlers
    
    # Model search functionality for HF
    model_search.change(
        update_model_dropdown,
        inputs=[model_search],
        outputs=[hf_model]
    )
    
    # Show/hide model selection based on source
    def toggle_model_source(choice):
        return (
            gr.Group(visible=choice == "Hugging Face Model"),
            gr.Group(visible=choice == "Local Path")
        )
    
    model_source.change(
        toggle_model_source,
        inputs=[model_source],
        outputs=[hf_group, local_group]
    )
    
    # Local model scanning
    refresh_local_btn.click(
        update_local_models_dropdown,
        inputs=[local_path],
        outputs=[local_models_dropdown]
    )
    
    # Auto-scan on path change
    local_path.change(
        update_local_models_dropdown,
        inputs=[local_path],
        outputs=[local_models_dropdown]
    )
    
    # Model loading with progress
    load_btn.click(
        load_model_with_progress,
        inputs=[model_source, hf_model, local_path, local_models_dropdown, quantization, memory_optimization],
        outputs=[progress_display]
    ).then(
        lambda: "✅ Model loaded successfully!" if model is not None else "❌ Model loading failed",
        outputs=[model_status]
    ).then(
        get_memory_stats,
        outputs=[memory_info]
    ).then(
        update_model_name,
        outputs=[title]
    )
    
    # Model unloading
    unload_btn.click(
        unload_model,
        outputs=[model_status]
    ).then(
        lambda: "Ready to load model",
        outputs=[progress_display]
    ).then(
        get_memory_stats,
        outputs=[memory_info]
    ).then(
        lambda: "# 🔮 AI Chat Assistant (No Model Loaded)",
        outputs=[title]
    )
    
    # Refresh memory stats
    refresh_btn.click(get_memory_stats, outputs=[memory_info])
    
    # Manual memory cleanup
    cleanup_btn.click(cleanup_memory, outputs=[]).then(
        get_memory_stats, outputs=[memory_info]
    )
    
    # Nuclear memory cleanup
    nuclear_btn.click(nuclear_memory_cleanup, outputs=[]).then(
        get_memory_stats, outputs=[memory_info]
    )
    
    # Initialize on load
    demo.load(get_memory_stats, outputs=[memory_info])
    demo.load(
        lambda: update_local_models_dropdown(LOCAL_MODELS_BASE),
        outputs=[local_models_dropdown]
    )

# Enable queue for streaming
demo.queue()