import random
import numpy as np
import traceback

try:
    import torch
except ImportError:
    torch = None
# TODO: Clarify audio data format conventions across sunoGPT project
# Currently there's inconsistency between (channels, samples) and (samples, channels) formats:
# - Some functions expect (channels, samples) - e.g. sox effects, some torch operations
# - Others expect (samples, channels) - e.g. soundfile, some numpy operations
# - This file uses heuristics to detect format, but a unified convention would be cleaner
# - Consider standardizing on one format throughout the project

try:
    import torchaudio.sox_effects as sox

    HAS_SOX = True
except:
    HAS_SOX = False

try:
    import librosa

    HAS_LIBROSA = True
except ImportError:
    HAS_LIBROSA = False

try:
    import sox as pysox

    HAS_PYSOX = True
except ImportError:
    HAS_PYSOX = False

# Log audio effects library availability once at import time
_audio_effects_priority = []
if HAS_PYSOX:
    _audio_effects_priority.append("pysox")
if HAS_LIBROSA:
    _audio_effects_priority.append("librosa")
if HAS_SOX:
    _audio_effects_priority.append("torchaudio_sox")

if _audio_effects_priority:
    print(f"Audio effects libraries available (priority order): {', '.join(_audio_effects_priority)}")
else:
    print("Warning: No audio effects libraries available (pysox, librosa, torchaudio.sox_effects)")


def apply_sample_permutation(audio_sample, sr, prob=0.3, pitch_range=(-8, 8), rate_range=(0.85, 1.15)):
    """
    Apply pitch and rate permutation to audio samples.

    Args:
        audio_sample: numpy array of audio data
        sr: sample rate (typically 24000)
        prob: probability of applying effects (default: 0.3)
        pitch_range: range of pitch shift in semitones (default: (-8, 8))
        rate_range: range of rate change factor (default: (0.85, 1.15))

    Returns:
        tuple: (modified_audio, new_duration_s, applied_effects_info)
    """
    # Calculate original duration correctly for both mono and stereo
    if audio_sample.ndim == 1:
        original_duration_s = len(audio_sample) / sr
    else:
        # For stereo, determine which dimension contains samples
        # If shape[0] > shape[1], assume format is (samples, channels), so use shape[0]
        # Otherwise, assume format is (channels, samples), so use shape[1]
        if audio_sample.shape[0] > audio_sample.shape[1]:
            original_duration_s = audio_sample.shape[0] / sr  # (samples, channels)
        else:
            original_duration_s = audio_sample.shape[1] / sr  # (channels, samples)

    # Skip if random chance doesn't trigger
    if random.random() >= prob:
        return audio_sample, original_duration_s, "none"

    # Randomly decide which effects to apply
    apply_pitch = random.choice([True, False])
    apply_rate = random.choice([True, False])

    # Skip if no effects chosen
    if not apply_pitch and not apply_rate:
        return audio_sample, original_duration_s, "none"

    # Generate random values
    pitch_semitones = random.uniform(*pitch_range) if apply_pitch else 0
    rate_factor = random.uniform(*rate_range) if apply_rate else 1.0

    effects_info = f"pitch={pitch_semitones:.1f}_rate={rate_factor:.2f}"

    # Convert to stereo if mono for processing
    if audio_sample.ndim == 1:
        audio_stereo = np.stack([audio_sample, audio_sample], axis=0)
        was_mono = True
    else:
        audio_stereo = audio_sample.T if audio_sample.shape[0] > audio_sample.shape[1] else audio_sample
        was_mono = False

    # Apply effects using the wrapper (pysox → librosa → torchaudio sox)
    try:
        modified = apply_audio_effects(audio_stereo, sr, pitch_semitones, rate_factor)

        # Convert back to original format
        if was_mono:
            modified = modified[0]  # Take first channel

        # Calculate new duration after rate change
        new_duration_s = len(modified) / sr if was_mono else modified.shape[1] / sr

        return modified, new_duration_s, effects_info

    except Exception as e:
        # If effects fail, return original audio
        return audio_sample, original_duration_s, "failed"


