# textpolicy/generation/mlx_generation.py
"""
Complete MLX-LM text generation functions for RL training.

This module provides proper integration with MLX-LM for text generation RL,
including correct logprob extraction for policy gradient training.

Key functions:
- load_model: Load MLX model and tokenizer
- generate_tokens: Generate text with logprob tracking  
- compute_logprobs: Extract logprobs for RL training
- create_policy: Create policy function for rollout collection
"""

from __future__ import annotations
from typing import Dict, Optional, Tuple, Any, Callable
import mlx.core as mx
import mlx.nn as nn
try:
    from mlx_lm import load, generate
    HAS_MLX_LM = True
except ImportError:
    HAS_MLX_LM = False
    print("Warning: mlx_lm not found. Using fallback implementations.")

try:
    from mlx_lm.sample_utils import make_sampler, make_logits_processors
# sampling utilities fallback when sample_utils is unavailable
except ImportError:
    _make_sampler = None
    _make_logits_processors = None


def _get_eos_configs_for_model(
    model_path: str,
    tokenizer_config: Optional[Dict]
) -> Tuple[Optional[Dict], Dict[str, Any]]:
    """
    Determine tokenizer_config and model_config for proper EOS handling based on model type.
    """
    model_config: Dict[str, Any] = {}
    if tokenizer_config is None and "Qwen" in model_path:
        tokenizer_config = {}
    if "Qwen" in model_path:
        # For Qwen Instruct variants, let tokenizer.eos_token_id (<|im_end|>) prevail;
        # override only for base Qwen to use <|endoftext|> (151643) as EOS.
        if "Instruct" not in model_path:
            eos_id = 151643
            model_config["eos_token_id"] = eos_id
    return tokenizer_config, model_config


def _prepare_tokenizer(tokenizer: Any, verbose: bool) -> None:
    """
    Configure tokenizer verbosity and ensure EOS token IDs for stopping.
    """
    tokenizer.verbose = verbose
    # Force tokenizer's EOS to Qwen's natural <|endoftext|> when available
    eos_id = getattr(tokenizer, 'eos_token_id', None)
    if eos_id is not None:
        # Override tokenizer.eos_token to match eos_token_id for natural stopping
        tokenizer.eos_token_id = eos_id
        tokenizer.eos_token = tokenizer.convert_ids_to_tokens(eos_id)
        tokenizer.eos_token_ids = [eos_id]
        # Align pad_token to EOS to ensure MLX-LM uses EOS for padding/stopping
        tokenizer.pad_token_id = eos_id
        tokenizer.pad_token = tokenizer.eos_token


def _make_eos_safe_sampler(temp: float, top_p: float) -> Any:
    """
    Build a sampler that does not prune low-probability tokens (e.g., EOS) and encourages natural stopping.
    """
    if make_sampler is not None:
        # Use more conservative sampling parameters to encourage natural EOS generation
        # Lower min_p and ensure we keep more tokens in consideration
        return make_sampler(
            temp=temp,
            top_p=top_p,
            min_p=0.0,  # Don't filter out low-probability tokens like EOS
            min_tokens_to_keep=2,  # Keep at least 2 tokens to ensure EOS has a chance
        )
    else:
        # Fallback implementation when mlx_lm.sample_utils is not available
        return None


def _make_logits_processors(repetition_penalty: float) -> Any:
    """
    Create logits processors to enforce repetition penalty.
    """
    if make_logits_processors is not None:
        return make_logits_processors(
            repetition_penalty=repetition_penalty,
            repetition_context_size=20,
        )
    else:
        # Fallback implementation when mlx_lm.sample_utils is not available
        return None


def _extract_response_tokens(
    response: Any,
    prompt_list: Any,
    tokenizer: Any,
) -> mx.array:
    """
    Extract response token IDs from a raw generation output (string or list).
    Enhanced to better handle EOS tokens and edge cases.
    """
    if isinstance(response, list):
        return mx.array(response)
    try:
        full_tokens = tokenizer.encode(response)
        eos_id = getattr(tokenizer, 'eos_token_id', None)
        
        # First, try to find EOS token and include it in response
        if eos_id is not None and eos_id in full_tokens:
            idx = full_tokens.index(eos_id)
            # Include EOS token in response for proper reward calculation
            resp = full_tokens[len(prompt_list): idx + 1]
        else:
            # No EOS found - extract response portion without EOS
            try:
                prompt_text = tokenizer.decode(prompt_list)
                if response.startswith(prompt_text):
                    tail = response[len(prompt_text):]
                else:
                    tail = response
                resp = tokenizer.encode(tail.strip()) if tail.strip() else []
            except Exception:
                # Fallback: encode the whole response and hope for the best
                resp = tokenizer.encode(response) if response else []
        
        return mx.array(resp) if resp else mx.array([])
    except Exception:
        return mx.array([])


