"""Cross-modal retrieval utilities."""
import logging
from typing import List, Dict, Optional
from datetime import datetime, timedelta
import numpy as np

logger = logging.getLogger(__name__)


def cross_modal_search(
    query_embedding: List[float],
    candidate_embeddings: List[Dict],
    modality_filter: str = None,
    top_k: int = 5,
    min_similarity: float = 0.0
) -> List[Dict]:
    """
    Search across modalities using embedding similarity.
    
    Enables finding semantically similar content regardless of modality.
    For example: finding images that match a text query, or finding text
    that describes an audio clip.
    
    Args:
        query_embedding: Query embedding vector
        candidate_embeddings: List of dicts with 'embedding', 'modality', 'content' keys
        modality_filter: Optional modality filter (e.g., 'image' to search only images)
        top_k: Number of results to return
        min_similarity: Minimum similarity threshold (0-1)
    
    Returns:
        Top-k most similar results across modalities, sorted by similarity
        
    Example:
        >>> query_emb = embedder.embed_text("red sunset over ocean")
        >>> candidates = [
        ...     {"embedding": img_emb, "modality": "image", "content": "sunset.jpg"},
        ...     {"embedding": txt_emb, "modality": "text", "content": "Beautiful sunset"}
        ... ]
        >>> results = cross_modal_search(query_emb, candidates, top_k=5)
    """
    if not candidate_embeddings:
        logger.warning("No candidate embeddings provided")
        return []
    
    query_vec = np.array(query_embedding)
    
    # Validate query vector
    if np.linalg.norm(query_vec) == 0:
        logger.warning("Query embedding has zero norm")
        return []
    
    scores = []
    for candidate in candidate_embeddings:
        # Filter by modality if specified
        if modality_filter and candidate.get('modality') != modality_filter:
            continue
        
        candidate_vec = np.array(candidate['embedding'])
        
        # Skip if candidate has zero norm
        candidate_norm = np.linalg.norm(candidate_vec)
        if candidate_norm == 0:
            logger.debug(f"Skipping candidate with zero norm: {candidate.get('content', 'unknown')}")
            continue
        
        # Cosine similarity
        similarity = np.dot(query_vec, candidate_vec) / (
            np.linalg.norm(query_vec) * candidate_norm
        )
        
        # Filter by minimum similarity
        if similarity < min_similarity:
            continue
        
        scores.append({
            'content': candidate.get('content', ''),
            'modality': candidate.get('modality', 'unknown'),
            'similarity': float(similarity),
            'metadata': candidate.get('metadata', {}),
            'id': candidate.get('id', None)
        })
    
    # Sort by similarity descending
    scores.sort(key=lambda x: x['similarity'], reverse=True)
    
    logger.info(f"Cross-modal search: {len(scores)} results (filtered from {len(candidate_embeddings)} candidates)")
    
    return scores[:top_k]


def find_image_mentions_in_text(
    image_hash: str,
    text_events: List[Dict],
    window_minutes: int = 60,
    image_timestamp: datetime = None
) -> List[Dict]:
    """
    Find text events that mention or reference an image.
    
    This enables finding conversations about images, descriptions,
    or any text that references visual content.
    
    Args:
        image_hash: Hash of the image to find mentions of
        text_events: List of text event dicts with 'content' and optionally 'timestamp'
        window_minutes: Time window to search (minutes) - only used if timestamps available
        image_timestamp: Timestamp of the image (for temporal filtering)
    
    Returns:
        List of dicts containing matched events with confidence scores
        
    Example:
        >>> text_events = [
        ...     {"content": "Look at this screenshot", "timestamp": "..."},
        ...     {"content": "The image shows a bug", "timestamp": "..."}
        ... ]
        >>> mentions = find_image_mentions_in_text("abc123", text_events)
    """
    mentions = []
    
    # Define image reference keywords
    image_keywords = [
        "image", "picture", "photo", "screenshot", "showed", "attached",
        "look at", "see this", "visual", "diagram", "chart", "graph",
        "illustration", "figure", "snapshot"
    ]
    
    for event in text_events:
        content_lower = event.get('content', '').lower()
        
        # Check for temporal proximity if timestamps available
        if image_timestamp and 'timestamp' in event:
            try:
                event_time = datetime.fromisoformat(event['timestamp']) if isinstance(event['timestamp'], str) else event['timestamp']
                time_diff = abs((event_time - image_timestamp).total_seconds() / 60)  # minutes
                
                if time_diff > window_minutes:
                    continue  # Outside time window
                    
                # Boost confidence for temporally close events
                temporal_boost = 1.0 - (time_diff / window_minutes) * 0.3
            except (ValueError, TypeError):
                temporal_boost = 1.0
        else:
            temporal_boost = 1.0
        
        # Check for image reference keywords
        keyword_matches = [kw for kw in image_keywords if kw in content_lower]
        
        if keyword_matches:
            confidence = 0.6 * temporal_boost + 0.1 * min(len(keyword_matches), 3)
            
            mentions.append({
                'event': event,
                'confidence': min(confidence, 1.0),
                'reason': f"Contains keywords: {', '.join(keyword_matches)}",
                'matched_keywords': keyword_matches
            })
        
        # Check for direct hash reference (exact match)
        if image_hash and image_hash[:8] in event.get('content', ''):
            mentions.append({
                'event': event,
                'confidence': 0.95,
                'reason': 'Direct hash reference',
                'matched_keywords': []
            })
    
    # Sort by confidence
    mentions.sort(key=lambda x: x['confidence'], reverse=True)
    
    logger.info(f"Found {len(mentions)} potential image mentions in {len(text_events)} text events")
    
    return mentions


