#!/usr/bin/env python3
"""
AirLLM Ollama-Compatible Wrapper

This wrapper provides an Ollama-like interface for AirLLM,
making it easy to replace Ollama in existing projects.
"""

import torch
from typing import List, Dict, Optional, Union

# Try to import airllm, preferring the local checkout if available
import sys
import os
import importlib.util

# Inject dummy BetterTransformer BEFORE importing airllm (local code needs it)
class DummyBetterTransformer:
    @staticmethod
    def transform(model):
        return model

if "optimum.bettertransformer" not in sys.modules:
    spec = importlib.util.spec_from_loader("optimum.bettertransformer", None)
    dummy_module = importlib.util.module_from_spec(spec)
    dummy_module.BetterTransformer = DummyBetterTransformer
    sys.modules["optimum.bettertransformer"] = dummy_module

# Fix RoPE scaling compatibility: patch transformers to handle "llama3" type
def _patch_rope_scaling():
    """Patch transformers LlamaConfig to handle unsupported 'llama3' RoPE scaling type."""
    try:
        from transformers import LlamaConfig
        from transformers.models.llama.configuration_llama import LlamaConfig as OriginalLlamaConfig
        
        # Store original __init__ if not already patched
        if not hasattr(OriginalLlamaConfig, '_rope_scaling_patched'):
            original_init = OriginalLlamaConfig.__init__
            
            def patched_init(self, *args, **kwargs):
                # Call original init
                original_init(self, *args, **kwargs)
                
                # Fix rope_scaling if it's "llama3" (unsupported in some transformers versions)
                if hasattr(self, 'rope_scaling') and self.rope_scaling is not None:
                    # Check if it's a dict or object
                    if isinstance(self.rope_scaling, dict):
                        if self.rope_scaling.get('type') == 'llama3':
                            print("Warning: Converting unsupported RoPE scaling 'llama3' to 'linear'")
                            self.rope_scaling['type'] = 'linear'
                            if 'factor' not in self.rope_scaling:
                                self.rope_scaling['factor'] = 1.0
                    elif hasattr(self.rope_scaling, 'type'):
                        if getattr(self.rope_scaling, 'type', None) == 'llama3':
                            print("Warning: Converting unsupported RoPE scaling 'llama3' to 'linear'")
                            # Convert to dict format
                            factor = getattr(self.rope_scaling, 'factor', 1.0)
                            self.rope_scaling = {'type': 'linear', 'factor': factor}
            
            OriginalLlamaConfig.__init__ = patched_init
            OriginalLlamaConfig._rope_scaling_patched = True
    except Exception as e:
        # If patching fails, we'll handle it in the error handler
        print(f"Warning: Could not patch RoPE scaling: {e}", file=sys.stderr)

def _patch_attention_position_embeddings():
    """Patch LlamaSdpaAttention to accept and ignore position_embeddings argument for AirLLM compatibility."""
    try:
        from transformers.models.llama import modeling_llama
        import functools
        
        # Check if LlamaSdpaAttention exists and hasn't been patched
        if hasattr(modeling_llama, 'LlamaSdpaAttention'):
            LlamaSdpaAttention = modeling_llama.LlamaSdpaAttention
            if not hasattr(LlamaSdpaAttention, '_position_embeddings_patched'):
                original_forward = LlamaSdpaAttention.forward
                
                @functools.wraps(original_forward)
                def patched_forward(self, *args, **kwargs):
                    # Remove position_embeddings if present (AirLLM compatibility)
                    kwargs.pop('position_embeddings', None)
                    # Call original forward
                    return original_forward(self, *args, **kwargs)
                
                LlamaSdpaAttention.forward = patched_forward
                LlamaSdpaAttention._position_embeddings_patched = True
    except Exception as e:
        # If patching fails, we'll handle it in the error handler
        print(f"Warning: Could not patch attention position_embeddings: {e}", file=sys.stderr)

# Apply the patches before importing airllm
_patch_rope_scaling()
_patch_attention_position_embeddings()

