"""
Unified Audio Processing Interface

Routes to backend-specific implementations (TensorFlow, PyTorch, or librosa)

License: BSD 3-Clause
"""

from pathlib import Path
from typing import Union, Optional, Tuple, List, Dict, Any, Callable, NamedTuple
import pandas as pd
from cocina.config_handler import ConfigHandler
from soundhub_utils.audio._shared import AudioData
from soundhub_utils.utils import names
from soundhub_utils.audio import _shared

ch = ConfigHandler(__file__)


#
# AUDIO FILE INFO
#
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,
    )


#
# AUDIO LOADING
#
def load_audio(
    src: Union[str, Path],
    sample_rate: Optional[int] = None,
    file_type: Optional[str] = None,
    backbone: Optional[str] = None
):
    """
    Load audio file.

    Args:
        src: Path to audio file (FLAC, WAV, etc.)
        sample_rate: Target sample rate in Hz. If None, uses original sample rate
        file_type: Audio file type. If None, auto-detected from extension
        backbone: Backend to use ('tensorflow', 'pytorch', or 'librosa')

    Returns:
        AudioData containing tensor/array, sample_rate, duration, and channels

    Raises:
        IOError: If file cannot be loaded
        ValueError: If invalid file type, sample rate, or backbone specified
    """
    backend = _get_backend(backbone)
    return backend.load_audio(src, sample_rate, file_type)


def spectrograms(
    src: Union[str, Path],
    dest_folder: str,
    spectrogram_shape: Tuple[int, int],
    spectrogram_func: Callable,
    basename: Optional[str] = None,
    directory: Optional[str] = None,
    segment_duration: Optional[float] = None,
    segment_overlap: Optional[float] = None,
    sample_rate: Optional[int] = None,
    create_parents: bool = True,
    overwrite: bool = True,
    output_format: Optional[str] = None,
    tfrecord_compression: Optional[str] = None,
    tfrecord_chunk_size: Optional[int] = None,
    metadata: Optional[Dict[str, Any]] = None
) -> Union[List[str], List[Any]]:
    """
    Generate multiple spectrograms from audio source with segmentation.

    This is the main interface that takes any spectrogram generation function
    and handles all the common logic: file loading, segmentation, naming,
    and output format handling (PNG/TFRecord).

    Args:
        src: Audio source - file path
        dest_folder: Folder to save spectrogram files
        spectrogram_shape: Output shape for each spectrogram (height, width)
        spectrogram_func: Function that generates spectrograms. Must accept:
            (audio_data, shape, dest, segment, segment_duration, segment_overlap,
             sample_rate, create_parents, overwrite) and return spectrogram data.
            Examples: tensorflow_audio.spectrogram, pytorch_audio.spectrogram,
                     custom functions, etc.
        basename: Base name for output files (auto-determined if None)
        directory: Directory name within dest_folder (optional)
        segment_duration: Duration of each segment in seconds
        segment_overlap: Overlap between segments in seconds
        sample_rate: Audio sample rate in Hz
        create_parents: Create parent directories if needed
        overwrite: Overwrite existing files if they exist
        output_format: Output format - 'png' or 'tfrecord' (default from config)
        tfrecord_compression: TFRecord compression - 'GZIP', 'ZLIB', or None
        tfrecord_chunk_size: Split TFRecord into chunks of this size (None = single file)
        metadata: Additional metadata to store in TFRecord (e.g., audio_uri, start_time)

    Returns:
        If output_format='png': List of spectrogram file paths (or tensors/arrays)
        If output_format='tfrecord': List of TFRecord file paths

    Raises:
        ValueError: If invalid parameters specified
        IOError: If audio loading or file operations fail
        FileExistsError: If files exist and overwrite=False

    Example:
        # Using TensorFlow mel spectrogram
        from soundhub_utils.audio import tensorflow_audio
        spectrograms('audio.wav', 'output/', (128, 1000),
                    tensorflow_audio.mel_spectrogram)

        # Using custom spectrogram function
        def my_spectrogram(audio, shape, **kwargs):
            # Custom implementation
            return my_stft(audio)
        spectrograms('audio.wav', 'output/', (257, 1000), my_spectrogram)
    """
    # Get configuration defaults
    if output_format is None:
        output_format = ch.get('TFRECORD_OUTPUT_FORMAT', 'png')
    if tfrecord_compression is None:
        tfrecord_compression = ch.get('TFRECORD_COMPRESSION', 'GZIP')
    if tfrecord_chunk_size is None:
        tfrecord_chunk_size = ch.get('TFRECORD_CHUNK_SIZE', None)
    if segment_duration is None:
        segment_duration = _shared.DEFAULT_SEGMENT_DURATION
    if segment_overlap is None:
        segment_overlap = 0.0

    # Validate output format
    if output_format not in ['png', 'tfrecord']:
        raise ValueError(f"output_format must be 'png' or 'tfrecord', got: {output_format}")

    # Handle source input - currently only file paths supported in main interface
    if not isinstance(src, (str, Path)):
        raise ValueError(f"Main spectrograms interface only supports file paths, got: {type(src)}")

    # Get audio file info without loading full file
    audio_info = get_audio_info(src)
    actual_sample_rate = sample_rate if sample_rate is not None else audio_info.sample_rate

    if basename is None:
        basename = Path(src).stem

    # Calculate number of segments
    num_segments = _shared.calculate_num_segments(
        audio_info.duration,
        segment_duration,
        segment_overlap
    )

    # Use internal spectrogram generation with streaming mode
    return _generate_spectrograms(
        spec_func=spectrogram_func,
        audio_data=None,  # Use streaming mode
        num_segments=num_segments,
        spectrogram_shape=spectrogram_shape,
        dest_folder=dest_folder,
        basename=basename,
        directory=directory,
        segment_duration=segment_duration,
        segment_overlap=segment_overlap,
        sample_rate=actual_sample_rate,
        output_format=output_format,
        tfrecord_compression=tfrecord_compression,
        tfrecord_chunk_size=tfrecord_chunk_size,
        metadata=metadata,
        create_parents=create_parents,
        overwrite=overwrite,
        audio_path=src,
        native_sample_rate=audio_info.sample_rate,
        total_frames=audio_info.frames,
    )


