"""
Shared Audio Utilities Module

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

License: BSD 3-Clause
"""

import os
from pathlib import Path
from typing import Union, Optional, Tuple, NamedTuple, List, Dict, Any, Callable
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


class AudioInfo(NamedTuple):
    """Audio file metadata without loading samples into memory."""
    sample_rate: int
    frames: int
    duration: float
    channels: int


def get_audio_info(audio_path: Union[str, Path]) -> AudioInfo:
    """
    Get audio file metadata without loading samples into memory.

    Uses soundfile.info() which reads only the file header.

    Args:
        audio_path: Path to audio file

    Returns:
        AudioInfo with sample_rate, frames, duration, channels
    """
    import soundfile as sf
    info = sf.info(str(audio_path))
    return AudioInfo(
        sample_rate=info.samplerate,
        frames=info.frames,
        duration=info.duration,
        channels=info.channels,
    )


#
# 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)


#
# 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()


#
# SOX BINARY SPECTROGRAM GENERATION
#
def sox_spectrogram(
    audio_array: Union['np.ndarray', 'tf.Tensor', 'torch.Tensor', AudioData],
    shape: Tuple[int, int],
    dest: Optional[Union[str, Path]] = None,
    segment: Optional[int] = None,
    segment_duration: Optional[float] = None,
    segment_overlap: Optional[float] = None,
    sample_rate: Optional[int] = None,
    create_parents: bool = True,
    overwrite: bool = True,
    backend: str = 'librosa'
) -> Union[str, np.ndarray]:
    """
    Generate spectrogram using the actual sox binary (exact match to Shiny_PNW-Cnet).

    This function saves audio to a temporary WAV file and calls sox to generate
    a spectrogram, ensuring exact compatibility with models trained on sox-generated
    spectrograms.

    Args:
        audio_array: Audio data (numpy array, tf.Tensor, torch.Tensor, or AudioData)
        shape: Output shape of spectrogram (height, width) - must be (257, 1000) for PNW-Cnet
        dest: Optional file path to save spectrogram as PNG (required)
        segment: Segment number for extracting specific audio portion
        segment_duration: Duration of segment in seconds
        sample_rate: Audio sample rate in Hz (required for tensor/array inputs)
        create_parents: Create parent directories if needed
        overwrite: Overwrite existing file if it exists
        backend: Backend type ('tensorflow', 'pytorch', 'librosa') for type conversion

    Returns:
        Path to saved PNG file (str)

    Raises:
        RuntimeError: If sox binary is not installed or not in PATH
        ValueError: If dest is None (sox requires output file)
        FileExistsError: If file exists and overwrite=False

    Notes:
        - Requires sox binary to be installed: `apt-get install sox` or `brew install sox`
        - Uses exact Shiny_PNW-Cnet parameters: -x 1000 -y 257 -z 90 -m -r
        - Only works for 12-second segments at 8kHz sampling rate (PNW-Cnet standard)
        - Audio is resampled to 8kHz and converted to mono automatically
    """
    import subprocess
    import tempfile
    import shutil
    import soundfile as sf

    # Check if sox is installed
    if shutil.which('sox') is None:
        raise RuntimeError(
            "sox binary not found in PATH. Please install sox:\n"
            "  Ubuntu/Debian: sudo apt-get install sox\n"
            "  macOS: brew install sox\n"
            "  For other systems, see: http://sox.sourceforge.net/"
        )

    # dest is required for sox (it generates a PNG file)
    if dest is None:
        raise ValueError("dest parameter is required when using sox binary (cannot return array)")

    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)

    # Extract AudioData if provided
    if isinstance(audio_array, AudioData):
        y = audio_array.tensor
        if sample_rate is None:
            sample_rate = audio_array.sample_rate
    else:
        y = audio_array

    if sample_rate is None:
        sample_rate = DEFAULT_SAMPLE_RATE

    # Convert tensor to numpy array based on backend
    if backend == 'tensorflow':
        try:
            import tensorflow as tf
            if isinstance(y, tf.Tensor):
                y = y.numpy()
        except ImportError:
            pass
    elif backend == 'pytorch':
        try:
            import torch
            if isinstance(y, torch.Tensor):
                y = y.cpu().numpy()
        except ImportError:
            pass

    # Now y should be a numpy array
    if not isinstance(y, np.ndarray):
        raise ValueError(f"Could not convert audio to numpy array (type: {type(y)})")

    # Extract segment if specified
    if segment is not None:
        start_idx, end_idx = calculate_segment_indices(
            len(y), sample_rate, segment, segment_duration, segment_overlap
        )
        y = y[start_idx:end_idx]

    # Create temporary WAV file
    with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as tmp_wav:
        tmp_wav_path = tmp_wav.name

    try:
        # Write audio to temporary WAV file
        # sox expects 16-bit PCM WAV files
        sf.write(tmp_wav_path, y, sample_rate, subtype='PCM_16')

        # Build sox command (matching Shiny_PNW-Cnet exactly)
        # sox -V1 "input.wav" -n trim 0 12.000 remix 1 rate 8k spectrogram -x 1000 -y 257 -z 90 -m -r -o "output.png"

        # Calculate duration (default 12 seconds for PNW-Cnet)
        if segment_duration is None:
            segment_duration = DEFAULT_SEGMENT_DURATION

        duration = min(segment_duration, len(y) / sample_rate)

        # Validate shape matches PNW-Cnet expectations
        target_height, target_width = shape
        if (target_height, target_width) != (257, 1000):
            raise ValueError(
                f"sox binary currently only supports (257, 1000) shape for PNW-Cnet compatibility. "
                f"Got: {shape}. Use 'pysox' spectrogram_type for other shapes."
            )

        sox_cmd = [
            'sox',
            '-V1',  # Verbosity level 1 (errors only)
            tmp_wav_path,
            '-n',  # Null output (no audio output, only spectrogram)
            'trim', '0', f'{duration:.3f}',  # Extract from start
            'remix', '1',  # Convert to mono (channel 1)
            'rate', '8k',  # Resample to 8kHz
            'spectrogram',
            '-x', str(target_width),  # 1000 pixels wide
            '-y', str(target_height),  # 257 pixels tall
            '-z', '90',  # 90 dB dynamic range
            '-m',  # Monochrome
            '-r',  # Raw (no axes)
            '-o', str(dest_path.absolute())
        ]

        # Run sox command
        result = subprocess.run(
            sox_cmd,
            capture_output=True,
            text=True,
            check=False
        )

        if result.returncode != 0:
            raise RuntimeError(
                f"sox command failed with return code {result.returncode}\n"
                f"Command: {' '.join(sox_cmd)}\n"
                f"stderr: {result.stderr}\n"
                f"stdout: {result.stdout}"
            )

    finally:
        # Clean up temporary WAV file
        if os.path.exists(tmp_wav_path):
            os.remove(tmp_wav_path)

    return str(dest_path)