LOCAL_AIRLLM_PATH = "/home/ladmin/code/airllm/airllm/air_llm"
if os.path.exists(LOCAL_AIRLLM_PATH) and LOCAL_AIRLLM_PATH not in sys.path:
    sys.path.insert(0, LOCAL_AIRLLM_PATH)

try:
    from airllm import AutoModel
    AIRLLM_AVAILABLE = True
except ImportError as e:
    AIRLLM_AVAILABLE = False
    AutoModel = None
    print(f"Warning: Failed to import AirLLM: {e}", file=sys.stderr)


class AirLLMOllamaWrapper:
    """
    A wrapper that provides an Ollama-like API for AirLLM.
    
    Usage:
        # Instead of: ollama.generate(model="llama2", prompt="Hello")
        # Use: airllm_wrapper.generate(model="llama2", prompt="Hello")
    """
    
    def __init__(self, model_name: str, compression: Optional[str] = None, **kwargs):
        """
        Initialize AirLLM model.
        
        Args:
            model_name: Hugging Face model name or path (e.g., "meta-llama/Llama-3.2-3B-Instruct")
            compression: Optional compression ('4bit' or '8bit') for 3x speed improvement
            **kwargs: Additional arguments for AutoModel.from_pretrained()
        """
        if not AIRLLM_AVAILABLE or AutoModel is None:
            raise ImportError(
                "AirLLM is not available. Please install it with: pip install airllm bitsandbytes\n"
                "If you see a BetterTransformer error, you may need to install: pip install optimum[bettertransformer]"
            )
        
        print(f"Loading AirLLM model: {model_name}")
        
        # Fix RoPE scaling compatibility issue: transformers 4.39.3 doesn't support "llama3" type
        # Modify config file if it's a local path and has unsupported rope_scaling
        model_path = model_name
        if os.path.exists(model_name) or model_name.startswith('/') or model_name.startswith('~'):
            if model_name.startswith('~'):
                model_path = os.path.expanduser(model_name)
            else:
                model_path = os.path.abspath(model_name)
            
            config_json_path = os.path.join(model_path, "config.json")
            if os.path.exists(config_json_path):
                try:
                    import json
                    with open(config_json_path, 'r') as f:
                        config_data = json.load(f)
                    
                    # Check and fix rope_scaling
                    if 'rope_scaling' in config_data and config_data['rope_scaling'] is not None:
                        rope_scaling = config_data['rope_scaling']
                        if isinstance(rope_scaling, dict) and rope_scaling.get('type') == 'llama3':
                            print("Warning: Fixing unsupported RoPE scaling type 'llama3' -> 'linear'")
                            # Backup original config
                            backup_path = config_json_path + ".backup"
                            if not os.path.exists(backup_path):
                                import shutil
                                shutil.copy2(config_json_path, backup_path)
                            
                            # Fix the rope_scaling type
                            config_data['rope_scaling']['type'] = 'linear'
                            if 'factor' not in config_data['rope_scaling']:
                                config_data['rope_scaling']['factor'] = 1.0
                            
                            # Save fixed config
                            with open(config_json_path, 'w') as f:
                                json.dump(config_data, f, indent=2)
                            print(f"Fixed config saved to {config_json_path}")
                except Exception as e:
                    print(f"Warning: Could not fix config file: {e}", file=sys.stderr)
        
        # Determine max_seq_len before loading model
        # AirLLM needs this at initialization time
        max_seq_len = 2048  # Default for Llama models
        
        # Check if this is a Llama model to determine appropriate max length
        # We need to load config first to check model type
        try:
            from transformers import AutoConfig
            config = AutoConfig.from_pretrained(model_name, **{k: v for k, v in kwargs.items() if k in ['token', 'trust_remote_code']})
            model_type = getattr(config, 'model_type', '').lower()
            is_llama = 'llama' in model_type or 'llama' in model_name.lower()
            
            # Also fix rope_scaling in the loaded config object if needed
            if is_llama and hasattr(config, 'rope_scaling') and config.rope_scaling is not None:
                if isinstance(config.rope_scaling, dict) and config.rope_scaling.get('type') == 'llama3':
                    print("Warning: Converting RoPE scaling 'llama3' to 'linear' in config object")
                    config.rope_scaling['type'] = 'linear'
                    if 'factor' not in config.rope_scaling:
                        config.rope_scaling['factor'] = 1.0
                elif hasattr(config.rope_scaling, 'type') and getattr(config.rope_scaling, 'type', None) == 'llama3':
                    # Convert object to dict
                    factor = getattr(config.rope_scaling, 'factor', 1.0)
                    config.rope_scaling = {'type': 'linear', 'factor': factor}
            
            if is_llama:
                config_max = getattr(config, 'max_position_embeddings', None)
                if config_max and config_max > 0:
                    max_seq_len = min(config_max, 2048)
                else:
                    max_seq_len = 2048
            else:
                config_max = getattr(config, 'max_position_embeddings', None)
                if config_max and config_max > 0 and config_max <= 2048:
                    max_seq_len = config_max
                else:
                    max_seq_len = 512
        except Exception:
            # Fallback to defaults if config loading fails
            pass
        
        # AutoModel.from_pretrained() accepts:
        # - Hugging Face model IDs (e.g., "meta-llama/Llama-3.1-8B-Instruct")
        # - Local paths (e.g., "/path/to/local/model")
        # - Can use local_dir parameter for local models
        try:
            self.model = AutoModel.from_pretrained(
                model_name,
                compression=compression,
                max_seq_len=max_seq_len,  # Pass max_seq_len to AirLLM
                **kwargs
            )
        except ValueError as e:
            # Handle specific RoPE scaling errors
            if "Unknown RoPE scaling type" in str(e) or "rope_scaling" in str(e).lower():
                import traceback
                error_msg = (
                    f"RoPE scaling compatibility error: {e}\n"
                    "The model config uses a RoPE scaling type not supported by your transformers version.\n"
                    "If this is a local model, the config file should have been fixed automatically.\n"
                    "If the error persists, try:\n"
                    "1. For local models: Check that config.json has rope_scaling.type='linear' instead of 'llama3'\n"
                    "2. Upgrade transformers: pip install --upgrade transformers\n"
                    "3. Or downgrade to a compatible version: pip install 'transformers==4.37.0'\n"
                    f"\nFull traceback:\n{traceback.format_exc()}"
                )
                raise RuntimeError(error_msg) from e
            raise
        except Exception as e:
            import traceback
            error_msg = (
                f"Failed to load AirLLM model '{model_name}': {e}\n"
                f"Error type: {type(e).__name__}\n"
                "This is often a transformers version compatibility issue.\n"
                "Try one of these solutions:\n"
                "1. Install an older transformers version: pip install 'transformers==4.37.0'\n"
                "2. Or try: pip install 'transformers==4.38.2'\n"
                "3. If using transformers 4.39.3, try downgrading: pip install 'transformers==4.37.0'\n"
                "4. Check AirLLM compatibility with your transformers version\n"
                f"\nFull traceback:\n{traceback.format_exc()}"
            )
            raise RuntimeError(error_msg) from e
        self.model_name = model_name
        
        # Store max_length for tokenization
        self.max_length = max_seq_len
        
        # Check if this is a Llama model to determine appropriate max length
        is_llama = False
        if hasattr(self.model, 'config'):
            model_type = getattr(self.model.config, 'model_type', '').lower()
            is_llama = 'llama' in model_type or 'llama' in self.model_name.lower()
        
        if is_llama:
            # Llama models: typically support 2048-4096 tokens
            # AirLLM works well with Llama, so we can use larger chunks
            if hasattr(self.model, 'config'):
                config_max = getattr(self.model.config, 'max_position_embeddings', None)
                if config_max and config_max > 0:
                    # Use config value, but cap at 2048 for AirLLM safety
                    self.max_length = min(config_max, 2048)
                else:
                    self.max_length = 2048  # Safe default for Llama
        else:
            # For other models (e.g., DeepSeek), use conservative default
            if hasattr(self.model, 'config'):
                config_max = getattr(self.model.config, 'max_position_embeddings', None)
                if config_max and config_max > 0 and config_max <= 2048:
                    self.max_length = config_max
                else:
                    self.max_length = 512  # Very conservative
        
        print(f"Using sequence length limit: {self.max_length} (AirLLM chunk size)")
        
        print("Model loaded successfully!")
    
    def generate(
        self,
        prompt: str,
        model: Optional[str] = None,  # Ignored, kept for API compatibility
        max_tokens: int = 50,
        temperature: float = 0.7,
        top_p: float = 0.9,
        stream: bool = False,
        **kwargs
    ) -> Union[str, Dict]:
        """
        Generate text from a prompt (Ollama-compatible interface).
        
        Args:
            prompt: Input text prompt
            model: Ignored (kept for compatibility)
            max_tokens: Maximum number of tokens to generate
            temperature: Sampling temperature (0.0 to 1.0)
            top_p: Nucleus sampling parameter
            stream: If True, return streaming response (not yet implemented)
            **kwargs: Additional generation parameters
        
        Returns:
            Generated text string or dict with response
        """
        # Tokenize input with attention mask
        # AirLLM processes sequences in chunks, but each chunk must fit within the model's
        # position embedding limits. We need to ensure we don't exceed the chunk size.
        # Use the model's max_length to ensure compatibility with position embeddings
        input_tokens = self.model.tokenizer(
            prompt,
            return_tensors="pt",
            return_attention_mask=True,
            truncation=True,
            max_length=self.max_length,  # Respect model's position embedding limit
            padding=False
        )
        
        # Move to GPU if available
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        input_ids = input_tokens['input_ids'].to(device)
        attention_mask = input_tokens.get('attention_mask', None)
        if attention_mask is not None:
            attention_mask = attention_mask.to(device)
        
        # Ensure we don't exceed max_length (manual truncation as safety check)
        seq_length = input_ids.shape[1]
        if seq_length > self.max_length:
            print(f"Warning: Sequence length ({seq_length}) exceeds limit ({self.max_length}), truncating...")
            input_ids = input_ids[:, :self.max_length]
            if attention_mask is not None:
                attention_mask = attention_mask[:, :self.max_length]
            seq_length = self.max_length
        
        if seq_length >= self.max_length:
            print(f"Note: Using sequence of {seq_length} tokens (at limit: {self.max_length})")
        
        # Prepare generation parameters
        # For Llama models, we can use more tokens
        max_gen_tokens = min(max_tokens, 512)
        
        gen_kwargs = {
            'max_new_tokens': max_gen_tokens,
            'use_cache': False,  # Disable cache to avoid DynamicCache compatibility issues
            'return_dict_in_generate': True,
            'temperature': temperature,
            'top_p': top_p,
            **kwargs
        }
        
        # Add attention mask if available
        if attention_mask is not None:
            gen_kwargs['attention_mask'] = attention_mask
        
        # Generate
        try:
            with torch.inference_mode():
                generation_output = self.model.generate(input_ids, **gen_kwargs)
        except (TypeError, RuntimeError) as e:
            if "position_embeddings" in str(e) or "cannot unpack" in str(e):
                error_msg = (
                    f"AirLLM compatibility error with transformers: {e}\n"
                    "This is a known issue with AirLLM and transformers version compatibility.\n"
                    "Try one of these solutions:\n"
                    "1. Install transformers 4.37.0: pip install 'transformers==4.37.0'\n"
                    "2. Or try transformers 4.38.2: pip install 'transformers==4.38.2'\n"
                    "3. If you're using 4.39.3, it may have compatibility issues - try downgrading\n"
                    "4. Or use Ollama instead: nanobot agent -m 'Hello' (with Ollama provider)"
                )
                raise RuntimeError(error_msg) from e
            raise
        
        # Decode output - get only the newly generated tokens
        if hasattr(generation_output, 'sequences'):
            # Extract only the new tokens (after input length)
            input_length = input_ids.shape[1]
            generated_ids = generation_output.sequences[0, input_length:]
            output = self.model.tokenizer.decode(generated_ids, skip_special_tokens=True)
        else:
            # Fallback for older output formats
            output = self.model.tokenizer.decode(generation_output.sequences[0], skip_special_tokens=True)
            # Remove the input prompt from output if present
            if output.startswith(prompt):
                output = output[len(prompt):].strip()
        
        if stream:
            # For streaming, return a generator (simplified version)
            return {"response": output}
        else:
            return output
    
    def chat(
        self,
        messages: List[Dict[str, str]],
        model: Optional[str] = None,
        max_tokens: int = 50,
        temperature: float = 0.7,
        **kwargs
    ) -> str:
        """
        Chat interface (Ollama-compatible).
        
        Args:
            messages: List of message dicts with 'role' and 'content' keys
            model: Ignored (kept for compatibility)
            max_tokens: Maximum tokens to generate
            temperature: Sampling temperature
            **kwargs: Additional parameters
        
        Returns:
            Generated response string
        """
        # Try to use the model's chat template if available (for Llama, etc.)
        if hasattr(self.model.tokenizer, 'apply_chat_template') and self.model.tokenizer.chat_template:
            try:
                # Use the model's native chat template
                prompt = self.model.tokenizer.apply_chat_template(
                    messages,
                    tokenize=False,
                    add_generation_prompt=True
                )
            except Exception:
                # Fallback to simple formatting if chat template fails
                prompt = self._format_messages(messages)
        else:
            # Fallback to simple formatting
            prompt = self._format_messages(messages)
        
        return self.generate(
            prompt=prompt,
            max_tokens=max_tokens,
            temperature=temperature,
            **kwargs
        )
    
    def _format_messages(self, messages: List[Dict[str, str]]) -> str:
        """Format chat messages into a single prompt (fallback method)."""
        formatted = []
        for msg in messages:
            role = msg.get('role', 'user')
            content = msg.get('content', '')
            if role == 'system':
                formatted.append(f"System: {content}")
            elif role == 'user':
                formatted.append(f"User: {content}")
            elif role == 'assistant':
                formatted.append(f"Assistant: {content}")
        return "\n".join(formatted) + "\nAssistant:"
    
    def embeddings(self, prompt: str) -> List[float]:
        """
        Get embeddings for a prompt (simplified - returns token embeddings).
        
        Note: This is a simplified version. For full embeddings,
        you may need to access model internals.
        """
        tokens = self.model.tokenizer(
            [prompt],
            return_tensors="pt",
            truncation=True,
            max_length=512,
            padding=False
        )
        # This is a placeholder - actual embeddings would require model forward pass
        return tokens['input_ids'].tolist()[0]


# Convenience function for easy migration
def create_ollama_client(model_name: str, compression: Optional[str] = None, **kwargs):
    """
    Create an Ollama-compatible client using AirLLM.
    
    Usage:
        client = create_ollama_client("meta-llama/Llama-3.2-3B-Instruct")
        response = client.generate("Hello, how are you?")
    """
    return AirLLMOllamaWrapper(model_name, compression=compression, **kwargs)


# Example usage
if __name__ == "__main__":
    # Example 1: Basic generation
    print("Example 1: Basic Generation")
    print("=" * 60)
    
    # Initialize (this will take time on first run)
    # client = create_ollama_client("meta-llama/Llama-3.2-3B-Instruct")
    
    # Generate
    # response = client.generate("What is the capital of France?")
    # print(f"Response: {response}")
    
    print("\nExample 2: Chat Interface")
    print("=" * 60)
    
    # Chat example
    # messages = [
    #     {"role": "user", "content": "Hello! How are you?"}
    # ]
    # response = client.chat(messages)
    # print(f"Response: {response}")
    
    print("\nUncomment the code above to test!")

