"""
Gain Staging Engine - Real-time level metering and gain staging for Ableton Live.

Provides professional gain staging analysis with:
- Peak and RMS level metering
- Headroom calculation
- Clipping detection
- Automatic gain adjustment recommendations
- Real-time level monitoring
"""
from __future__ import absolute_import, print_function, unicode_literals

import logging
import math
from dataclasses import dataclass, field
from typing import Dict, List, Any, Optional, Tuple
from enum import Enum

logger = logging.getLogger("GainStaging")


class LevelAlert(Enum):
    """Alert levels for track metering."""
    OK = "ok"
    WARNING = "warning"
    CLIPPING = "clipping"
    SILENT = "silent"


@dataclass
class TrackLevelInfo:
    """Comprehensive level information for a track."""
    track_index: int
    track_name: str
    peak_level: float = 0.0  # 0.0 - 1.0 (linear)
    peak_db: float = -60.0
    rms_level: float = 0.0  # 0.0 - 1.0 (linear)
    rms_db: float = -60.0
    headroom_db: float = 6.0
    is_clipping: bool = False
    is_silent: bool = False
    alert_status: LevelAlert = LevelAlert.OK
    alert_message: str = ""
    volume_fader: float = 0.0  # Current volume fader position
    recommended_volume: Optional[float] = None


@dataclass
class MasterLevelInfo:
    """Level information for master track."""
    peak_level: float = 0.0
    peak_db: float = -60.0
    rms_level: float = 0.0
    rms_db: float = -60.0
    headroom_db: float = 6.0
    is_clipping: bool = False
    lufs_estimate: float = -14.0


@dataclass
class ProjectLevelAnalysis:
    """Complete project level analysis."""
    tracks: List[TrackLevelInfo] = field(default_factory=list)
    master: MasterLevelInfo = field(default_factory=MasterLevelInfo)
    global_clipping_detected: bool = False
    tracks_needing_attention: List[int] = field(default_factory=list)
    recommendations: List[str] = field(default_factory=list)