#
# SPECTROGRAM SAVING - TFRECORD
#
class TFRecordWriter:
    """
    Context manager for streaming TFRecord writes.

    Writes spectrograms incrementally to avoid memory issues with large audio files.
    Supports optional chunking to split into multiple TFRecord files.

    Example:
        with TFRecordWriter('output.tfrecord', compression='GZIP') as writer:
            for i, spec in enumerate(spectrograms):
                writer.write_spectrogram(spec, i, 'audio_file')
    """

    def __init__(
        self,
        output_path: Union[str, Path],
        compression: Optional[str] = 'GZIP',
        chunk_size: Optional[int] = None,
        basename: str = 'spectrograms'
    ):
        """
        Initialize TFRecord writer.

        Args:
            output_path: Path to output TFRecord file (or directory if chunking)
            compression: Compression type ('GZIP', 'ZLIB', or None)
            chunk_size: Number of spectrograms per TFRecord file (None = single file)
            basename: Base name for chunked files (e.g., 'audio_001.tfrecord')
        """
        self.output_path = Path(output_path)
        self.compression = compression
        self.chunk_size = chunk_size
        self.basename = basename
        self.tf = None
        self.writer = None
        self.current_chunk = 0
        self.chunk_count = 0
        self.written_files = []

    def __enter__(self):
        # Import tensorflow only when needed
        try:
            import tensorflow as tf
            self.tf = tf
        except ImportError:
            raise ImportError(
                "TensorFlow required for TFRecord output. "
                "Install with: pip install tensorflow"
            )

        # Create parent directories
        if self.chunk_size is None:
            # Single file mode
            self.output_path.parent.mkdir(parents=True, exist_ok=True)
            self._open_writer(self.output_path)
        else:
            # Chunked mode - output_path is directory
            self.output_path.mkdir(parents=True, exist_ok=True)

        return self

    def _open_writer(self, path: Path):
        """Open a new TFRecord writer."""
        options = None
        if self.compression:
            options = self.tf.io.TFRecordOptions(compression_type=self.compression)

        self.writer = self.tf.io.TFRecordWriter(str(path), options=options)
        self.written_files.append(str(path))

    def _close_writer(self):
        """Close current TFRecord writer."""
        if self.writer:
            self.writer.close()
            self.writer = None

    def _maybe_rotate_chunk(self):
        """Check if we need to start a new chunk file."""
        if self.chunk_size is None:
            return  # Single file mode

        if self.chunk_count >= self.chunk_size:
            # Close current writer and open new one
            self._close_writer()
            self.current_chunk += 1
            self.chunk_count = 0

            # Generate new filename
            chunk_filename = f"{self.basename}_chunk{self.current_chunk:04d}.tfrecord"
            chunk_path = self.output_path / chunk_filename
            self._open_writer(chunk_path)

    def write_spectrogram(
        self,
        spectrogram: np.ndarray,
        segment_index: int,
        filename: str,
        metadata: Optional[Dict[str, Any]] = None
    ):
        """
        Write a single spectrogram to TFRecord.

        Args:
            spectrogram: Spectrogram array (uint8, shape: [height, width])
            segment_index: Segment index in the audio file
            filename: Original audio filename (for tracking)
            metadata: Optional metadata dict (e.g., {'audio_uri': 's3://...', 'start_time': 500})
        """
        # Check if we need to open first writer (chunked mode)
        if self.chunk_size is not None and self.writer is None:
            chunk_filename = f"{self.basename}_chunk{self.current_chunk:04d}.tfrecord"
            chunk_path = self.output_path / chunk_filename
            self._open_writer(chunk_path)

        # Check if we need to rotate to new chunk
        self._maybe_rotate_chunk()

        height, width = spectrogram.shape

        # Build feature dictionary
        feature = {
            'spectrogram': self.tf.train.Feature(
                bytes_list=self.tf.train.BytesList(value=[spectrogram.tobytes()])
            ),
            'height': self.tf.train.Feature(
                int64_list=self.tf.train.Int64List(value=[height])
            ),
            'width': self.tf.train.Feature(
                int64_list=self.tf.train.Int64List(value=[width])
            ),
            'segment_index': self.tf.train.Feature(
                int64_list=self.tf.train.Int64List(value=[segment_index])
            ),
            'filename': self.tf.train.Feature(
                bytes_list=self.tf.train.BytesList(value=[filename.encode('utf-8')])
            ),
        }

        # Add metadata fields if provided
        if metadata:
            for key, value in metadata.items():
                if isinstance(value, (int, bool)):
                    feature[key] = self.tf.train.Feature(
                        int64_list=self.tf.train.Int64List(value=[int(value)])
                    )
                elif isinstance(value, float):
                    feature[key] = self.tf.train.Feature(
                        float_list=self.tf.train.FloatList(value=[value])
                    )
                elif isinstance(value, str):
                    feature[key] = self.tf.train.Feature(
                        bytes_list=self.tf.train.BytesList(value=[value.encode('utf-8')])
                    )
                elif value is None:
                    # Skip None values
                    continue
                else:
                    # Convert to string for unsupported types
                    feature[key] = self.tf.train.Feature(
                        bytes_list=self.tf.train.BytesList(value=[str(value).encode('utf-8')])
                    )

        # Create and write example
        example = self.tf.train.Example(features=self.tf.train.Features(feature=feature))
        self.writer.write(example.SerializeToString())
        self.chunk_count += 1

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._close_writer()

    def get_written_files(self) -> List[str]:
        """Get list of all TFRecord files written."""
        return self.written_files