def _apply_sox_effects_sample(y_stereo, sr, pitch_semitones, rate_factor):
    """Apply effects using sox."""
    y_tensor = torch.from_numpy(y_stereo.astype(np.float32))
    effects = []

    # Add pitch shift effect (sox expects cents - 100ths of semitone)
    if pitch_semitones != 0:
        effects.append(["pitch", str(int(pitch_semitones * 100))])

    # Add tempo change effect (preserves pitch)
    if rate_factor != 1.0:
        effects.append(["tempo", str(rate_factor)])

    if not effects:
        return y_stereo

    # Apply effects
    modified, new_sr = sox.apply_effects_tensor(y_tensor, sr, effects)

    # Note: tempo effect preserves sample rate, pitch effect may change it
    # Resample back to original sample rate if needed
    if new_sr != sr:
        modified = torch.nn.functional.interpolate(
            modified.unsqueeze(0),
            size=int(modified.shape[1] * sr / new_sr),
            mode="linear",
            align_corners=False,
        ).squeeze(0)

    return modified.numpy()


def _apply_librosa_effects_sample(y_stereo, sr, pitch_semitones, rate_factor):
    """Apply effects using librosa."""
    channels_processed = []

    for ch in range(y_stereo.shape[0]):
        channel_data = y_stereo[ch]

        # Apply pitch shift
        if pitch_semitones != 0:
            channel_data = librosa.effects.pitch_shift(channel_data, sr=sr, n_steps=pitch_semitones)

        # Apply time stretch (rate change)
        if rate_factor != 1.0:
            channel_data = librosa.effects.time_stretch(channel_data, rate=rate_factor)

        channels_processed.append(channel_data)

    # Handle variable lengths - find the minimum length and crop all channels
    if len(channels_processed) > 1:
        min_length = min(len(ch) for ch in channels_processed)
        channels_processed = [ch[:min_length] for ch in channels_processed]

    # Stack channels back to (channels, samples) format
    modified = np.stack(channels_processed, axis=0)

    return modified


def _apply_pysox_effects_sample(y_stereo, sr, pitch_semitones, rate_factor):
    """
    Apply effects using pysox (import sox).

    Uses WSOLA algorithm for both pitch and tempo, same quality as torchaudio sox.
    """
    # Pysox expects (samples, channels) format
    y_transposed = y_stereo.T  # Convert from (channels, samples) to (samples, channels)

    # Create transformer and add effects
    tfm = pysox.Transformer()

    # Add pitch shift effect if needed
    if pitch_semitones != 0:
        tfm.pitch(pitch_semitones)

    # Add tempo change effect if needed
    if rate_factor != 1.0:
        if abs(rate_factor - 1.0) <= 0.1:
            tfm.stretch(rate_factor, audio_type="m")
        else:
            tfm.tempo(rate_factor, audio_type="m")

    # Apply effects and get output array
    # build_array expects (n_samples, n_channels) and returns same format
    modified_transposed = tfm.build_array(input_array=y_transposed, sample_rate_in=sr)

    # Convert back to (channels, samples) format
    modified = modified_transposed.T

    return modified