def find_audio_mentions_in_text(
    audio_hash: str,
    text_events: List[Dict],
    window_minutes: int = 60,
    audio_timestamp: datetime = None
) -> List[Dict]:
    """
    Find text events that mention or reference an audio clip.
    
    Args:
        audio_hash: Hash of the audio to find mentions of
        text_events: List of text event dicts
        window_minutes: Time window to search (minutes)
        audio_timestamp: Timestamp of the audio (for temporal filtering)
    
    Returns:
        List of dicts containing matched events with confidence scores
    """
    mentions = []
    
    # Audio reference keywords
    audio_keywords = [
        "audio", "sound", "voice", "recording", "listen", "heard",
        "said", "speaking", "music", "clip", "transcript"
    ]
    
    for event in text_events:
        content_lower = event.get('content', '').lower()
        
        # Check for temporal proximity
        if audio_timestamp and 'timestamp' in event:
            try:
                event_time = datetime.fromisoformat(event['timestamp']) if isinstance(event['timestamp'], str) else event['timestamp']
                time_diff = abs((event_time - audio_timestamp).total_seconds() / 60)
                
                if time_diff > window_minutes:
                    continue
                    
                temporal_boost = 1.0 - (time_diff / window_minutes) * 0.3
            except (ValueError, TypeError):
                temporal_boost = 1.0
        else:
            temporal_boost = 1.0
        
        # Check for keywords
        keyword_matches = [kw for kw in audio_keywords if kw in content_lower]
        
        if keyword_matches:
            confidence = 0.6 * temporal_boost + 0.1 * min(len(keyword_matches), 3)
            
            mentions.append({
                'event': event,
                'confidence': min(confidence, 1.0),
                'reason': f"Contains keywords: {', '.join(keyword_matches)}",
                'matched_keywords': keyword_matches
            })
        
        # Direct hash reference
        if audio_hash and audio_hash[:8] in event.get('content', ''):
            mentions.append({
                'event': event,
                'confidence': 0.95,
                'reason': 'Direct hash reference',
                'matched_keywords': []
            })
    
    mentions.sort(key=lambda x: x['confidence'], reverse=True)
    
    logger.info(f"Found {len(mentions)} potential audio mentions in {len(text_events)} text events")
    
    return mentions


def compute_multimodal_relevance(
    query_modality: str,
    result_modality: str,
    similarity_score: float,
    same_session: bool = True
) -> float:
    """
    Compute relevance score for cross-modal retrieval with modality penalties.
    
    Some modality combinations are naturally more related than others.
    For example, text→text is more direct than text→audio.
    
    Args:
        query_modality: Modality of the query
        result_modality: Modality of the result
        similarity_score: Base similarity score (0-1)
        same_session: Whether query and result are from same session
    
    Returns:
        Adjusted relevance score (0-1)
    """
    # Base score
    relevance = similarity_score
    
    # Same modality gets a bonus, no cross-modal penalty
    if query_modality == result_modality:
        relevance *= 1.1
    else:
        # Cross-modal penalties/bonuses (only for different modalities)
        cross_modal_factors = {
            ('text', 'image'): 0.95,  # Text queries for images are common
            ('image', 'text'): 0.95,  # Images can be described by text
            ('text', 'audio'): 0.90,  # Text to audio is less direct
            ('audio', 'text'): 0.95,  # Audio transcripts to text are common
            ('image', 'audio'): 0.85,  # Image to audio is more distant
            ('audio', 'image'): 0.85,
        }
        
        factor = cross_modal_factors.get((query_modality, result_modality), 0.80)
        relevance *= factor
    
    # Session bonus
    if same_session:
        relevance *= 1.05
    
    # Clamp to [0, 1]
    return min(max(relevance, 0.0), 1.0)


def group_by_modality(results: List[Dict]) -> Dict[str, List[Dict]]:
    """
    Group search results by modality.
    
    Args:
        results: List of search result dicts with 'modality' key
    
    Returns:
        Dict mapping modality to list of results
        
    Example:
        >>> results = [{"modality": "text", ...}, {"modality": "image", ...}]
        >>> grouped = group_by_modality(results)
        >>> print(grouped.keys())  # ['text', 'image']
    """
    grouped = {}
    for result in results:
        modality = result.get('modality', 'unknown')
        if modality not in grouped:
            grouped[modality] = []
        grouped[modality].append(result)
    
    return grouped