#
# SPECTROGRAM GENERATION - UNIFIED INTERFACE
#
def generate_spectrograms(
    spec_func: Callable,
    audio_data: Any,
    num_segments: int,
    spectrogram_shape: Tuple[int, int],
    dest_folder: str,
    basename: str,
    directory: Optional[str] = None,
    segment_duration: Optional[float] = None,
    segment_overlap: Optional[float] = None,
    sample_rate: Optional[int] = None,
    output_format: str = 'png',
    tfrecord_compression: Optional[str] = 'GZIP',
    tfrecord_chunk_size: Optional[int] = None,
    metadata: Optional[Dict[str, Any]] = None,
    create_parents: bool = True,
    overwrite: bool = True,
    audio_path: Optional[Union[str, Path]] = None,
    native_sample_rate: Optional[int] = None,
    total_frames: Optional[int] = None,
) -> Union[List[Union[Any, str]], List[str]]:
    """
    Unified spectrogram generation interface supporting both PNG and TFRecord output.

    This function eliminates code duplication across TensorFlow, PyTorch, and librosa backends
    by providing a common interface for generating spectrograms with configurable output format.

    Supports two modes:
    - Legacy mode: audio_data contains the full audio array (all segments sliced from memory)
    - Streaming mode: audio_path is set, audio_data is None. Each segment is loaded from disk
      on demand via load_audio_segment(), keeping memory usage constant regardless of file length.

    Args:
        spec_func: Backend-specific spectrogram function (sox_spectrogram, mel_spectrogram, etc.)
                   Must accept: (audio_data, shape, dest, segment, segment_duration,
                                segment_overlap, sample_rate, create_parents, overwrite)
        audio_data: Backend-specific audio data (tf.Tensor, torch.Tensor, np.ndarray, etc.)
                    Can be None when audio_path is provided (streaming mode).
        num_segments: Total number of segments to generate
        spectrogram_shape: Output shape (height, width)
        dest_folder: Output folder for spectrograms
        basename: Base name for output files
        directory: Optional subdirectory within dest_folder
        segment_duration: Duration of each segment in seconds
        segment_overlap: Overlap between segments in seconds
        sample_rate: Audio sample rate in Hz (target sample rate in streaming mode)
        output_format: 'png' or 'tfrecord'
        tfrecord_compression: TFRecord compression type ('GZIP', 'ZLIB', None)
        tfrecord_chunk_size: Split TFRecord into chunks of this size (None = single file)
        metadata: Additional metadata for TFRecord (e.g., audio_uri, start_time)
        create_parents: Create parent directories if needed
        overwrite: Overwrite existing files
        audio_path: Path to audio file for streaming mode (None = legacy mode)
        native_sample_rate: Original sample rate of audio file (required if audio_path is set)
        total_frames: Total frames in audio file (required if audio_path is set)

    Returns:
        If output_format='png': List of spectrogram data or file paths (backend-dependent)
        If output_format='tfrecord': List of TFRecord file paths (always List[str])

    Example:
        # TensorFlow backend
        specs = generate_spectrograms(
            spec_func=sox_spectrogram,
            audio_data=audio_tensor,
            num_segments=100,
            spectrogram_shape=(257, 1000),
            dest_folder='output',
            basename='audio_file',
            output_format='tfrecord'
        )
    """
    from soundhub_utils import names

    streaming = audio_path is not None

    def _get_audio_and_segment(segment_idx):
        """Return (audio_data, segment_index) for spec_func call."""
        if streaming:
            chunk = load_audio_segment(
                audio_path=audio_path,
                segment=segment_idx,
                native_sample_rate=native_sample_rate,
                target_sample_rate=sample_rate,
                segment_duration=segment_duration,
                segment_overlap=segment_overlap,
                total_frames=total_frames,
            )
            # segment=None tells spec_func to use the entire chunk
            return chunk, None
        else:
            return audio_data, segment_idx

    # TFRecord output
    if output_format == 'tfrecord':
        # Build TFRecord path
        if directory:
            tfrecord_path = Path(dest_folder) / directory / f"{basename}.tfrecord"
        else:
            tfrecord_path = Path(dest_folder) / f"{basename}.tfrecord"

        # Use TFRecord writer for streaming writes
        with TFRecordWriter(
            output_path=tfrecord_path,
            compression=tfrecord_compression,
            chunk_size=tfrecord_chunk_size,
            basename=basename
        ) as writer:
            for segment_idx in range(num_segments):
                audio_input, seg_idx = _get_audio_and_segment(segment_idx)

                # Generate spectrogram without saving to PNG
                spec = spec_func(
                    audio_input,
                    spectrogram_shape,
                    dest=None,  # Don't save PNG
                    segment=seg_idx,
                    segment_duration=segment_duration,
                    segment_overlap=segment_overlap,
                    sample_rate=sample_rate,
                    create_parents=create_parents,
                    overwrite=overwrite
                )

                # Convert to numpy for TFRecord
                spec_np = _to_numpy(spec)
                writer.write_spectrogram(spec_np, segment_idx, basename, metadata)

        # Return list of written TFRecord files
        return writer.get_written_files()

    # PNG output (original behavior)
    spectrograms_list = []
    for segment_idx in range(num_segments):
        # Generate filename using names module
        if dest_folder:
            spec_filename = names.filename(
                basename,
                directory=directory,
                segment=segment_idx,
                ext='png'
            )
            spec_path = Path(dest_folder) / spec_filename
        else:
            spec_path = None

        audio_input, seg_idx = _get_audio_and_segment(segment_idx)

        # Generate spectrogram
        spec = spec_func(
            audio_input,
            spectrogram_shape,
            dest=spec_path,
            segment=seg_idx,
            segment_duration=segment_duration,
            segment_overlap=segment_overlap,
            sample_rate=sample_rate,
            create_parents=create_parents,
            overwrite=overwrite
        )
        spectrograms_list.append(spec)

    return spectrograms_list