def inference(
    model: str,
    files: Union[List[str], pd.DataFrame],
    spectrogram_shape: Optional[Tuple[int, int]] = None,
    confidence_threshold: float = 0.1,
    top_k: int = 5,
    class_names: Optional[List[str]] = None,
    return_format: str = 'dict',
    backbone: Optional[str] = None
) -> Union[Dict[str, Any], pd.DataFrame]:
    """
    Run model inference on spectrogram files.

    Args:
        model: Path to model file
        files: List of spectrogram file paths or DataFrame with file paths
        spectrogram_shape: Expected spectrogram shape (height, width)
        confidence_threshold: Minimum confidence for predictions
        top_k: Number of top predictions to return per file
        class_names: List of class names (if None, uses indices)
        return_format: Output format ('dict' or 'dataframe')
        backbone: Backend to use ('tensorflow', 'pytorch', or 'librosa')

    Returns:
        Dictionary or DataFrame containing inference results

    Raises:
        FileNotFoundError: If model file not found
        ValueError: If invalid parameters or backbone specified
        IOError: If inference fails
    """
    backend = _get_backend(backbone)
    return backend.inference(
        model=model,
        files=files,
        spectrogram_shape=spectrogram_shape,
        confidence_threshold=confidence_threshold,
        top_k=top_k,
        class_names=class_names,
        return_format=return_format
    )


#
# INTERNAL
#
def _load_backbone_modules():
    backbones = []
    modules = {}
    try:
        from soundhub_utils.audio import pytorch_audio
        backbones.append('pytorch')
        modules['pytorch'] = pytorch_audio
    except ImportError:
        pass
    try:
        from soundhub_utils.audio import librosa_audio
        backbones.append('librosa')
        modules['librosa'] = librosa_audio
    except ImportError:
        pass
    try:
        from soundhub_utils.audio import tensorflow_audio
        backbones.append('tensorflow')
        modules['tensorflow'] = tensorflow_audio
    except ImportError:
        pass
    default = ch.get('DEFAULT_AUDIO_BACKBONE', backbones[0]) 
    return backbones, modules, default


def _get_backend(backbone: Optional[str] = None):
    """
    Get the audio backend module.

    Args:
        backbone: Backend to use ('tensorflow', 'pytorch', or 'librosa').
                 If None, uses DEFAULT_BACKBONE.

    Returns:
        Backend module

    Raises:
        ValueError: If backbone is not available
    """
    if backbone is None:
        backbone = DEFAULT_BACKBONE

    if backbone not in AVAILABLE_BACKBONES:
        available_str = ', '.join(AVAILABLE_BACKBONES) if AVAILABLE_BACKBONES else 'none'
        source_map = {
            'tensorflow': 'tensorflow (conda or pip)',
            'pytorch': 'pytorch (conda or pip)',
            'librosa': 'librosa (conda or pip)'
        }
        install_msg = source_map.get(backbone, backbone)
        raise ValueError(
            f"Selected backbone '{backbone}' is not available. "
            f"Available backbones: {available_str}. "
            f"To use '{backbone}', install {install_msg}"
        )

    return BACKBONE_MODULES[backbone]


#
# INTERNAL SPECTROGRAM GENERATION
#
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]]:
    """
    Internal spectrogram generation function supporting both PNG and TFRecord output.
    """
    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 = _shared.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
        from soundhub_utils.utils.tf import TFRecordSpectrogramWriter
        with TFRecordSpectrogramWriter(
            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_array(audio_data: AudioData, to_numpy: bool = False) -> Union['tf.Tensor', 'torch.Tensor', 'np.ndarray']:
    """
    Extract audio tensor from AudioData.

    Args:
        audio_data: AudioData container
        to_numpy: If True, return numpy array; otherwise return original tensor type

    Returns:
        Audio array as numpy array or original tensor type
    """
    if to_numpy:
        try:
            return audio_data.tensor.numpy()
        except:
            # Fallback for arrays that don't have .numpy() method
            import numpy as np
            return np.asarray(audio_data.tensor)
    return audio_data.tensor


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

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


#
# Detect available backends
#
AVAILABLE_BACKBONES, BACKBONE_MODULES, DEFAULT_BACKBONE = _load_backbone_modules()