def load_model(
    model_path: str, 
    adapter_path: Optional[str] = None,
    tokenizer_config: Optional[Dict] = None,
    verbose: bool = False
) -> Tuple[nn.Module, Any]:
    """
    Load MLX model and tokenizer for RL training.
    
    This function properly loads MLX-LM models with support for LoRA adapters
    and ensures compatibility with our training system. Automatically configures
    proper EOS tokens for Qwen models to ensure correct generation stopping.
    
    Args:
        model_path: Path or HuggingFace model ID  
        adapter_path: Optional LoRA adapter path
        tokenizer_config: Optional tokenizer configuration for EOS tokens
        verbose: Enable debug logging for chat template application
        
    Returns:
        (model, tokenizer): MLX model and tokenizer instances
    """
    if not HAS_MLX_LM:
        raise ImportError("mlx_lm is required. Install with: pip install mlx-lm")
    
    print(f"Loading MLX model: {model_path}")
    if adapter_path:
        print(f"Loading with LoRA adapters: {adapter_path}")
    
    # Configure model & tokenizer for EOS handling based on model type
    tokenizer_config, model_config = _get_eos_configs_for_model(model_path, tokenizer_config)
    model, tokenizer = load(
        path_or_hf_repo=model_path,
        adapter_path=adapter_path,
        tokenizer_config=tokenizer_config,
        model_config=model_config,
        lazy=False,
    )
    _prepare_tokenizer(tokenizer, verbose)
    print("✓ Model loaded successfully")
    return model, tokenizer


def generate_tokens(
    model: nn.Module,
    tokenizer: Any,
    prompt_tokens: mx.array,
    max_tokens: int = 50,
    temperature: float = 0.7,  # Lower default temperature for more stable generation
    top_p: float = 0.9,        # Lower top_p for more focused sampling
    repetition_penalty: float = 1.1  # Add repetition penalty to prevent loops
) -> Tuple[mx.array, Dict[str, Any]]:
    """Generate response tokens with proper MLX-LM integration and EOS token support."""
    if not HAS_MLX_LM:
        return _simple_generate(model, prompt_tokens, max_tokens, temperature)
    
    prompt_list = prompt_tokens.tolist()
    
    # Use stream_generate instead of generate to get proper EOS token handling
    # This is the core fix - stream_generate respects EOS tokens, generate() does not
    try:
        from mlx_lm import stream_generate
        
        # EOS-safe sampling with reduced temperature for more predictable stopping
        optimized_temperature = min(temperature, 0.7)
        sampler = _make_eos_safe_sampler(optimized_temperature, top_p)
        logits_processors = _make_logits_processors(repetition_penalty) if _make_logits_processors is not None else None


        # Use stream_generate to get token-by-token generation with EOS detection
        response_segments = list(stream_generate(
            model=model,
            tokenizer=tokenizer,
            prompt=prompt_list,  # type: ignore
            max_tokens=max_tokens,
            sampler=sampler,
            logits_processors=logits_processors,
        ))
        
        # Extract tokens from response segments and detect natural EOS stopping
        response_token_list = []
        
        for segment in response_segments:
            response_token_list.append(segment.token)
            # Check if this segment indicates natural stopping (EOS token)
            if hasattr(segment, 'finish_reason') and segment.finish_reason == "stop":
                break
        
        # Convert to MLX array
        response_tokens = mx.array(response_token_list) if response_token_list else mx.array([])
        
        
    except ImportError:
        # Fallback to original generate method if stream_generate unavailable
        print("WARNING: stream_generate not available, using fallback generate method")
        optimized_temperature = min(temperature, 0.7)
        sampler = _make_eos_safe_sampler(optimized_temperature, top_p)
        logits_processors = _make_logits_processors(repetition_penalty) if _make_logits_processors is not None else None

        response = generate(
            model=model,
            tokenizer=tokenizer,
            prompt=prompt_list,  # type: ignore
            max_tokens=max_tokens,
            sampler=sampler,
            logits_processors=logits_processors,
            verbose=False,
        )
        
        response_tokens = _extract_response_tokens(response, prompt_list, tokenizer)
    
    # Compute logprobs for the response tokens
    logprobs = compute_logprobs(model, prompt_tokens, response_tokens)
    return response_tokens, {'logprob': logprobs}


