import os
import torch
import torchaudio
import numpy as np
from typing import Tuple, Optional, List, Dict, Any
import soundfile as sf
import librosa
import tempfile

def load_audio(file_path: str) -> Tuple[torch.Tensor, int]:
    """
    Load audio file and return waveform and sample rate.

    Args:
        file_path: Path to audio file

    Returns:
        Tuple of (waveform, sample_rate)
    """
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"Audio file not found: {file_path}")

    try:
        waveform, sample_rate = torchaudio.load(file_path)
        return waveform, sample_rate
    except Exception as e:
        # Fallback to librosa if torchaudio fails
        try:
            waveform, sample_rate = librosa.load(file_path, sr=None, mono=True)
            waveform = torch.tensor(waveform).unsqueeze(0)
            return waveform, sample_rate
        except Exception as e2:
            raise RuntimeError(f"Failed to load audio file: {e2}")

def normalize_audio(waveform: torch.Tensor) -> torch.Tensor:
    """
    Normalize audio waveform to range [-1, 1].

    Args:
        waveform: Audio waveform tensor

    Returns:
        Normalized waveform
    """
    if torch.max(torch.abs(waveform)) > 0:
        return waveform / torch.max(torch.abs(waveform))
    return waveform

def get_audio_duration(file_path: str) -> float:
    """
    Get the duration of an audio file in seconds.

    Args:
        file_path: Path to audio file

    Returns:
        Duration in seconds
    """
    try:
        waveform, sample_rate = load_audio(file_path)
        duration = waveform.shape[1] / sample_rate
        return duration
    except Exception as e:
        # Fallback to librosa
        try:
            duration = librosa.get_duration(filename=file_path)
            return duration
        except Exception as e2:
            print(f"Failed to get audio duration: {e2}")
            return 0.0

def convert_sample_rate(waveform: torch.Tensor, orig_freq: int, new_freq: int) -> torch.Tensor:
    """
    Convert audio sample rate.

    Args:
        waveform: Audio waveform tensor
        orig_freq: Original sample rate
        new_freq: Target sample rate

    Returns:
        Resampled waveform
    """
    if orig_freq == new_freq:
        return waveform

    resampler = torchaudio.transforms.Resample(orig_freq, new_freq)
    return resampler(waveform)

def convert_to_mono(waveform: torch.Tensor) -> torch.Tensor:
    """
    Convert audio to mono if it's stereo.

    Args:
        waveform: Audio waveform tensor

    Returns:
        Mono waveform
    """
    if waveform.shape[0] > 1:
        return torch.mean(waveform, dim=0, keepdim=True)
    return waveform

def apply_vad(waveform: torch.Tensor, sample_rate: int) -> torch.Tensor:
    """
    Apply Voice Activity Detection to remove silence.

    Args:
        waveform: Audio waveform tensor
        sample_rate: Sample rate

    Returns:
        Waveform with silence removed
    """
    # For now, use a simple energy-based VAD
    # In a production system, you would use a more sophisticated VAD model

    # Convert to numpy for easier processing
    waveform_np = waveform.numpy().flatten()

    # Calculate frame energy
    frame_length = int(sample_rate * 0.025)  # 25ms frames
    hop_length = int(sample_rate * 0.010)    # 10ms hop

    energy = librosa.feature.rms(y=waveform_np, frame_length=frame_length, hop_length=hop_length)[0]

    # Set threshold as percentage of max energy
    threshold = 0.05 * np.max(energy)

    # Create mask for frames above threshold
    mask = energy > threshold

    # Convert frame-level mask to sample-level mask
    sample_mask = np.zeros_like(waveform_np, dtype=bool)

    for i, m in enumerate(mask):
        start = i * hop_length
        end = min(start + frame_length, len(waveform_np))
        if m:
            sample_mask[start:end] = True

    # Apply mask to get active speech segments
    active_speech = waveform_np[sample_mask]

    # Convert back to torch tensor
    return torch.tensor(active_speech).unsqueeze(0)