def apply_audio_effects(audio, sr, pitch_semitones=0, rate_factor=1.0):
    """
    Thin wrapper to apply audio effects (pitch shift and rate change).

    Fallback order: pysox → librosa → torchaudio sox → original audio

    Args:
        audio: numpy array or torch tensor of audio data (channels, samples) format
        sr: sample rate
        pitch_semitones: pitch shift in semitones (default: 0)
        rate_factor: rate change factor (default: 1.0)

    Returns:
        Modified audio in the same format as input

    Note:
        - pysox: High quality, uses WSOLA (best for transients, preserves stereo)
        - librosa: Fast, uses phase vocoder (can introduce artifacts)
        - torchaudio sox: High quality, uses WSOLA (fallback for pysox)
    """
    # Check if input is torch tensor
    is_torch = False
    if torch is not None and isinstance(audio, torch.Tensor):
        is_torch = True
        audio_np = audio.numpy()
    else:
        audio_np = audio

    # Ensure stereo format (channels, samples)
    if audio_np.ndim == 1:
        audio_np = np.stack([audio_np, audio_np], axis=0)

    # Try pysox first, then librosa, then torchaudio sox
    try:
        if HAS_PYSOX:
            try:
                modified = _apply_pysox_effects_sample(audio_np, sr, pitch_semitones, rate_factor)
            except Exception:
                if HAS_LIBROSA:
                    try:
                        modified = _apply_librosa_effects_sample(
                            audio_np, sr, pitch_semitones, rate_factor
                        )
                    except Exception:
                        if HAS_SOX:
                            modified = _apply_sox_effects_sample(
                                audio_np, sr, pitch_semitones, rate_factor
                            )
                        else:
                            raise
                elif HAS_SOX:
                    modified = _apply_sox_effects_sample(audio_np, sr, pitch_semitones, rate_factor)
                else:
                    raise
        elif HAS_LIBROSA:
            try:
                modified = _apply_librosa_effects_sample(audio_np, sr, pitch_semitones, rate_factor)
            except Exception:
                if HAS_SOX:
                    modified = _apply_sox_effects_sample(audio_np, sr, pitch_semitones, rate_factor)
                else:
                    raise
        elif HAS_SOX:
            modified = _apply_sox_effects_sample(audio_np, sr, pitch_semitones, rate_factor)
        else:
            # No audio processing library available, return original
            modified = audio_np
    except Exception as e:
        # If all methods fail, log error and return original
        print(f"Warning: All audio effects methods failed: {e}")
        modified = audio_np

    # Convert back to torch tensor if input was torch
    if is_torch:
        modified = torch.from_numpy(modified.astype(np.float32))

    return modified


def create_sample_from_stems(
    audioloader_instance, main_meta, full_song_audio, sample_rate, sample_permutation_prob=0.0
):
    """Create sample from stems using intelligent cropping algorithm."""
    try:
        # Import the stem sampling functions
        from stem_to_sample import find_content_regions

        # Try to get stems from metadata
        stem_candidates = []
        stems_data = main_meta.get("stems", {})

        # Process all available stems and categorize them
        vocal_stems = []
        other_stems = []

        for stem_type, stem_info in stems_data.items():
            try:
                # Handle different stem info formats
                if isinstance(stem_info, dict):
                    stem_path = stem_info.get("local_filepath") or stem_info.get("path")
                    stem_s3_path = stem_info.get("s3_filepath")
                    duration = stem_info.get("duration_s", main_meta.get("duration_s", 180))
                else:
                    # If stem_info is just a path string
                    stem_path = stem_info
                    stem_s3_path = None
                    duration = main_meta.get("duration_s", 180)

                if stem_path:
                    stem_meta = {
                        "local_filepath": stem_path,
                        "s3_filepath": stem_s3_path,
                        "duration_s": duration,
                    }

                    # Check if this is a vocal-related stem
                    stem_name_lower = stem_type.lower()
                    if any(
                        vocal_keyword in stem_name_lower
                        for vocal_keyword in ["vocal", "voice", "sing", "lead"]
                    ):
                        vocal_stems.append((stem_type, stem_path, stem_meta))
                    else:
                        other_stems.append((stem_type, stem_path, stem_meta))

            except Exception as e:
                print(f"🎧 AUDIOLOADER: Error processing {stem_type} stem: {e}")
                continue

        # Prioritize vocal stems, then use any other stems
        stem_candidates = vocal_stems + other_stems

        if not stem_candidates:
            return None

        # Pick a random stem (vocals are first in the list, so they're preferred)
        stem_type, stem_path, stem_meta = random.choice(stem_candidates)

        # Log stem selection info
        vocal_count = len(
            [
                s
                for s in stem_candidates
                if any(kw in s[0].lower() for kw in ["vocal", "voice", "sing", "lead"])
            ]
        )
        other_count = len(stem_candidates) - vocal_count
        # print(f"🎧 AUDIOLOADER: Using '{stem_type}' stem (available: {vocal_count} vocal, {other_count} other stems)")

        # Load stem audio using the same method as regular loading
        try:
            stem_audio, _ = audioloader_instance._load_audio_for_semantic(
                stem_meta["local_filepath"],
                stem_meta["s3_filepath"],
                stem_meta["duration_s"],
                mock=audioloader_instance.sampling_params.mock_data,
            )
            stem_audio = stem_audio.T

            if stem_audio is None or len(stem_audio) < sample_rate * 5:  # Less than 5 seconds
                # print(f"🎧 AUDIOLOADER: {stem_type} stem too short, skipping")
                return None
            stem_audio = stem_audio.T

            # Use intelligent cropping to find good sample regions
            sample_duration_s = random.uniform(3, 10)

            # Find content regions using the stem_to_sample algorithm
            # Use first channel if stereo, otherwise use the mono audio directly
            audio_for_analysis = stem_audio[0] if len(stem_audio.shape) > 1 else stem_audio
            candidate_times = find_content_regions(
                audio_for_analysis, sample_rate, sample_duration_s, num_samples=3
            )

            if len(candidate_times) == 0:
                # print(f"🎧 AUDIOLOADER: No good content regions found in {stem_type} stem")
                return None

            # Pick random start time from candidates
            start_time = random.choice(candidate_times)
            start_sample = int(start_time * sample_rate)
            end_sample = int((start_time + sample_duration_s) * sample_rate)

            # Extract sample
            data_row_sample = stem_audio.T[start_sample:end_sample].T

            # Apply pitch and rate permutation to the extracted sample
            source_type = "stem"
            if sample_permutation_prob > 0:
                # Apply permutation effects
                data_row_sample, actual_duration_s, effects_info = apply_sample_permutation(
                    data_row_sample,
                    sample_rate,  # 24kHz sample rate
                    prob=sample_permutation_prob,
                    pitch_range=(-8, 8),  # ±8 semitones (reasonable range for samples)
                    rate_range=(0.85, 1.15),  # ±15% rate change
                )

                # Always use permuted tag when permutation is enabled, regardless of the effects_info
                source_type = "stem_permuted"

                # Update sample duration if rate was changed
                sample_duration_s = actual_duration_s

            # print(f"🎧 AUDIOLOADER: Created {sample_duration_s:.1f}s sample from {stem_type} stem at {start_time:.1f}s")

            # Sanitize stem type name for control tags (lowercase, replace spaces/special chars with underscores)
            # sanitized_stem_type = stem_type.lower().replace(" ", "_").replace("-", "_")

            return data_row_sample, start_time, source_type

        except Exception as e:
            print(f"🎧 AUDIOLOADER: Error loading {stem_type} stem: {e}")
            return None

    except Exception as e:
        print(f"🎧 AUDIOLOADER: Error in stem-based sampling: {e}")
        return None