def _to_numpy(spec: Any) -> np.ndarray:
    """
    Convert spectrogram to numpy array (handles TensorFlow Tensor, PyTorch Tensor, numpy array).

    Args:
        spec: Spectrogram as tf.Tensor, torch.Tensor, or np.ndarray

    Returns:
        Numpy array
    """
    # Check if it's already numpy
    if isinstance(spec, np.ndarray):
        return spec

    # Try TensorFlow
    try:
        import tensorflow as tf
        if isinstance(spec, tf.Tensor):
            return spec.numpy()
    except ImportError:
        pass

    # Try PyTorch
    try:
        import torch
        if isinstance(spec, torch.Tensor):
            return spec.cpu().numpy()
    except ImportError:
        pass

    # Fallback: assume it has a numpy() method
    if hasattr(spec, 'numpy'):
        return spec.numpy()

    raise TypeError(f"Cannot convert {type(spec)} to numpy array")


#
# INFERENCE UTILITIES
#
def load_spectrograms_tf(file_paths: List[str], spectrogram_shape: Tuple[int, int], tf) -> 'tf.Tensor':
    """
    Load spectrogram images using TensorFlow for model compatibility.

    Args:
        file_paths: List of paths to spectrogram PNG files
        spectrogram_shape: Target shape (height, width)
        tf: TensorFlow module

    Returns:
        TensorFlow tensor with shape (batch_size, height, width, channels)
    """
    def load_and_preprocess_image(path):
        image = tf.io.read_file(path)
        image = tf.image.decode_png(image, channels=1)
        image = tf.cast(image, tf.float32) / 255.0

        # Don't resize - images should already be correct size
        # Resizing causes interpolation artifacts that affect predictions
        return image

    # Create dataset from file paths and load images
    path_ds = tf.data.Dataset.from_tensor_slices(file_paths)
    image_ds = path_ds.map(load_and_preprocess_image, num_parallel_calls=tf.data.AUTOTUNE)

    # Batch all images
    # batch_ds = image_ds.batch(len(file_paths))
    batch_ds = image_ds.batch(100)

    return batch_ds
    # # Get single batch
    # for batch in batch_ds:
    #     return batch