def _truncate_repetitive_text(text: str, max_repetitions: int = 3) -> str:
    """
    Truncate text if it contains excessive repetitions.
    
    This helps prevent the model from generating endless loops of the same tokens.
    """
    words = text.split()
    if len(words) < 4:
        return text
    
    # Check for word repetition
    for i in range(len(words) - max_repetitions):
        if len(set(words[i:i+max_repetitions])) == 1:
            # Found repetition, truncate here
            return ' '.join(words[:i])
    
    # Check for character repetition (like "5555555")
    for i in range(len(text) - max_repetitions):
        if len(set(text[i:i+max_repetitions])) == 1:
            # Found character repetition, truncate here
            return text[:i]
    
    return text


def _simple_generate(
    model: nn.Module, 
    prompt_tokens: mx.array, 
    max_tokens: int, 
    temperature: float
) -> Tuple[mx.array, Dict[str, Any]]:
    """
    Simple fallback generation for development without MLX-LM.
    
    This provides basic autoregressive generation for testing when
    MLX-LM is not available.
    """
    current_tokens = prompt_tokens
    generated = []
    
    for _ in range(max_tokens):
        # Model forward pass
        logits = model(current_tokens[None])  # Add batch dimension
        next_token_logits = logits[0, -1, :]  # Last token logits
        
        # Temperature scaling
        if temperature > 0:
            scaled_logits = next_token_logits / temperature
        else:
            scaled_logits = next_token_logits
        
        # Sample next token
        probs = mx.softmax(scaled_logits)
        next_token = mx.random.categorical(probs[None])[0]
        
        # Add to sequence
        generated.append(next_token)
        current_tokens = mx.concatenate([current_tokens, next_token[None]])
        
        # Stop on EOS (approximate) - avoid .item() calls
        if len(generated) > 5 and next_token < 5:  # Simple stop condition
            break
    
    response_tokens = mx.array(generated) if generated else mx.array([2])
    
    # Compute simple logprobs
    logprobs = compute_logprobs(model, prompt_tokens, response_tokens)
    
    return response_tokens, {'logprob': logprobs}


def compute_logprobs(
    model: nn.Module,
    prompt_tokens: mx.array,
    response_tokens: mx.array
) -> mx.array:
    """
    Extract log-probabilities of response_tokens under model via teacher-forcing.
    Raises on dimension mismatch or invalid (nan/inf/positive) values.
    """
    if len(response_tokens) == 0:
        return mx.array([])
    
    full_sequence = mx.concatenate([prompt_tokens, response_tokens])
    model_input = full_sequence[None] if full_sequence.ndim == 1 else full_sequence
    logits = model(model_input)
    prompt_len, response_len = len(prompt_tokens), len(response_tokens)
    prediction_logits = logits[0, prompt_len-1:prompt_len-1+response_len, :]
    if prediction_logits.shape[0] != response_len:
        raise ValueError(
            f"Logits/tokens mismatch: {prediction_logits.shape[0]} vs {response_len}"
        )
    
    log_probs = prediction_logits - mx.logsumexp(prediction_logits, axis=-1, keepdims=True)
    selected = log_probs[mx.arange(response_len), response_tokens]
    if mx.any(mx.isnan(selected)) or mx.any(mx.isinf(selected)):
        raise ValueError("Invalid logprobs (nan/inf)")
    if mx.any(selected > 0):
        print("Warning: positive logprobs detected")
    return selected


def encode(tokenizer: Any, text: str) -> mx.array:
    """
    Convert text to MLX token array.
    
    Args:
        tokenizer: MLX tokenizer
        text: Input text string
        
    Returns:
        Token array as MLX array
    """
    tokens = tokenizer.encode(text)
    return mx.array(tokens, dtype=mx.int32)


def decode(tokenizer: Any, tokens: mx.array) -> str:
    """
    Convert MLX token array to text.
    
    Args:
        tokenizer: MLX tokenizer  
        tokens: Token array
        
    Returns:
        Decoded text string
    """
    token_list = tokens.tolist()
    return tokenizer.decode(token_list)


