"""
MassiveSelector - Intelligent sample selection engine for maximum variety.

Selects 330+ samples from the available library for DJ Extended and Radio Edit versions,
ensuring maximum variety while maintaining coherence.
"""

import random
import hashlib
from typing import Dict, List, Set, Tuple, Optional
from collections import defaultdict


class MassiveSelector:
    """
    Selects 330+ samples intelligently for maximum variety and coherence.
    
    Strategy:
    - DJ Extended (~165 samples): Full energy, extended sections
    - Radio Edit (~165 samples): Condensed, different samples where possible
    - All samples distributed using round-robin, layering, and strategic placement
    """
    
    # Sample counts per role (based on actual library)
    ROLE_COUNTS = {
        'kick': 13,
        'snare': 13,
        'drum_loops': 47,
        'percussion': 37,
        'bells': 14,
        'plucks': 15,
        'fx': 20,
        'midi_arpeggios': 107,
        'bass': 18,
        'synths': 45,
        'vocals': 28,
        'fills': 12,
        'transitions': 8
    }
    
    # Total: 377 samples available
    TOTAL_SAMPLES = sum(ROLE_COUNTS.values())
    
    def __init__(self, library_indexer):
        """
        Initialize with library indexer dependency.
        
        Args:
            library_indexer: Object with methods like get_samples_by_role(), 
                           get_samples_by_tempo_key(), get_all_samples()
        """
        self.indexer = library_indexer
        self.selected_history = defaultdict(int)  # Track sample usage
        self.dj_selection = {}  # Cache for DJ version
        self.radio_selection = {}  # Cache for Radio version
        
    def select_for_dj_extended(self, tempo: int, key: str) -> Dict[str, List[str]]:
        """
        Select ~165 samples optimized for DJ Extended version.
        
        DJ version characteristics:
        - Extended intros/outros (8-16 bars)
        - Full drum loops throughout
        - All percussion layers
        - Extended FX sections
        - Maximum variety in arpeggios
        
        Args:
            tempo: BPM (typically 95 for reggaeton)
            key: Musical key (e.g., 'Am', 'Cm')
            
        Returns:
            Dict mapping roles to lists of sample paths
        """
        selection = {}
        
        # Get all available samples filtered by tempo/key if possible
        try:
            all_samples = self.indexer.get_samples_by_tempo_key(tempo, key)
        except AttributeError:
            all_samples = self.indexer.get_all_samples()
        
        # Strategy: Use ALL available samples in each category
        # Kicks: ALL 13 - round-robin every 2 bars
        kicks = self._get_samples_by_role('kick', all_samples, limit=13)
        selection['kick'] = kicks
        
        # Snares: ALL 13 - round-robin + fills
        snares = self._get_samples_by_role('snare', all_samples, limit=13)
        selection['snare'] = snares
        
        # Drum loops: ALL 47 - layered in sections A/B
        drum_loops = self._get_samples_by_role('drum_loops', all_samples, limit=47)
        selection['drum_loops'] = self._split_for_layering(drum_loops, layers=2)
        
        # Percussion: ALL 37 - loops + one-shots distributed
        percussion = self._get_samples_by_role('percussion', all_samples, limit=37)
        selection['percussion'] = percussion
        selection['percussion_loops'] = percussion[:20]  # First 20 as loops
        selection['percussion_oneshots'] = percussion[20:]  # Rest as one-shots
        
        # Bells: ALL 14 - lead melodies
        bells = self._get_samples_by_role('bells', all_samples, limit=14)
        selection['bells'] = bells
        
        # Plucks: ALL 15 - lead melodies
        plucks = self._get_samples_by_role('plucks', all_samples, limit=15)
        selection['plucks'] = plucks
        
        # FX: ALL 20 - risers and impacts
        fx = self._get_samples_by_role('fx', all_samples, limit=20)
        selection['fx_risers'] = fx[:12]  # First 12 as risers
        selection['fx_impacts'] = fx[12:]  # Rest as impacts
        
        # MIDI Arpeggios: ALL 107 - maximum variety
        midi = self._get_samples_by_role('midi_arpeggios', all_samples, limit=107)
        selection['midi_arpeggios'] = midi
        
        # Additional elements for DJ version
        bass = self._get_samples_by_role('bass', all_samples, limit=18)
        selection['bass'] = bass
        
        synths = self._get_samples_by_role('synths', all_samples, limit=45)
        selection['synths'] = synths
        
        vocals = self._get_samples_by_role('vocals', all_samples, limit=28)
        selection['vocals'] = vocals
        
        fills = self._get_samples_by_role('fills', all_samples, limit=12)
        selection['fills'] = fills
        
        transitions = self._get_samples_by_role('transitions', all_samples, limit=8)
        selection['transitions'] = transitions
        
        # Mark as selected for history
        self._mark_selected(selection)
        self.dj_selection = selection
        
        return selection
    
    def select_for_radio_edit(self, tempo: int, key: str) -> Dict[str, List[str]]:
        """
        Select ~165 samples optimized for Radio Edit version.
        
        Radio Edit characteristics:
        - Shorter intros/outros (4-8 bars)
        - Condensed but still varied
        - Different samples from DJ version where possible
        - Focus on hooks and immediate impact
        
        Args:
            tempo: BPM (typically 95 for reggaeton)
            key: Musical key (e.g., 'Am', 'Cm')
            
        Returns:
            Dict mapping roles to lists of sample paths
        """
        selection = {}
        
        # Get all available samples
        try:
            all_samples = self.indexer.get_samples_by_tempo_key(tempo, key)
        except AttributeError:
            all_samples = self.indexer.get_all_samples()
        
        # Get DJ selection to avoid duplicates where possible
        dj_samples = set()
        if self.dj_selection:
            for samples in self.dj_selection.values():
                if isinstance(samples, list):
                    dj_samples.update(samples)
                elif isinstance(samples, dict):
                    for sublist in samples.values():
                        dj_samples.update(sublist)
        
        # Strategy: Use different samples from DJ version where possible
        # If not enough unique samples, alternate with DJ selections
        
        # Kicks: Alternate selection - use different kicks if available
        kicks = self._get_alternate_samples('kick', all_samples, dj_samples, 13)
        selection['kick'] = kicks
        
        # Snares: Alternate selection
        snares = self._get_alternate_samples('snare', all_samples, dj_samples, 13)
        selection['snare'] = snares
        
        # Drum loops: Select different set or alternate
        drum_loops = self._get_alternate_samples('drum_loops', all_samples, dj_samples, 47)
        selection['drum_loops'] = self._split_for_layering(drum_loops, layers=2)
        
        # Percussion: Different set
        percussion = self._get_alternate_samples('percussion', all_samples, dj_samples, 37)
        selection['percussion'] = percussion
        selection['percussion_loops'] = percussion[:20]
        selection['percussion_oneshots'] = percussion[20:]
        
        # Bells: Different set
        bells = self._get_alternate_samples('bells', all_samples, dj_samples, 14)
        selection['bells'] = bells
        
        # Plucks: Different set
        plucks = self._get_alternate_samples('plucks', all_samples, dj_samples, 15)
        selection['plucks'] = plucks
        
        # FX: Different set
        fx = self._get_alternate_samples('fx', all_samples, dj_samples, 20)
        selection['fx_risers'] = fx[:12]
        selection['fx_impacts'] = fx[12:]
        
        # MIDI: Different arpeggios or alternate every other
        midi = self._get_alternate_samples('midi_arpeggios', all_samples, dj_samples, 107)
        selection['midi_arpeggios'] = midi
        
        # Additional - use different where possible
        bass = self._get_alternate_samples('bass', all_samples, dj_samples, 18)
        selection['bass'] = bass
        
        synths = self._get_alternate_samples('synths', all_samples, dj_samples, 45)
        selection['synths'] = synths
        
        vocals = self._get_alternate_samples('vocals', all_samples, dj_samples, 28)
        selection['vocals'] = vocals
        
        fills = self._get_alternate_samples('fills', all_samples, dj_samples, 12)
        selection['fills'] = fills
        
        transitions = self._get_alternate_samples('transitions', all_samples, dj_samples, 8)
        selection['transitions'] = transitions
        
        # Mark as selected
        self._mark_selected(selection)
        self.radio_selection = selection
        
        return selection
    
    def distribute_samples_across_sections(
        self, 
        samples: List[str], 
        section_count: int
    ) -> Dict[str, List[str]]:
        """
        Distribute samples evenly across song sections.
        
        Uses a round-robin distribution strategy to ensure variety
        throughout the song sections.
        
        Args:
            samples: List of sample paths to distribute
            section_count: Number of sections (e.g., 5 for intro/verse/chorus/bridge/outro)
            
        Returns:
            Dict mapping section names to lists of samples
        """
        if not samples or section_count <= 0:
            return {}
        
        sections = {}
        section_names = self._generate_section_names(section_count)
        
        # Round-robin distribution
        for i, sample in enumerate(samples):
            section_idx = i % section_count
            section_name = section_names[section_idx]
            
            if section_name not in sections:
                sections[section_name] = []
            sections[section_name].append(sample)
        
        return sections
    
    def avoid_sample_repetition(
        self, 
        selected: List[str], 
        max_uses: int = 3
    ) -> List[str]:
        """
        Ensure no sample is overused across the selection.
        
        Checks selection history and removes or replaces samples
        that have exceeded max usage count.
        
        Args:
            selected: List of selected sample paths
            max_uses: Maximum times a sample can be used (default 3)
            
        Returns:
            Filtered list with overused samples removed
        """
        filtered = []
        
        for sample in selected:
            usage_count = self.selected_history.get(sample, 0)
            if usage_count < max_uses:
                filtered.append(sample)
            # Else: skip overused sample
        
        return filtered
    
    def validate_selection_size(self, selection: Dict) -> bool:
        """
        Validate that selection has sufficient samples per role.
        
        Minimum thresholds for professional production:
        - kick: >= 8
        - snare: >= 8
        - drum_loops: >= 20
        - percussion: >= 15
        - bells: >= 8
        - plucks: >= 8
        - fx: >= 10
        - midi: >= 50
        
        Args:
            selection: Dict mapping roles to sample lists
            
        Returns:
            True if selection meets all minimums, False otherwise
        """
        minimums = {
            'kick': 8,
            'snare': 8,
            'drum_loops': 20,
            'percussion': 15,
            'bells': 8,
            'plucks': 8,
            'fx': 10,
            'midi_arpeggios': 50
        }
        
        # Handle nested dicts (like drum_loops with layers)
        def get_count(value):
            if isinstance(value, list):
                return len(value)
            elif isinstance(value, dict):
                return sum(get_count(v) for v in value.values())
            return 0
        
        for role, min_count in minimums.items():
            if role not in selection:
                print(f"Validation failed: Missing role '{role}'")
                return False
            
            actual_count = get_count(selection[role])
            if actual_count < min_count:
                print(f"Validation failed: '{role}' has {actual_count}, need {min_count}")
                return False
        
        # Check total sample count
        total = sum(get_count(v) for v in selection.values())
        if total < 150:
            print(f"Validation failed: Total samples {total}, need >= 150")
            return False
        
        print(f"Validation passed: {total} samples across {len(selection)} roles")
        return True
    
    def get_selection_stats(self, selection: Dict) -> Dict:
        """
        Get statistics about a selection.
        
        Args:
            selection: Dict mapping roles to sample lists
            
        Returns:
            Dict with statistics
        """
        stats = {
            'total_samples': 0,
            'roles': {},
            'by_category': defaultdict(int)
        }
        
        for role, samples in selection.items():
            if isinstance(samples, list):
                count = len(samples)
                stats['roles'][role] = count
                stats['total_samples'] += count
                
                # Categorize
                if role in ['kick', 'snare', 'drum_loops', 'percussion', 'fills']:
                    stats['by_category']['drums'] += count
                elif role in ['bass']:
                    stats['by_category']['bass'] += count
                elif role in ['bells', 'plucks', 'synths']:
                    stats['by_category']['melodic'] += count
                elif role in ['fx', 'fx_risers', 'fx_impacts', 'transitions']:
                    stats['by_category']['fx'] += count
                elif role in ['vocals']:
                    stats['by_category']['vocals'] += count
                elif role in ['midi_arpeggios']:
                    stats['by_category']['midi'] += count
        
        return dict(stats)
    
    def reset_history(self):
        """Reset selection history for fresh start."""
        self.selected_history.clear()
        self.dj_selection = {}
        self.radio_selection = {}
    
    # Helper methods
    
    def _get_samples_by_role(
        self, 
        role: str, 
        all_samples: List[str], 
        limit: int = None
    ) -> List[str]:
        """Get samples filtered by role with optional limit."""
        try:
            role_samples = self.indexer.get_samples_by_role(role)
        except AttributeError:
            # Fallback: filter by path containing role name
            role_samples = [s for s in all_samples if role.replace('_', '') in s.lower()]
        
        if limit and len(role_samples) > limit:
            # Use deterministic selection based on hash for consistency
            sorted_samples = sorted(role_samples)
            return sorted_samples[:limit]
        
        return role_samples
    
    def _get_alternate_samples(
        self,
        role: str,
        all_samples: List[str],
        exclude_set: Set[str],
        limit: int
    ) -> List[str]:
        """Get samples for role, preferring ones not in exclude_set."""
        role_samples = self._get_samples_by_role(role, all_samples)
        
        # Split into unique and shared
        unique = [s for s in role_samples if s not in exclude_set]
        shared = [s for s in role_samples if s in exclude_set]
        
        # Combine: unique first, then shared if needed
        combined = unique + shared
        
        if len(combined) >= limit:
            return combined[:limit]
        
        # If not enough, just return what we have
        return combined
    
    def _split_for_layering(self, samples: List[str], layers: int = 2) -> Dict[str, List[str]]:
        """Split samples into layers (A/B) for layering strategy."""
        result = {}
        chunk_size = len(samples) // layers
        
        for i in range(layers):
            layer_name = f'layer_{chr(65 + i)}'  # layer_A, layer_B, etc.
            start = i * chunk_size
            end = start + chunk_size if i < layers - 1 else len(samples)
            result[layer_name] = samples[start:end]
        
        return result
    
    def _generate_section_names(self, count: int) -> List[str]:
        """Generate section names based on count."""
        default_names = ['intro', 'verse', 'chorus', 'bridge', 'outro', 
                        'build', 'drop', 'breakdown', 'pre_chorus', 'post_chorus']
        
        if count <= len(default_names):
            return default_names[:count]
        
        # Generate numbered sections if needed
        names = default_names.copy()
        for i in range(len(default_names), count):
            names.append(f'section_{i + 1}')
        
        return names
    
    def _mark_selected(self, selection: Dict):
        """Mark samples as selected in history."""
        def mark_samples(value):
            if isinstance(value, list):
                for sample in value:
                    self.selected_history[sample] += 1
            elif isinstance(value, dict):
                for sublist in value.values():
                    mark_samples(sublist)
        
        for samples in selection.values():
            mark_samples(samples)
    
    def generate_round_robin_pattern(
        self,
        samples: List[str],
        bars_per_sample: int = 2,
        total_bars: int = 128
    ) -> List[Tuple[int, str]]:
        """
        Generate a round-robin pattern for sample playback.
        
        Args:
            samples: List of samples to cycle through
            bars_per_sample: How many bars each sample plays
            total_bars: Total song length in bars
            
        Returns:
            List of (bar_position, sample_path) tuples
        """
        pattern = []
        bar = 0
        sample_idx = 0
        
        while bar < total_bars:
            sample = samples[sample_idx % len(samples)]
            pattern.append((bar, sample))
            bar += bars_per_sample
            sample_idx += 1
        
        return pattern
    
    def create_dj_extended_structure(self, tempo: int, key: str) -> Dict:
        """
        Create complete DJ Extended structure with distributed samples.
        
        Returns structure ready for arrangement creation.
        """
        # Select samples
        selection = self.select_for_dj_extended(tempo, key)
        
        # Define DJ Extended sections (longer)
        sections = [
            ('intro', 16),
            ('build_1', 8),
            ('verse_1', 16),
            ('pre_chorus', 8),
            ('chorus_1', 16),
            ('verse_2', 16),
            ('pre_chorus', 8),
            ('chorus_2', 16),
            ('bridge', 16),
            ('final_chorus', 16),
            ('outro', 16)
        ]
        
        structure = {
            'version': 'dj_extended',
            'tempo': tempo,
            'key': key,
            'total_bars': sum(s[1] for s in sections),
            'selection': selection,
            'sections': {}
        }
        
        # Distribute samples across sections
        for role, samples in selection.items():
            if isinstance(samples, list):
                distributed = self.distribute_samples_across_sections(
                    samples, len(sections)
                )
                for section_name in distributed:
                    if section_name not in structure['sections']:
                        structure['sections'][section_name] = {}
                    structure['sections'][section_name][role] = distributed[section_name]
        
        return structure
    
    def create_radio_edit_structure(self, tempo: int, key: str) -> Dict:
        """
        Create complete Radio Edit structure with distributed samples.
        
        Returns structure ready for arrangement creation.
        """
        # Select samples (different from DJ version)
        selection = self.select_for_radio_edit(tempo, key)
        
        # Define Radio Edit sections (shorter)
        sections = [
            ('intro', 8),
            ('verse_1', 16),
            ('pre_chorus', 4),
            ('chorus_1', 8),
            ('verse_2', 16),
            ('chorus_2', 8),
            ('bridge', 8),
            ('final_chorus', 8),
            ('outro', 8)
        ]
        
        structure = {
            'version': 'radio_edit',
            'tempo': tempo,
            'key': key,
            'total_bars': sum(s[1] for s in sections),
            'selection': selection,
            'sections': {}
        }
        
        # Distribute samples across sections
        for role, samples in selection.items():
            if isinstance(samples, list):
                distributed = self.distribute_samples_across_sections(
                    samples, len(sections)
                )
                for section_name in distributed:
                    if section_name not in structure['sections']:
                        structure['sections'][section_name] = {}
                    structure['sections'][section_name][role] = distributed[section_name]
        
        return structure