def load_spectrograms_from_tfrecord(
    tfrecord_paths: Union[str, List[str]],
    spectrogram_shape: Tuple[int, int],
    batch_size: int = 100,
    compression: Optional[str] = 'GZIP'
):
    """
    Load spectrograms from TFRecord file(s) for inference.

    Args:
        tfrecord_paths: Path to TFRecord file or list of TFRecord files (for chunked mode)
        spectrogram_shape: Expected spectrogram shape (height, width)
        batch_size: Batch size for dataset
        compression: Compression type used ('GZIP', 'ZLIB', or None)

    Returns:
        TensorFlow dataset yielding batches of (spectrograms, metadata_dict)
        - spectrograms: shape (batch_size, height, width, 1), float32, range [0, 1]
        - metadata_dict: dict with 'segment_index', 'filename', and any custom metadata

    Example:
        dataset = load_spectrograms_from_tfrecord('audio.tfrecord', (257, 1000))
        for spectrograms, metadata in dataset:
            predictions = model.predict(spectrograms)
    """
    import tensorflow as tf

    # Ensure tfrecord_paths is a list
    if isinstance(tfrecord_paths, str):
        tfrecord_paths = [tfrecord_paths]

    def parse_example(example_proto):
        """Parse a single TFRecord example."""
        # Define base feature description
        feature_description = {
            'spectrogram': tf.io.FixedLenFeature([], tf.string),
            'height': tf.io.FixedLenFeature([], tf.int64),
            'width': tf.io.FixedLenFeature([], tf.int64),
            'segment_index': tf.io.FixedLenFeature([], tf.int64),
            'filename': tf.io.FixedLenFeature([], tf.string),
        }

        parsed = tf.io.parse_single_example(example_proto, feature_description)

        # Decode spectrogram
        spectrogram = tf.io.decode_raw(parsed['spectrogram'], tf.uint8)
        spectrogram = tf.reshape(spectrogram, [spectrogram_shape[0], spectrogram_shape[1]])
        spectrogram = tf.cast(spectrogram, tf.float32) / 255.0
        spectrogram = tf.expand_dims(spectrogram, -1)  # Add channel dimension

        # Build metadata dictionary
        metadata = {
            'segment_index': parsed['segment_index'],
            'filename': parsed['filename'],
            'height': parsed['height'],
            'width': parsed['width'],
        }

        return spectrogram, metadata

    # Create dataset from TFRecord file(s)
    dataset = tf.data.TFRecordDataset(
        tfrecord_paths,
        compression_type=compression if compression else ''
    )
    dataset = dataset.map(parse_example, num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.batch(batch_size)
    dataset = dataset.prefetch(tf.data.AUTOTUNE)  # Prefetch for better performance

    return dataset


def is_tfrecord_path(path: Union[str, Path, List[str], List[Path]]) -> bool:
    """
    Check if path(s) point to TFRecord file(s).

    Args:
        path: File path, list of paths, or directory path

    Returns:
        True if path(s) are TFRecord files, False otherwise
    """
    # Handle list of paths
    if isinstance(path, list):
        if len(path) == 0:
            return False
        # Check first item
        path = path[0]

    path = Path(path)

    # Check extension
    return path.suffix.lower() == '.tfrecord' or path.name.endswith('.tfrecord')


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


#
# SEGMENTATION UTILITIES
#
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)))
