"""Speaker profile management system for Trax platform.

This module provides speaker profile management with embedding storage,
similarity comparison, and persistent caching for improved speaker recognition.
"""

import logging
from pathlib import Path
from typing import Any, Dict, List, Optional

import numpy as np

from ..base.services import BaseService
from .diarization_types import (
    SpeakerProfile, ProfileMatch, SpeakerProfileManagerProtocol,
    ProfileNotFoundError, ProfileValidationError
)
from .diarization_utils import find_similar_speakers
from .speaker_profile_utils import (
    load_existing_profiles, add_speaker_profile, update_speaker_profile,
    remove_speaker_profile, save_all_profiles, load_all_profiles,
    get_profile_statistics
)

logger = logging.getLogger(__name__)


class SpeakerProfileManager(BaseService):
    """Manages speaker profiles with embedding storage and similarity matching.
    
    Provides efficient speaker recognition through embedding comparison,
    persistent storage, and caching mechanisms.
    """
    
    def __init__(self, storage_dir: Optional[Path] = None):
        """Initialize the SpeakerProfileManager.
        
        Args:
            storage_dir: Directory for storing profile files
        """
        super().__init__(name="SpeakerProfileManager")
        self.storage_dir = storage_dir or Path("data/speaker_profiles")
        self.storage_dir.mkdir(parents=True, exist_ok=True)
        
        self.profiles: Dict[str, SpeakerProfile] = {}
        self.embeddings_cache: Dict[str, np.ndarray] = {}
        self.similarity_threshold = 0.7
        self.max_profiles = 1000  # Prevent memory issues
        
        # Load existing profiles
        load_existing_profiles(self.storage_dir, self.profiles, self.embeddings_cache)
        
        logger.info(f"SpeakerProfileManager initialized with {len(self.profiles)} profiles")
    
    async def _initialize_impl(self) -> None:
        """Initialize the speaker profile manager."""
        try:
            # Load existing profiles
            load_existing_profiles(self.storage_dir, self.profiles, self.embeddings_cache)
            logger.info(f"SpeakerProfileManager initialized with {len(self.profiles)} profiles")
        except Exception as e:
            logger.error(f"Failed to initialize SpeakerProfileManager: {e}")
            raise
    
    def add_speaker(
        self, 
        speaker_id: str, 
        embedding: np.ndarray, 
        **kwargs
    ) -> SpeakerProfile:
        """Add a new speaker profile.
        
        Args:
            speaker_id: Unique identifier for the speaker
            embedding: Speaker embedding vector
            **kwargs: Additional profile data (name, segments, etc.)
            
        Returns:
            Created speaker profile
            
        Raises:
            ProfileValidationError: If profile data is invalid
        """
        return add_speaker_profile(
            speaker_id, embedding, self.profiles, self.embeddings_cache,
            self.storage_dir, self.max_profiles, **kwargs
        )
    
    def get_speaker(self, speaker_id: str) -> Optional[SpeakerProfile]:
        """Get a speaker profile by ID.
        
        Args:
            speaker_id: ID of the speaker to retrieve
            
        Returns:
            Speaker profile if found, None otherwise
        """
        return self.profiles.get(speaker_id)
    
    def find_similar_speakers(
        self, 
        embedding: np.ndarray, 
        threshold: float = 0.7
    ) -> List[ProfileMatch]:
        """Find speakers with similar embeddings.
        
        Args:
            embedding: Query embedding to compare against
            threshold: Minimum similarity threshold
            
        Returns:
            List of profile matches sorted by similarity score
        """
        return find_similar_speakers(
            embedding, self.embeddings_cache, self.profiles, threshold
        )
    
    def update_speaker(
        self, 
        speaker_id: str, 
        embedding: np.ndarray, 
        **kwargs
    ) -> SpeakerProfile:
        """Update an existing speaker profile.
        
        Args:
            speaker_id: ID of the speaker to update
            embedding: New embedding vector
            **kwargs: Additional profile data to update
            
        Returns:
            Updated speaker profile
            
        Raises:
            ProfileNotFoundError: If speaker profile doesn't exist
        """
        return update_speaker_profile(
            speaker_id, embedding, self.profiles, self.embeddings_cache,
            self.storage_dir, **kwargs
        )
    
    def remove_speaker(self, speaker_id: str) -> bool:
        """Remove a speaker profile.
        
        Args:
            speaker_id: ID of the speaker to remove
            
        Returns:
            True if profile was removed, False if not found
        """
        return remove_speaker_profile(
            speaker_id, self.profiles, self.embeddings_cache, self.storage_dir
        )
    
    def save_profiles(self, file_path: Optional[Path] = None) -> bool:
        """Save all profiles to disk.
        
        Args:
            file_path: Optional custom file path
            
        Returns:
            True if successful, False otherwise
        """
        if file_path is None:
            file_path = self.storage_dir / "profiles_backup.json"
        
        return save_all_profiles(self.profiles, file_path)
    
    def load_profiles(self, file_path: Path) -> bool:
        """Load profiles from disk.
        
        Args:
            file_path: Path to the profiles file
            
        Returns:
            True if successful, False otherwise
        """
        return load_all_profiles(file_path, self.profiles, self.embeddings_cache)
    
    def get_profile_stats(self) -> Dict[str, Any]:
        """Get statistics about stored profiles."""
        return get_profile_statistics(self.profiles)
    
    def cleanup(self):
        """Clean up resources and save profiles."""
        try:
            self.save_profiles()
            logger.info("SpeakerProfileManager cleanup completed")
        except Exception as e:
            logger.error(f"Failed to cleanup SpeakerProfileManager: {e}")