def create_policy(
    model: nn.Module,
    tokenizer: Any,
    generation_params: Optional[Dict[str, Any]] = None
) -> Callable[[mx.array], Tuple[mx.array, Dict[str, Any]]]:
    """
    Create a policy function for RL training with automatic chat template support.
    
    This returns a pure function that can be used by rollout systems
    to generate responses and collect the data needed for training.
    
    Automatically applies chat templates for instruction models
    to enable proper EOS token generation and natural stopping behavior.
    
    Args:
        model: MLX model
        tokenizer: MLX tokenizer
        generation_params: Generation parameters (max_tokens, temperature, etc.)
        
    Returns:
        Policy function: (prompt_tokens) -> (response_tokens, info)
    """
    params = generation_params or {}
    max_tokens = params.get('max_tokens', 50)
    temperature = params.get('temperature', 0.8)
    top_p = params.get('top_p', 0.95)
    
    def policy_fn(prompt_tokens: mx.array, deterministic: bool = False) -> Tuple[mx.array, Dict[str, Any]]:
        """
        Policy function that generates responses for RL training with automatic chat template support.
        
        Automatically applies chat templates for instruction models to enable
        proper EOS token generation. This allows models to naturally end responses with
        appropriate end-of-sequence tokens instead of being artificially truncated.
        
        Args:
            prompt_tokens: Input prompt tokens
            deterministic: Whether to use deterministic generation
            
        Returns:
            (response_tokens, generation_info): Response and metadata for training
        """
        # Auto-apply chat template for instruction models
        processed_tokens = prompt_tokens
        
        try:
            # Decode tokens to check if chat template is needed
            if hasattr(tokenizer, 'decode'):
                raw_prompt = tokenizer.decode(prompt_tokens.tolist())
            else:
                # Fallback for tokenizers without decode method
                raw_prompt = str(prompt_tokens.tolist())
            
            # Let the tokenizer decide if chat template is needed
            # This works for ANY instruction model (Qwen, Llama, Mistral, etc.)
            needs_formatting = (
                hasattr(tokenizer, 'apply_chat_template') and
                # Only apply if not already formatted (avoid double-formatting)
                not any(marker in raw_prompt for marker in ['<|im_start|>', '<|endoftext|>', '<|assistant|>'])
            )
            
            if needs_formatting:
                # Convert to messages format and apply chat template
                # This uses the tokenizer's built-in knowledge of its own chat format
                messages = [{"role": "user", "content": raw_prompt.strip()}]
                formatted_prompt = tokenizer.apply_chat_template(
                    messages,
                    tokenize=False,
                    add_generation_prompt=True  # Adds <|im_start|>assistant\n for response generation
                )
                
                # Re-encode with proper formatting for EOS generation
                if hasattr(tokenizer, 'encode'):
                    processed_tokens = mx.array(tokenizer.encode(formatted_prompt))
                else:
                    # Fallback if tokenizer doesn't have encode method
                    processed_tokens = prompt_tokens
                
                # Debug logging (only in verbose mode to avoid noise)
                if hasattr(tokenizer, 'verbose') and tokenizer.verbose:
                    print(f"Applied chat template: '{formatted_prompt[:100]}...'")  # Show first 100 chars for debugging
            
        except Exception:
            # Fallback to original tokens if formatting fails
            # This ensures robustness and backward compatibility
            pass
        
        # Generate response with processed tokens
        temp = 0.0 if deterministic else temperature
        return generate_tokens(
            model=model,
            tokenizer=tokenizer,
            prompt_tokens=processed_tokens,  # Use formatted tokens for proper EOS generation
            max_tokens=max_tokens,
            temperature=temp,
            top_p=top_p
        )
    
    return policy_fn


def compute_reward(
    prompt: str,
    response: str,
    reward_type: str = "length",
    **kwargs
) -> float:
    """
    Simple reward computation for RL training.
    
    This provides basic reward functions for testing. In practice,
    you would use the sophisticated reward system from textpolicy.rewards.
    
    Args:
        prompt: Input prompt text
        response: Generated response text
        reward_type: Type of reward to compute
        **kwargs: Additional parameters for reward computation
        
    Returns:
        Reward score
    """
    if reward_type == "length":
        target_length = kwargs.get('target_length', 30)
        actual_length = len(response.split())
        # Simple length-based reward
        diff = abs(actual_length - target_length)
        return max(0.0, 1.0 - diff / target_length)
    
    elif reward_type == "keyword":
        keywords = kwargs.get('keywords', ['good', 'great', 'excellent'])
        count = sum(1 for kw in keywords if kw.lower() in response.lower())
        return count / len(keywords)
    
    else:
        # Default: simple response quality heuristic
        if len(response.strip()) == 0:
            return 0.0
        if len(response.split()) < 5:
            return 0.2
        return 0.5


# Convenience function for complete setup
def create_setup(
    model_path: str,
    generation_params: Optional[Dict[str, Any]] = None,
    adapter_path: Optional[str] = None
) -> Tuple[Callable, nn.Module, Any]:
    """
    Complete setup for MLX-LM RL training.
    
    This function combines model loading and policy creation for
    convenient setup of RL training systems.
    
    Args:
        model_path: Path or HuggingFace model ID
        generation_params: Generation parameters
        adapter_path: Optional LoRA adapter path
        
    Returns:
        (policy_fn, model, tokenizer): Complete setup for RL training
    """
    # Load model and tokenizer
    model, tokenizer = load_model(model_path, adapter_path)
    
    # Create policy function
    policy_fn = create_policy(model, tokenizer, generation_params)
    
    return policy_fn, model, tokenizer
