"""
SampleRotator - Intelligent sample rotation system for Session View production.

Provides energy-based sample selection with usage tracking to avoid repetition
across scenes while maintaining sonic consistency.

Features:
- Energy-based filtering (RMS) for soft/medium/hard samples
- Usage tracking to prevent consecutive scene repetition
- BPM-aware selection with coherence validation
- Automatic sample variation across scenes

Usage:
    from engines.sample_rotator import SampleRotator
    
    rotator = SampleRotator(metadata_store)
    
    # Select samples for scene with specific energy level
    kicks = rotator.select_for_scene("kick", scene_energy=0.3, scene_index=0, count=2)
    
    # Select BPM-coherent samples
    samples = rotator.select_bpm_coherent("snare", target_bpm=95, scene_energy=0.8)
"""

import logging
import random
from pathlib import Path
from typing import Optional, List, Dict, Any, Tuple
from dataclasses import dataclass, field

from .metadata_store import SampleMetadataStore, SampleFeatures

logger = logging.getLogger("SampleRotator")


@dataclass
class SampleUsage:
    """Tracks sample usage across scenes."""
    path: str
    scene_indices: List[int] = field(default_factory=list)
    category: str = ""
    energy_levels: List[float] = field(default_factory=list)


class SampleRotator:
    """
    Intelligent sample rotation with energy-based filtering and usage tracking.
    
    Prevents sample fatigue by:
    1. Tracking which samples were used in previous scenes
    2. Avoiding same sample in consecutive scenes (configurable cooldown)
    3. Filtering samples by energy (RMS) to match scene intensity
    4. Maintaining BPM coherence across selections
    """
    
    # Energy level thresholds (RMS in dB)
    ENERGY_THRESHOLDS = {
        "low": (-60.0, -25.0),      # Soft samples for intros/breakdowns
        "medium": (-30.0, -15.0),   # Medium punch for verses
        "high": (-20.0, -5.0),      # Hard samples for drops/choruses
    }
    
    # Cooldown: minimum scenes before sample can be reused
    DEFAULT_COOLDOWN = 2
    
    def __init__(
        self,
        metadata_store: Optional[SampleMetadataStore] = None,
        cooldown_scenes: int = DEFAULT_COOLDOWN,
        bpm_tolerance: float = 5.0,
        verbose: bool = False
    ):
        """
        Initialize sample rotator.
        
        Args:
            metadata_store: SQLite metadata store for sample features
            cooldown_scenes: Minimum scenes before sample reuse (default 2)
            bpm_tolerance: BPM tolerance for coherent selection (default ±5)
            verbose: Enable verbose logging
        """
        self.metadata_store = metadata_store
        self.cooldown_scenes = cooldown_scenes
        self.bpm_tolerance = bpm_tolerance
        self.verbose = verbose
        
        # Usage tracking: category -> {path -> SampleUsage}
        self.usage_tracker: Dict[str, Dict[str, SampleUsage]] = {}
        
        # Scene counter
        self.current_scene_index = 0
        
        if verbose:
            logger.info(f"[SampleRotator] Initialized with {cooldown_scenes}-scene cooldown")
    
    def _get_energy_category(self, energy: float) -> str:
        """
        Map scene energy (0.0-1.0) to energy category.
        
        Args:
            energy: Scene energy level (0.0-1.0)
            
        Returns:
            Energy category: "low", "medium", or "high"
        """
        if energy < 0.4:
            return "low"
        elif energy < 0.75:
            return "medium"
        else:
            return "high"
    
    def _filter_by_rms(
        self,
        candidates: List[SampleFeatures],
        energy_category: str
    ) -> List[SampleFeatures]:
        """
        Filter samples by RMS based on energy category.
        
        Args:
            candidates: List of SampleFeatures
            energy_category: "low", "medium", or "high"
            
        Returns:
            Filtered list matching energy criteria
        """
        if not candidates:
            return []
        
        rms_min, rms_max = self.ENERGY_THRESHOLDS.get(energy_category, (-30.0, -15.0))
        
        filtered = []
        for sample in candidates:
            if sample.rms is None:
                # No RMS data, include as fallback
                filtered.append(sample)
            elif rms_min <= sample.rms <= rms_max:
                filtered.append(sample)
        
        # If no matches, relax criteria
        if not filtered and energy_category != "medium":
            logger.debug(f"No {energy_category} energy samples found, relaxing criteria")
            return candidates[:max(1, len(candidates) // 2)]
        
        return filtered
    
    def _exclude_recently_used(
        self,
        candidates: List[SampleFeatures],
        category: str,
        current_scene: int
    ) -> List[SampleFeatures]:
        """
        Exclude samples used within cooldown period.
        
        Args:
            candidates: List of SampleFeatures
            category: Sample category (kick, snare, etc.)
            current_scene: Current scene index
            
        Returns:
            Filtered list excluding recently used samples
        """
        if category not in self.usage_tracker:
            return candidates
        
        usage_dict = self.usage_tracker[category]
        
        filtered = []
        for sample in candidates:
            path = sample.path
            
            if path not in usage_dict:
                filtered.append(sample)
                continue
            
            usage = usage_dict[path]
            last_used_scene = max(usage.scene_indices) if usage.scene_indices else -self.cooldown_scenes
            
            # Check if sample is off cooldown
            if current_scene - last_used_scene >= self.cooldown_scenes:
                filtered.append(sample)
            elif self.verbose:
                logger.debug(f"Excluding {Path(path).name} (used in scene {last_used_scene})")
        
        # If all samples excluded (unlikely), allow recently used
        if not filtered:
            logger.warning(f"All {category} samples on cooldown, allowing recent usage")
            return candidates
        
        return filtered
    
    def _track_usage(
        self,
        selected: List[SampleFeatures],
        category: str,
        scene_index: int,
        energy: float
    ):
        """
        Track sample usage for future exclusion.
        
        Args:
            selected: List of selected SampleFeatures
            category: Sample category
            scene_index: Current scene index
            energy: Scene energy level
        """
        if category not in self.usage_tracker:
            self.usage_tracker[category] = {}
        
        for sample in selected:
            path = sample.path
            
            if path not in self.usage_tracker[category]:
                self.usage_tracker[category][path] = SampleUsage(
                    path=path,
                    category=category
                )
            
            usage = self.usage_tracker[category][path]
            usage.scene_indices.append(scene_index)
            usage.energy_levels.append(energy)
    
    def select_for_scene(
        self,
        category: str,
        scene_energy: float,
        scene_index: int,
        count: int = 1,
        bpm_range: Optional[Tuple[float, float]] = None,
        key: Optional[str] = None
    ) -> List[SampleFeatures]:
        """
        Select samples for a scene with energy-based filtering and usage tracking.
        
        Args:
            category: Sample category (kick, snare, bass, etc.)
            scene_energy: Scene energy level (0.0-1.0)
            scene_index: Current scene index
            count: Number of samples to select
            bpm_range: Optional (min_bpm, max_bpm) tuple
            key: Optional musical key filter
            
        Returns:
            List of selected SampleFeatures
        """
        if not self.metadata_store:
            logger.error("Metadata store not available")
            return []
        
        # Determine energy category
        energy_cat = self._get_energy_category(scene_energy)
        
        if self.verbose:
            logger.info(f"Selecting {count} {category} for scene {scene_index} "
                       f"(energy={scene_energy:.2f} → {energy_cat})")
        
        # Get candidates from database
        candidates = self.metadata_store.get_samples_by_category(category)
        
        if not candidates:
            logger.warning(f"No samples found in database for category: {category}")
            return []
        
        # Filter by BPM range if specified
        if bpm_range:
            min_bpm, max_bpm = bpm_range
            candidates = [s for s in candidates 
                         if s.bpm and min_bpm <= s.bpm <= max_bpm]
        
        # Filter by key if specified
        if key:
            candidates = [s for s in candidates if s.key == key]
        
        # Filter by energy (RMS)
        candidates = self._filter_by_rms(candidates, energy_cat)
        
        # Exclude recently used samples
        candidates = self._exclude_recently_used(candidates, category, scene_index)
        
        if not candidates:
            logger.warning(f"No available {category} samples after filtering")
            return []
        
        # Sort by RMS (prefer samples closest to energy target)
        rms_target = sum(self.ENERGY_THRESHOLDS[energy_cat]) / 2
        candidates.sort(key=lambda s: abs((s.rms or rms_target) - rms_target))
        
        # Select top candidates
        selected = candidates[:count]
        
        # Track usage
        self._track_usage(selected, category, scene_index, scene_energy)
        
        if self.verbose:
            names = [Path(s.path).name for s in selected]
            logger.info(f"Selected {len(selected)} {category}: {names}")
        
        return selected
    
    def select_bpm_coherent(
        self,
        category: str,
        target_bpm: float,
        scene_energy: float,
        scene_index: int,
        count: int = 1
    ) -> List[SampleFeatures]:
        """
        Select BPM-coherent samples for a scene.
        
        Uses the metadata store's coherent pool method with energy filtering.
        
        Args:
            category: Sample category
            target_bpm: Target BPM
            scene_energy: Scene energy level (0.0-1.0)
            scene_index: Current scene index
            count: Number of samples to select
            
        Returns:
            List of BPM-coherent SampleFeatures
        """
        if not self.metadata_store:
            return []
        
        # Get BPM-coherent pool
        bpm_min = target_bpm - self.bpm_tolerance
        bpm_max = target_bpm + self.bpm_tolerance
        
        return self.select_for_scene(
            category=category,
            scene_energy=scene_energy,
            scene_index=scene_index,
            count=count,
            bpm_range=(bpm_min, bpm_max)
        )
    
    def get_usage_report(self) -> Dict[str, Any]:
        """
        Generate usage report showing sample distribution across scenes.
        
        Returns:
            Dictionary with usage statistics by category
        """
        report = {
            "total_scenes": self.current_scene_index + 1,
            "categories": {},
            "most_used": [],
            "least_used": [],
        }
        
        for category, usage_dict in self.usage_tracker.items():
            cat_stats = {
                "total_samples": len(usage_dict),
                "samples_used_once": 0,
                "samples_used_multiple": 0,
                "samples": []
            }
            
            for path, usage in usage_dict.items():
                usage_count = len(usage.scene_indices)
                cat_stats["samples"].append({
                    "path": path,
                    "count": usage_count,
                    "scenes": usage.scene_indices,
                    "energies": usage.energy_levels
                })
                
                if usage_count == 1:
                    cat_stats["samples_used_once"] += 1
                else:
                    cat_stats["samples_used_multiple"] += 1
            
            report["categories"][category] = cat_stats
        
        return report
    
    def reset(self):
        """Reset usage tracking for fresh session."""
        self.usage_tracker.clear()
        self.current_scene_index = 0
        logger.info("[SampleRotator] Reset usage tracking")
    
    def advance_scene(self):
        """Advance to next scene index."""
        self.current_scene_index += 1


def create_rotator(
    db_path: str,
    cooldown_scenes: int = 2,
    bpm_tolerance: float = 5.0,
    verbose: bool = False
) -> SampleRotator:
    """
    Create and initialize a SampleRotator instance.
    
    Args:
        db_path: Path to metadata database
        cooldown_scenes: Sample reuse cooldown
        bpm_tolerance: BPM tolerance
        verbose: Enable logging
        
    Returns:
        Initialized SampleRotator
    """
    store = SampleMetadataStore(db_path)
    store.init_database()
    
    rotator = SampleRotator(
        metadata_store=store,
        cooldown_scenes=cooldown_scenes,
        bpm_tolerance=bpm_tolerance,
        verbose=verbose
    )
    
    return rotator


if __name__ == "__main__":
    # Test the SampleRotator
    import tempfile
    import os
    
    logging.basicConfig(level=logging.INFO)
    
    # Create test database
    with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
        test_db = f.name
    
    try:
        rotator = create_rotator(test_db, verbose=True)
        
        # Create test samples
        from .metadata_store import SampleFeatures
        
        test_samples = [
            SampleFeatures(
                path="/test/kick_soft.wav",
                bpm=95.0,
                rms=-35.0,
                categories=["kick"]
            ),
            SampleFeatures(
                path="/test/kick_medium.wav",
                bpm=96.0,
                rms=-20.0,
                categories=["kick"]
            ),
            SampleFeatures(
                path="/test/kick_hard.wav",
                bpm=94.0,
                rms=-10.0,
                categories=["kick"]
            ),
        ]
        
        for sample in test_samples:
            rotator.metadata_store.save_sample_features(sample.path, sample)
        
        print("\n=== Testing Energy-Based Selection ===")
        
        # Test low energy selection
        low_samples = rotator.select_for_scene(
            category="kick",
            scene_energy=0.3,
            scene_index=0,
            count=1
        )
        print(f"Low energy (0.3): {[Path(s.path).name for s in low_samples]}")
        
        # Test high energy selection
        high_samples = rotator.select_for_scene(
            category="kick",
            scene_energy=0.9,
            scene_index=1,
            count=1
        )
        print(f"High energy (0.9): {[Path(s.path).name for s in high_samples]}")
        
        # Test cooldown
        print("\n=== Testing Cooldown ===")
        rotator.current_scene_index = 2
        again_samples = rotator.select_for_scene(
            category="kick",
            scene_energy=0.9,
            scene_index=2,
            count=1
        )
        print(f"Scene 2 (cooldown active): {[Path(s.path).name for s in again_samples]}")
        
        # Get usage report
        print("\n=== Usage Report ===")
        report = rotator.get_usage_report()
        print(f"Total scenes: {report['total_scenes']}")
        for cat, stats in report['categories'].items():
            print(f"{cat}: {stats['total_samples']} samples tracked")
        
        print("\n✓ Tests completed successfully")
        
    finally:
        # Cleanup
        if os.path.exists(test_db):
            os.unlink(test_db)