class GainStaging:
    """
    Real-time gain staging with metering for Ableton Live.
    
    Provides:
    - Peak level measurement using track.output_meter_level
    - RMS level estimation from peak
    - Headroom calculation
    - Clipping detection
    - Automatic volume adjustment recommendations
    """
    
    # Thresholds
    CLIPPING_THRESHOLD_DB = 0.0  # dBFS
    WARNING_THRESHOLD_DB = -3.0  # Warning when approaching 0 dB
    SILENT_THRESHOLD_DB = -60.0  # Considered silent below this
    TARGET_HEADROOM_DB = -6.0  # Recommended headroom
    
    def __init__(self, ableton_connection=None):
        self.connection = ableton_connection
        self._level_history: Dict[int, List[Tuple[float, float]]] = {}  # track_index -> [(peak, rms), ...]
        self._history_size = 10
    
    def measure_peak_level(self, track_index: int, song=None) -> Dict[str, Any]:
        """
        Measure peak level of a track using Live's output_meter_level.
        
        Args:
            track_index: Index of the track to measure
            song: Ableton Live song object (optional, uses connection if not provided)
            
        Returns:
            Dict with peak_level (0.0-1.0), peak_db, and status
        """
        try:
            if song is not None and hasattr(song, "tracks"):
                track = song.tracks[track_index]
            elif self.connection and hasattr(self.connection, "_song"):
                track = self.connection._song.tracks[track_index]
            else:
                return {
                    "track_index": track_index,
                    "peak_level": 0.0,
                    "peak_db": -60.0,
                    "error": "No song object available"
                }
            
            # Try to get output_meter_level from the track's mixer_device
            mixer = getattr(track, "mixer_device", None)
            if mixer is not None:
                # output_meter_level is the live meter reading (0.0 - 1.0)
                meter_level = getattr(mixer, "output_meter_level", 0.0)
                
                # Convert to dB
                peak_db = self._linear_to_db(meter_level)
                
                return {
                    "track_index": track_index,
                    "track_name": str(getattr(track, "name", f"Track {track_index}")),
                    "peak_level": float(meter_level),
                    "peak_db": float(peak_db),
                    "volume_fader": float(getattr(mixer.volume, "value", 0.0)) if hasattr(mixer, "volume") else 0.0,
                    "measured": True
                }
            else:
                return {
                    "track_index": track_index,
                    "peak_level": 0.0,
                    "peak_db": -60.0,
                    "error": "No mixer device found"
                }
                
        except Exception as e:
            logger.error(f"Error measuring peak level for track {track_index}: {e}")
            return {
                "track_index": track_index,
                "peak_level": 0.0,
                "peak_db": -60.0,
                "error": str(e)
            }
    
    def measure_rms_level(self, track_index: int, song=None) -> Dict[str, Any]:
        """
        Estimate RMS level from peak measurement.
        
        Note: True RMS requires audio analysis. This uses a typical 
        crest factor estimate for dynamic content.
        
        Args:
            track_index: Index of the track to measure
            song: Ableton Live song object
            
        Returns:
            Dict with rms_level, rms_db, and estimation info
        """
        peak_info = self.measure_peak_level(track_index, song)
        
        if peak_info.get("error"):
            return {**peak_info, "rms_level": 0.0, "rms_db": -60.0}
        
        peak_level = peak_info["peak_level"]
        peak_db = peak_info["peak_db"]
        
        # Estimate RMS using typical crest factor of 12 dB for music
        # This is a reasonable estimate for dynamic audio
        crest_factor_db = 12.0
        rms_db = peak_db - crest_factor_db
        rms_level = self._db_to_linear(rms_db)
        
        # Update history for more accurate estimation
        if track_index not in self._level_history:
            self._level_history[track_index] = []
        
        self._level_history[track_index].append((peak_level, rms_level))
        if len(self._level_history[track_index]) > self._history_size:
            self._level_history[track_index].pop(0)
        
        # Calculate average for more stable reading
        if len(self._level_history[track_index]) > 1:
            avg_peak = sum(p for p, r in self._level_history[track_index]) / len(self._level_history[track_index])
            avg_rms = sum(r for p, r in self._level_history[track_index]) / len(self._level_history[track_index])
            rms_level = avg_rms
            rms_db = self._linear_to_db(rms_level)
        
        return {
            "track_index": track_index,
            "track_name": peak_info.get("track_name", ""),
            "rms_level": float(rms_level),
            "rms_db": float(rms_db),
            "peak_level": float(peak_level),
            "peak_db": float(peak_db),
            "crest_factor_db": crest_factor_db,
            "estimated": True,
            "samples": len(self._level_history[track_index])
        }
    
    def calculate_headroom(self, track_index: int, target_peak_db: float = -6.0, 
                          song=None) -> Dict[str, Any]:
        """
        Calculate available headroom for a track.
        
        Args:
            track_index: Index of the track
            target_peak_db: Target peak level in dB (default -6 dB for headroom)
            song: Ableton Live song object
            
        Returns:
            Dict with headroom analysis and recommendations
        """
        peak_info = self.measure_peak_level(track_index, song)
        
        if peak_info.get("error"):
            return {**peak_info, "headroom_db": 0.0}
        
        peak_db = peak_info["peak_db"]
        current_volume = peak_info.get("volume_fader", 0.85)
        
        # Calculate headroom (distance from 0 dBFS)
        headroom_db = -peak_db if peak_db < 0 else 0.0
        
        # Calculate how much we can increase before hitting target
        adjustment_needed = target_peak_db - peak_db
        
        # Calculate recommended volume
        recommended_volume = None
        if adjustment_needed < 0:
            # Need to reduce gain
            recommended_volume = current_volume * self._db_to_linear(adjustment_needed)
            recommended_volume = max(0.0, min(1.0, recommended_volume))
        elif adjustment_needed > 3:
            # Can increase gain (but be conservative)
            recommended_volume = min(1.0, current_volume * self._db_to_linear(adjustment_needed * 0.5))
        
        return {
            "track_index": track_index,
            "track_name": peak_info.get("track_name", ""),
            "current_peak_db": float(peak_db),
            "target_peak_db": float(target_peak_db),
            "headroom_db": float(headroom_db),
            "adjustment_needed_db": float(adjustment_needed),
            "current_volume": float(current_volume),
            "recommended_volume": recommended_volume,
            "has_headroom": headroom_db > abs(target_peak_db)
        }
    
    def detect_clipping(self, track_index: int, song=None) -> Dict[str, Any]:
        """
        Detect if a track is clipping.
        
        Args:
            track_index: Index of the track
            song: Ableton Live song object
            
        Returns:
            Dict with clipping status and alerts
        """
        peak_info = self.measure_peak_level(track_index, song)
        
        if peak_info.get("error"):
            return {**peak_info, "is_clipping": False, "is_warning": False}
        
        peak_db = peak_info["peak_db"]
        
        is_clipping = peak_db >= self.CLIPPING_THRESHOLD_DB
        is_warning = peak_db >= self.WARNING_THRESHOLD_DB and not is_clipping
        is_silent = peak_db <= self.SILENT_THRESHOLD_DB
        
        alert_status = LevelAlert.OK
        alert_message = ""
        
        if is_clipping:
            alert_status = LevelAlert.CLIPPING
            alert_message = "CLIPPING DETECTED! Reduce gain immediately."
        elif is_warning:
            alert_status = LevelAlert.WARNING
            alert_message = f"Approaching clipping ({peak_db:.1f} dB). Consider reducing gain."
        elif is_silent:
            alert_status = LevelAlert.SILENT
            alert_message = "Track appears silent. Check input/mute status."
        
        return {
            "track_index": track_index,
            "track_name": peak_info.get("track_name", ""),
            "peak_db": float(peak_db),
            "is_clipping": is_clipping,
            "is_warning": is_warning,
            "is_silent": is_silent,
            "alert_status": alert_status.value,
            "alert_message": alert_message,
            "needs_attention": is_clipping or is_warning
        }
    
    def analyze_all_tracks(self, song=None) -> ProjectLevelAnalysis:
        """
        Analyze levels for all tracks in the project.
        
        Args:
            song: Ableton Live song object
            
        Returns:
            ProjectLevelAnalysis with complete level information
        """
        if song is None and self.connection and hasattr(self.connection, "_song"):
            song = self.connection._song
        
        if song is None:
            return ProjectLevelAnalysis(
                recommendations=["No song object available for analysis"]
            )
        
        analysis = ProjectLevelAnalysis()
        tracks_needing_attention = []
        recommendations = []
        
        try:
            for i, track in enumerate(song.tracks):
                # Get comprehensive level info
                peak_info = self.measure_peak_level(i, song)
                rms_info = self.measure_rms_level(i, song)
                headroom_info = self.calculate_headroom(i, self.TARGET_HEADROOM_DB, song)
                clip_info = self.detect_clipping(i, song)
                
                if peak_info.get("error"):
                    continue
                
                track_info = TrackLevelInfo(
                    track_index=i,
                    track_name=peak_info.get("track_name", f"Track {i}"),
                    peak_level=peak_info.get("peak_level", 0.0),
                    peak_db=peak_info.get("peak_db", -60.0),
                    rms_level=rms_info.get("rms_level", 0.0),
                    rms_db=rms_info.get("rms_db", -60.0),
                    headroom_db=headroom_info.get("headroom_db", 6.0),
                    is_clipping=clip_info.get("is_clipping", False),
                    is_silent=clip_info.get("is_silent", False),
                    alert_status=LevelAlert(clip_info.get("alert_status", "ok")),
                    alert_message=clip_info.get("alert_message", ""),
                    volume_fader=peak_info.get("volume_fader", 0.0),
                    recommended_volume=headroom_info.get("recommended_volume")
                )
                
                analysis.tracks.append(track_info)
                
                if track_info.is_clipping:
                    tracks_needing_attention.append(i)
                    analysis.global_clipping_detected = True
                elif track_info.alert_status == LevelAlert.WARNING:
                    tracks_needing_attention.append(i)
                
                # Build recommendations
                if track_info.recommended_volume is not None:
                    vol_change = track_info.recommended_volume - track_info.volume_fader
                    if abs(vol_change) > 0.05:
                        if vol_change < 0:
                            recommendations.append(
                                f"Track {i} ({track_info.track_name}): "
                                f"Reduce volume by {abs(vol_change):.2f} to prevent clipping"
                            )
                        else:
                            recommendations.append(
                                f"Track {i} ({track_info.track_name}): "
                                f"Can increase volume by {vol_change:.2f} for better presence"
                            )
            
            # Analyze master track
            master_info = self._analyze_master(song)
            analysis.master = master_info
            
            if master_info.is_clipping:
                analysis.global_clipping_detected = True
                recommendations.insert(0, "MASTER TRACK CLIPPING! Reduce individual track levels immediately.")
            
            analysis.tracks_needing_attention = tracks_needing_attention
            analysis.recommendations = recommendations if recommendations else ["All tracks have proper gain staging."]
            
        except Exception as e:
            logger.error(f"Error analyzing all tracks: {e}")
            analysis.recommendations.append(f"Analysis error: {str(e)}")
        
        return analysis
    
    def _analyze_master(self, song) -> MasterLevelInfo:
        """Analyze master track levels."""
        try:
            master = song.master_track
            mixer = getattr(master, "mixer_device", None)
            
            if mixer is not None:
                meter_level = getattr(mixer, "output_meter_level", 0.0)
                peak_db = self._linear_to_db(meter_level)
                
                # Estimate RMS and LUFS
                crest_factor = 12.0
                rms_db = peak_db - crest_factor
                # Rough LUFS estimate (usually 1-3 dB lower than RMS)
                lufs_estimate = rms_db - 2.0
                
                return MasterLevelInfo(
                    peak_level=float(meter_level),
                    peak_db=float(peak_db),
                    rms_level=self._db_to_linear(rms_db),
                    rms_db=float(rms_db),
                    headroom_db=float(-peak_db) if peak_db < 0 else 0.0,
                    is_clipping=peak_db >= self.CLIPPING_THRESHOLD_DB,
                    lufs_estimate=float(lufs_estimate)
                )
        except Exception as e:
            logger.error(f"Error analyzing master: {e}")
        
        return MasterLevelInfo()
    
    def auto_gain_staging_with_analysis(self, track_configs: List[Dict[str, Any]], 
                                        song=None) -> Dict[str, Any]:
        """
        Enhanced auto gain staging with real level analysis.
        
        Args:
            track_configs: List of track config dicts with track_index, name, role
            song: Ableton Live song object
            
        Returns:
            Dict with gain staging results including level analysis
        """
        if song is None and self.connection and hasattr(self.connection, "_song"):
            song = self.connection._song
        
        results = {
            "tracks_adjusted": [],
            "tracks_analyzed": [],
            "clipping_detected": False,
            "alerts": [],
            "recommendations": []
        }
        
        for track_config in track_configs:
            track_index = track_config.get("track_index", 0)
            track_name = track_config.get("name", f"Track {track_index}")
            role = track_config.get("role", "").lower()
            
            # Get current level analysis
            clip_info = self.detect_clipping(track_index, song)
            headroom_info = self.calculate_headroom(track_index, self.TARGET_HEADROOM_DB, song)
            
            track_result = {
                "track_index": track_index,
                "track_name": track_name,
                "role": role,
                "current_peak_db": headroom_info.get("current_peak_db", -60.0),
                "current_volume": headroom_info.get("current_volume", 0.0),
                "is_clipping": clip_info.get("is_clipping", False),
                "alert_message": clip_info.get("alert_message", "")
            }
            
            if clip_info.get("is_clipping"):
                results["clipping_detected"] = True
                results["alerts"].append(f"Track {track_index} ({track_name}): CLIPPING!")
            
            # Determine target level based on role
            target_db = self._get_target_db_for_role(role)
            current_db = headroom_info.get("current_peak_db", -60.0)
            
            # Calculate adjustment
            adjustment_db = target_db - current_db
            current_volume = headroom_info.get("current_volume", 0.85)
            
            if abs(adjustment_db) > 1.0:  # Only adjust if difference > 1 dB
                new_volume = current_volume * self._db_to_linear(adjustment_db)
                new_volume = max(0.0, min(1.0, new_volume))
                
                track_result["recommended_volume"] = new_volume
                track_result["adjustment_db"] = adjustment_db
                track_result["volume_changed"] = True
                
                results["tracks_adjusted"].append(track_result)
                results["recommendations"].append(
                    f"Track {track_index} ({track_name}): Adjust from {current_volume:.2f} to {new_volume:.2f}"
                )
            else:
                track_result["volume_changed"] = False
                results["tracks_analyzed"].append(track_result)
        
        results["total_tracks_processed"] = len(results["tracks_adjusted"]) + len(results["tracks_analyzed"])
        results["success"] = True
        
        return results
    
    def _get_target_db_for_role(self, role: str) -> float:
        """Get target dB level based on track role."""
        role_targets = {
            "kick": -1.0,
            "snare": -2.0,
            "drums": -3.0,
            "bass": -4.0,
            "vocal": -6.0,
            "lead": -5.0,
            "synth": -7.0,
            "pad": -10.0,
            "fx": -12.0,
        }
        
        for key, target in role_targets.items():
            if key in role:
                return target
        
        return -6.0  # Default target
    
    def _linear_to_db(self, linear: float) -> float:
        """Convert linear amplitude to dB."""
        if linear <= 0:
            return -60.0
        return 20.0 * math.log10(linear)
    
    def _db_to_linear(self, db: float) -> float:
        """Convert dB to linear amplitude."""
        return 10.0 ** (db / 20.0)