def segment_audio(waveform: torch.Tensor, sample_rate: int, segment_length_sec: float = 3.0) -> List[torch.Tensor]:
    """
    Segment audio into fixed-length chunks.

    Args:
        waveform: Audio waveform tensor
        sample_rate: Sample rate
        segment_length_sec: Segment length in seconds

    Returns:
        List of audio segments
    """
    # Ensure waveform is mono
    waveform = convert_to_mono(waveform)

    # Calculate segment length in samples
    segment_length = int(segment_length_sec * sample_rate)

    # Flatten waveform for easier processing
    waveform_flat = waveform.squeeze()

    # Calculate number of segments
    num_segments = max(1, int(waveform_flat.shape[0] / segment_length))

    segments = []
    for i in range(num_segments):
        start = i * segment_length
        end = min((i + 1) * segment_length, waveform_flat.shape[0])

        segment = waveform_flat[start:end]

        # Pad if necessary
        if segment.shape[0] < segment_length:
            padding = segment_length - segment.shape[0]
            segment = torch.nn.functional.pad(segment, (0, padding))

        segments.append(segment.unsqueeze(0))

    return segments

def save_audio(waveform: torch.Tensor, sample_rate: int, file_path: str) -> None:
    """
    Save audio waveform to file.

    Args:
        waveform: Audio waveform tensor
        sample_rate: Sample rate
        file_path: Output file path
    """
    try:
        torchaudio.save(file_path, waveform, sample_rate)
    except Exception as e:
        # Fallback to soundfile
        try:
            sf.write(file_path, waveform.squeeze().numpy(), sample_rate)
        except Exception as e2:
            raise RuntimeError(f"Failed to save audio file: {e2}")

def extract_audio_from_video(video_path: str) -> str:
    """
    Extract audio track from video file.

    Args:
        video_path: Path to video file

    Returns:
        Path to extracted audio file
    """
    try:
        import ffmpeg

        # Create temporary file for audio
        temp_file = tempfile.NamedTemporaryFile(suffix='.wav', delete=False)
        audio_path = temp_file.name
        temp_file.close()

        # Extract audio using ffmpeg
        (
            ffmpeg
            .input(video_path)
            .output(audio_path, acodec='pcm_s16le', ac=1, ar='16k')
            .run(quiet=True, overwrite_output=True)
        )

        return audio_path
    except Exception as e:
        raise RuntimeError(f"Failed to extract audio from video: {e}")

def get_speech_rate(waveform: torch.Tensor, sample_rate: int) -> float:
    """
    Estimate speaking rate in words per minute.
    This is a simplified implementation that uses energy peaks
    as a proxy for syllables, then converts to estimated WPM.

    Args:
        waveform: Audio waveform tensor
        sample_rate: Sample rate

    Returns:
        Estimated speech rate in words per minute
    """
    # Typical ratio of syllables to words in English
    SYLLABLE_TO_WORD_RATIO = 1.5

    # Convert to mono and numpy
    waveform = convert_to_mono(waveform)
    waveform_np = waveform.numpy().flatten()

    # Calculate duration in minutes
    duration_minutes = len(waveform_np) / sample_rate / 60

    if duration_minutes <= 0:
        return 0

    # Calculate energy
    energy = librosa.feature.rms(y=waveform_np, frame_length=int(sample_rate * 0.025), hop_length=int(sample_rate * 0.010))[0]

    # Detect peaks in energy as proxy for syllables
    from scipy.signal import find_peaks
    peaks, _ = find_peaks(energy, height=0.1*np.max(energy), distance=int(0.1 * len(energy)))

    # Estimate syllable count
    syllable_count = len(peaks)

    # Convert syllables to words
    estimated_word_count = syllable_count / SYLLABLE_TO_WORD_RATIO

    # Calculate words per minute
    wpm = estimated_word_count / duration_minutes

    # Cap at realistic values
    return min(max(wpm, 50), 200)  # Normal range is about 100-150 WPM