"""
Shared Audio Utilities Module

Shared utilities and helper functions for audio processing modules (librosa, pytorch, tensorflow).
Contains functionality that doesn't depend on specific backend implementations.

License: BSD 3-Clause
"""

import os
from pathlib import Path
from typing import Union, Optional, Tuple, NamedTuple, List, Dict, Any
import numpy as np
import matplotlib.pyplot as plt
from cocina.config_handler import ConfigHandler

ch = ConfigHandler(__file__)


#
# DEFAULTS
#
DEFAULT_SAMPLE_RATE: int = ch.get('DEFAULT_SAMPLE_RATE', 8000)
DEFAULT_SEGMENT_DURATION: int = ch.get('DEFAULT_SEGMENT_DURATION', 12.0)
DEFAULT_SPECTROGRAM_SHAPE: Tuple[int, int] = ch.get('DEFAULT_SPECTROGRAM_SHAPE', (257, 1000))


#
# NAMED-TUPLE
#
class AudioData(NamedTuple):
    """Audio data container with array and metadata."""
    tensor: Union['np.ndarray', 'tf.Tensor', 'torch.Tensor']
    sample_rate: int
    duration: float
    channels: int




#
# SEGMENT CALCULATION
#
def calculate_segment_indices(
    audio_length: int,
    sample_rate: int,
    segment: Optional[int] = None,
    segment_duration: Optional[float] = None,
    segment_overlap: Optional[float] = None
) -> Tuple[int, int]:
    """
    Calculate start and end indices for audio segment extraction.

    Args:
        audio_length: Total length of audio in samples
        sample_rate: Audio sample rate in Hz
        segment: Segment number (0-based)
        segment_duration: Duration of each segment in seconds
        segment_overlap: Overlap between segments in seconds

    Returns:
        Tuple of (start_index, end_index)
    """
    if segment is None:
        return 0, audio_length

    if segment_duration is None:
        segment_duration = DEFAULT_SEGMENT_DURATION
    if segment_overlap is None:
        segment_overlap = 0.0

    segment_samples = int(segment_duration * sample_rate)
    overlap_samples = int(segment_overlap * sample_rate)
    step_samples = segment_samples - overlap_samples

    start_index = segment * step_samples
    end_index = min(start_index + segment_samples, audio_length)

    return start_index, end_index


# Resampling padding: extra samples on each side to avoid edge artifacts
_RESAMPLE_PAD_SECONDS = 0.1  # 100ms padding on each boundary


def load_audio_segment(
    audio_path: Union[str, Path],
    segment: int,
    native_sample_rate: int,
    target_sample_rate: int,
    segment_duration: Optional[float] = None,
    segment_overlap: Optional[float] = None,
    total_frames: Optional[int] = None,
) -> np.ndarray:
    """
    Load a single audio segment from a file without loading the entire file.

    Uses soundfile for random-access reading, then resamples the small chunk
    to the target sample rate if needed.

    Args:
        audio_path: Path to audio file
        segment: Segment index (0-based)
        native_sample_rate: Original sample rate of the file
        target_sample_rate: Desired output sample rate
        segment_duration: Duration of each segment in seconds
        segment_overlap: Overlap between segments in seconds
        total_frames: Total number of frames in the file

    Returns:
        Numpy float32 array of audio samples at target_sample_rate, mono
    """
    import soundfile as sf

    if segment_duration is None:
        segment_duration = DEFAULT_SEGMENT_DURATION
    if segment_overlap is None:
        segment_overlap = 0.0

    # Calculate segment boundaries in target sample rate space
    target_segment_samples = int(segment_duration * target_sample_rate)
    target_overlap_samples = int(segment_overlap * target_sample_rate)
    target_step_samples = target_segment_samples - target_overlap_samples

    target_start = segment * target_step_samples
    target_end = target_start + target_segment_samples

    # Convert to native sample rate frame positions
    needs_resample = (native_sample_rate != target_sample_rate)
    rate_ratio = native_sample_rate / target_sample_rate

    native_start = int(target_start * rate_ratio)
    native_end = int(target_end * rate_ratio)

    # Add resampling padding to avoid edge artifacts
    if needs_resample:
        pad_samples = int(_RESAMPLE_PAD_SECONDS * native_sample_rate)
    else:
        pad_samples = 0

    padded_native_start = max(0, native_start - pad_samples)
    padded_native_end = native_end + pad_samples
    if total_frames is not None:
        padded_native_end = min(padded_native_end, total_frames)

    # Read only the needed frames from disk
    num_frames_to_read = padded_native_end - padded_native_start

    with sf.SoundFile(str(audio_path), 'r') as f:
        f.seek(padded_native_start)
        audio_chunk = f.read(frames=num_frames_to_read, dtype='float32', always_2d=False)

    # Convert to mono if stereo
    if audio_chunk.ndim > 1:
        audio_chunk = np.mean(audio_chunk, axis=1)

    # Resample if needed
    if needs_resample:
        import librosa
        audio_chunk = librosa.resample(
            audio_chunk,
            orig_sr=native_sample_rate,
            target_sr=target_sample_rate,
        )

        # Trim the resampling padding
        left_pad_native = native_start - padded_native_start
        right_pad_native = padded_native_end - native_end
        left_pad_target = int(left_pad_native / rate_ratio)
        right_pad_target = int(right_pad_native / rate_ratio)

        if right_pad_target > 0:
            audio_chunk = audio_chunk[left_pad_target:-right_pad_target]
        else:
            audio_chunk = audio_chunk[left_pad_target:]

    return audio_chunk.astype(np.float32)