# Singleton instance
_gain_staging_instance: Optional[GainStaging] = None


def get_gain_staging(ableton_connection=None) -> GainStaging:
    """Get global gain staging instance."""
    global _gain_staging_instance
    if _gain_staging_instance is None:
        _gain_staging_instance = GainStaging(ableton_connection)
    elif ableton_connection is not None:
        _gain_staging_instance.connection = ableton_connection
    return _gain_staging_instance


def reset_gain_staging():
    """Reset global gain staging instance."""
    global _gain_staging_instance
    _gain_staging_instance = None
    logger.info("Gain staging reset")


def analyze_levels(track_index: Optional[int] = None, song=None, 
                   ableton_connection=None) -> Dict[str, Any]:
    """
    Convenience function to analyze track levels.
    
    Args:
        track_index: Specific track to analyze (None for all tracks)
        song: Ableton Live song object
        ableton_connection: TCP connection to Ableton
        
    Returns:
        Level analysis results
    """
    gs = get_gain_staging(ableton_connection)
    
    if track_index is not None:
        # Analyze single track
        peak_info = gs.measure_peak_level(track_index, song)
        rms_info = gs.measure_rms_level(track_index, song)
        headroom_info = gs.calculate_headroom(track_index, -6.0, song)
        clip_info = gs.detect_clipping(track_index, song)
        
        return {
            "track_index": track_index,
            "track_name": peak_info.get("track_name", ""),
            "peak_level": peak_info.get("peak_level", 0.0),
            "peak_db": peak_info.get("peak_db", -60.0),
            "rms_level": rms_info.get("rms_level", 0.0),
            "rms_db": rms_info.get("rms_db", -60.0),
            "headroom_db": headroom_info.get("headroom_db", 6.0),
            "current_volume": headroom_info.get("current_volume", 0.0),
            "recommended_volume": headroom_info.get("recommended_volume"),
            "is_clipping": clip_info.get("is_clipping", False),
            "alert_status": clip_info.get("alert_status", "ok"),
            "alert_message": clip_info.get("alert_message", ""),
            "needs_attention": clip_info.get("needs_attention", False)
        }
    else:
        # Analyze all tracks
        analysis = gs.analyze_all_tracks(song)
        
        return {
            "tracks": [
                {
                    "track_index": t.track_index,
                    "track_name": t.track_name,
                    "peak_db": t.peak_db,
                    "rms_db": t.rms_db,
                    "headroom_db": t.headroom_db,
                    "current_volume": t.volume_fader,
                    "recommended_volume": t.recommended_volume,
                    "is_clipping": t.is_clipping,
                    "is_silent": t.is_silent,
                    "alert_status": t.alert_status.value,
                    "alert_message": t.alert_message
                }
                for t in analysis.tracks
            ],
            "master": {
                "peak_db": analysis.master.peak_db,
                "rms_db": analysis.master.rms_db,
                "headroom_db": analysis.master.headroom_db,
                "is_clipping": analysis.master.is_clipping,
                "lufs_estimate": analysis.master.lufs_estimate
            },
            "global_clipping_detected": analysis.global_clipping_detected,
            "tracks_needing_attention": analysis.tracks_needing_attention,
            "recommendations": analysis.recommendations
        }
