"""
Response reranker for selecting best outputs from multiple candidates.

This module implements response ranking and selection based on various
criteria including relevance, quality, and user feedback.
"""

from typing import List, Dict, Any, Optional, Tuple
import numpy as np
from sentence_transformers import SentenceTransformer

from ..core.config import ToxoConfig
from ..utils.logger import get_logger
from ..utils.exceptions import ModelError


class ResponseReranker:
    """
    Reranker for selecting the best response from multiple candidates.
    
    Uses various scoring methods including:
    - Semantic similarity to query
    - Quality heuristics
    - Learned preferences from feedback
    """
    
    def __init__(self, config: ToxoConfig):
        """
        Initialize response reranker.
        
        Args:
            config: Toxo configuration
        """
        self.config = config
        self.logger = get_logger("toxo.reranker")
        self.enabled = True
        
        # Embedding model for semantic similarity
        self.embedding_model = None
        self._initialize_embedding_model()
        
        # Learned preferences
        self.preference_weights = {
            "semantic_similarity": 0.4,
            "length_preference": 0.2,
            "fluency_score": 0.2,
            "feedback_score": 0.2
        }
        
        # Feedback history for learning
        self.feedback_history = []
        
        self.logger.info("ResponseReranker initialized")
    
    def _initialize_embedding_model(self):
        """Initialize embedding model for semantic similarity."""
        try:
            # Use same model as memory system for consistency
            self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
            self.logger.info("Reranker embedding model loaded")
        except Exception as e:
            self.logger.warning(f"Failed to load embedding model: {str(e)}")
            self.embedding_model = None
            self.enabled = False
    
    async def rank_responses(
        self,
        query: str,
        responses: List[str],
        context: Optional[Dict[str, Any]] = None
    ) -> str:
        """
        Rank responses and return the best one.
        
        Args:
            query: Original query
            responses: List of candidate responses
            context: Additional context for ranking
            
        Returns:
            Best response
        """
        if not responses:
            return ""
        
        if len(responses) == 1:
            return responses[0]
        
        if not self.enabled:
            self.logger.warning("Reranker disabled, returning first response")
            return responses[0]
        
        try:
            # Score all responses
            scores: List[float] = []
            for response in responses:
                try:
                    score = await self._score_response(query, response, context)
                    # Defensive: ensure score is a real float.
                    score_f = float(score) if score is not None else float("-inf")
                except Exception:
                    score_f = float("-inf")
                scores.append(score_f)
            
            # Find best response
            best_idx = np.argmax(scores)
            best_response = responses[best_idx]
            
            self.logger.debug(f"Ranked {len(responses)} responses, best score: {scores[best_idx]:.3f}")
            return best_response
            
        except Exception as e:
            self.logger.error(f"Ranking failed: {str(e)}")
            return responses[0]  # Fallback to first response
    
    async def _score_response(
        self,
        query: str,
        response: str,
        context: Optional[Dict[str, Any]] = None
    ) -> float:
        """
        Score a single response.
        
        Args:
            query: Original query
            response: Response to score
            context: Additional context
            
        Returns:
            Response score (higher is better)
        """
        scores = {}
        
        # Semantic similarity score
        scores["semantic_similarity"] = self._compute_semantic_similarity(query, response)
        
        # Length preference score
        scores["length_preference"] = self._compute_length_score(response)
        
        # Fluency score
        scores["fluency_score"] = self._compute_fluency_score(response)
        
        # Feedback-based score
        scores["feedback_score"] = self._compute_feedback_score(query, response)
        
        # Weighted combination
        total_score = sum(
            self.preference_weights[key] * score
            for key, score in scores.items()
            if key in self.preference_weights
        )
        
        return total_score
    
    def _compute_semantic_similarity(self, query: str, response: str) -> float:
        """Compute semantic similarity between query and response."""
        if not self.embedding_model:
            # Fallback to simple word overlap
            query_words = set(query.lower().split())
            response_words = set(response.lower().split())
            
            if not query_words:
                return 0.0
            
            overlap = len(query_words.intersection(response_words))
            return overlap / len(query_words)
        
        try:
            query_embedding = self.embedding_model.encode(query)
            response_embedding = self.embedding_model.encode(response)
            
            # Cosine similarity
            similarity = np.dot(query_embedding, response_embedding) / (
                np.linalg.norm(query_embedding) * np.linalg.norm(response_embedding)
            )
            
            return max(0.0, similarity)  # Ensure non-negative
            
        except Exception as e:
            self.logger.warning(f"Semantic similarity computation failed: {str(e)}")
            return 0.0
    
    def _compute_length_score(self, response: str) -> float:
        """Compute score based on response length preferences."""
        length = len(response.split())
        
        # Prefer moderate length responses (50-200 words)
        if 50 <= length <= 200:
            return 1.0
        elif length < 50:
            return length / 50.0  # Penalize very short responses
        else:
            return max(0.1, 200.0 / length)  # Penalize very long responses
    
    def _compute_fluency_score(self, response: str) -> float:
        """Compute basic fluency score based on heuristics."""
        if not response:
            return 0.0
        
        score = 1.0
        
        # Check for basic sentence structure
        sentences = response.split('.')
        if len(sentences) < 1:
            score *= 0.8
        
        # Check for reasonable word distribution
        words = response.split()
        if len(words) < 3:
            score *= 0.7
        
        # Check for excessive repetition
        unique_words = len(set(words))
        if len(words) > 0:
            repetition_ratio = unique_words / len(words)
            if repetition_ratio < 0.5:  # Too much repetition
                score *= repetition_ratio
        
        # Check for proper capitalization
        if response and response[0].islower():
            score *= 0.9
        
        return max(0.0, min(1.0, score))
    
    def _compute_feedback_score(self, query: str, response: str) -> float:
        """Compute score based on historical feedback."""
        if not self.feedback_history:
            return 0.5  # Neutral score if no feedback yet
        
        # Simple approach: check similarity to previously positively rated responses
        positive_responses = [
            fb["response"] for fb in self.feedback_history
            if fb.get("rating", 0) > 0.7
        ]
        
        if not positive_responses:
            return 0.5
        
        # Find most similar positive response
        max_similarity = 0.0
        for pos_response in positive_responses:
            similarity = self._simple_text_similarity(response, pos_response)
            max_similarity = max(max_similarity, similarity)
        
        return max_similarity
    
    def _simple_text_similarity(self, text1: str, text2: str) -> float:
        """Simple text similarity based on word overlap."""
        words1 = set(text1.lower().split())
        words2 = set(text2.lower().split())
        
        if not words1 and not words2:
            return 1.0
        if not words1 or not words2:
            return 0.0
        
        intersection = len(words1.intersection(words2))
        union = len(words1.union(words2))
        
        return intersection / union if union > 0 else 0.0
    
    async def update_from_feedback(
        self,
        input_text: str,
        expected_output: str,
        actual_output: str,
        feedback_type: str,
        rating: Optional[float] = None
    ):
        """
        Update reranker based on user feedback.
        
        Args:
            input_text: Original input
            expected_output: What user wanted
            actual_output: What was produced
            feedback_type: Type of feedback
            rating: Numeric rating if available
        """
        feedback_entry = {
            "input": input_text,
            "expected": expected_output,
            "actual": actual_output,
            "feedback_type": feedback_type,
            "rating": rating,
            "timestamp": str(np.datetime64('now'))
        }
        
        self.feedback_history.append(feedback_entry)
        
        # Keep only recent feedback (last 1000 entries)
        if len(self.feedback_history) > 1000:
            self.feedback_history = self.feedback_history[-1000:]
        
        # Adjust preference weights based on feedback
        await self._adjust_preferences(feedback_entry)
        
        self.logger.debug(f"Updated reranker with {feedback_type} feedback")
    
    async def _adjust_preferences(self, feedback: Dict[str, Any]):
        """Adjust preference weights based on feedback."""
        if feedback.get("rating") is None:
            return
        
        rating = feedback["rating"]
        
        # Simple adaptation: if rating is low, slightly adjust weights
        if rating < 0.3:
            # Poor rating - adjust weights to favor different criteria
            adjustment = 0.05
            
            # Slightly increase weight for semantic similarity
            self.preference_weights["semantic_similarity"] = min(
                0.6, self.preference_weights["semantic_similarity"] + adjustment
            )
            
            # Slightly decrease other weights
            for key in ["length_preference", "fluency_score"]:
                self.preference_weights[key] = max(
                    0.1, self.preference_weights[key] - adjustment/2
                )
        
        # Ensure weights sum to 1.0
        total_weight = sum(self.preference_weights.values())
        if total_weight > 0:
            for key in self.preference_weights:
                self.preference_weights[key] /= total_weight
    
    def get_stats(self) -> Dict[str, Any]:
        """Get reranker statistics."""
        return {
            "enabled": self.enabled,
            "feedback_count": len(self.feedback_history),
            "preference_weights": self.preference_weights.copy(),
            "embedding_model_loaded": self.embedding_model is not None,
            "avg_rating": np.mean([
                fb.get("rating", 0.5) for fb in self.feedback_history 
                if fb.get("rating") is not None
            ]) if self.feedback_history else None
        }
    
    def reset_preferences(self):
        """Reset preference weights to defaults."""
        self.preference_weights = {
            "semantic_similarity": 0.4,
            "length_preference": 0.2,
            "fluency_score": 0.2,
            "feedback_score": 0.2
        }
        self.logger.info("Reranker preferences reset to defaults")
    
    def get_state(self) -> Dict[str, Any]:
        """
        Get reranker state for serialization.
        
        Returns:
            State dictionary containing all reranker data
        """
        return {
            "enabled": self.enabled,
            "preference_weights": self.preference_weights.copy(),
            "feedback_history": self.feedback_history.copy(),
            "embedding_model_loaded": self.embedding_model is not None,
            "stats": self.get_stats()
        }
    
    def load_state(self, state: Dict[str, Any]) -> None:
        """
        Load reranker state from serialization.
        
        Args:
            state: State dictionary to load
        """
        self.enabled = state.get("enabled", True)
        self.preference_weights = state.get("preference_weights", {
            "semantic_similarity": 0.4,
            "length_preference": 0.2,
            "fluency_score": 0.2,
            "feedback_score": 0.2
        })
        self.feedback_history = state.get("feedback_history", [])
        
        self.logger.info(f"Loaded reranker state: {len(self.feedback_history)} feedback entries") 