def calculate_num_segments(
    audio_length_seconds: float,
    segment_duration: float,
    segment_overlap: float = 0.0
) -> int:
    """
    Calculate the number of segments for a given audio length.

    Args:
        audio_length_seconds: Total audio length in seconds
        segment_duration: Duration of each segment in seconds
        segment_overlap: Overlap between segments in seconds

    Returns:
        Number of segments
    """
    step_duration = segment_duration - segment_overlap
    return max(1, int(np.ceil((audio_length_seconds - segment_overlap) / step_duration)))


#
# SPECTROGRAM SAVING - PNG
#
def save_spectrogram_png(
    spectrogram_array: np.ndarray,
    dest: Union[str, Path],
    create_parents: bool = True,
    overwrite: bool = True,
    colormap: Optional[str] = 'viridis'
) -> None:
    """
    Save spectrogram array as PNG file.

    Args:
        spectrogram_array: Spectrogram data as numpy array (uint8 or float)
        dest: Destination file path
        create_parents: Create parent directories if needed
        overwrite: Overwrite existing file if it exists
        colormap: Matplotlib colormap name ('gray', 'viridis', etc.) or None for raw grayscale

    Raises:
        FileExistsError: If file exists and overwrite=False
    """
    dest_path = Path(dest)

    # Check if file exists and handle overwrite
    if dest_path.exists() and not overwrite:
        raise FileExistsError(f"Destination file {dest_path} already exists and overwrite=False")

    # Create parent directories if needed
    if create_parents:
        dest_path.parent.mkdir(parents=True, exist_ok=True)

    if colormap is None:
        # Save as raw grayscale PNG (no colorbar, no axes) - for Sox-matching
        from PIL import Image
        img = Image.fromarray(spectrogram_array, mode='L')
        img.save(dest_path)
    else:
        # Save with matplotlib (with colorbar and axes) - for visualization
        plt.figure(figsize=(10, 6))
        plt.imshow(spectrogram_array, aspect='auto', origin='lower', cmap=colormap)
        plt.colorbar()
        plt.title('Spectrogram')
        plt.xlabel('Time')
        plt.ylabel('Frequency')
        plt.tight_layout()
        plt.savefig(dest_path, dpi=150, bbox_inches='tight')
        plt.close()


#
# INFERENCE UTILITIES
#


def process_inference_results(
    predictions: np.ndarray,
    file_paths: List[str],
    confidence_threshold: float,
    top_k: int,
    class_names: Optional[List[str]]
) -> List[Dict[str, Any]]:
    """
    Process model predictions into structured results.

    Args:
        predictions: Model prediction outputs (batch_size, num_classes)
        file_paths: List of input file paths
        confidence_threshold: Minimum confidence for predictions
        top_k: Number of top predictions to include
        class_names: List of class names (if None, uses indices)

    Returns:
        List of result dictionaries, one per input file
    """
    results = []

    for i, (prediction, file_path) in enumerate(zip(predictions, file_paths)):
        # Get top-k predictions
        top_indices = np.argsort(prediction)[-top_k:][::-1]
        top_confidences = prediction[top_indices]

        # Filter by confidence threshold
        valid_mask = top_confidences >= confidence_threshold
        filtered_indices = top_indices[valid_mask]
        filtered_confidences = top_confidences[valid_mask]

        # Create predictions list
        prediction_list = []
        for idx, conf in zip(filtered_indices, filtered_confidences):
            pred_name = class_names[idx] if class_names else f"class_{idx}"
            prediction_list.append({
                'class': pred_name,
                'confidence': float(conf),
                'class_index': int(idx)
            })

        # Create result entry
        result = {
            'file_path': file_path,
            'file_name': os.path.basename(file_path),
            'predictions': prediction_list,
            'top_prediction': prediction_list[0] if prediction_list else None,
            'max_confidence': float(np.max(prediction)),
            'num_valid_predictions': len(prediction_list)
        }

        results.append(result)

    return results