def create_sample_from_full_song(data_row, main_meta, sample_rate, sample_permutation_prob=0.0):
    """Create sample from full song (fallback method)."""
    data_row = data_row.T
    data_row_s = len(data_row) / sample_rate
    if data_row is not None and data_row_s > 10:  # 10 seconds at 24kHz
        # Take a 3-35 second segment from a random position
        sample_duration_s = random.uniform(3, min(35, data_row_s - 1))
        sample_duration_samples = int(sample_duration_s * sample_rate)

        # Calculate maximum start position to ensure we don't exceed audio length
        max_start_sample = len(data_row) - sample_duration_samples
        start_sample = random.randint(0, max_start_sample)
        start_time_s = start_sample / sample_rate  # Convert to seconds

        # Extract sample from random position
        data_row_sample = data_row[start_sample : start_sample + sample_duration_samples].T

        # Apply pitch and rate permutation if probability is set
        source_type = "full_mix"
        if sample_permutation_prob > 0:
            # Apply permutation effects
            data_row_sample, actual_duration_s, effects_info = apply_sample_permutation(
                data_row_sample,
                sample_rate,
                prob=sample_permutation_prob,
                pitch_range=(-8, 8),  # ±8 semitones
                rate_range=(0.85, 1.15),  # ±15% rate change
            )

            # Always use permuted tag when permutation is enabled
            source_type = "full_mix_permuted"

            # Update sample duration if rate was changed
            sample_duration_s = actual_duration_s

        # print(f"🎧 AUDIOLOADER: Created sample excerpt {len(data_row_sample)} samples ({sample_duration_s:.1f}s) from full song at {start_time_s:.1f}s")

        return data_row_sample, start_time_s, source_type
    return None