"""
PyTorch Audio Processing Module

PyTorch-based audio processing and spectrogram generation utilities for soundhub

License: BSD 3-Clause
"""

import ast
import os
from pathlib import Path
from typing import Union, Optional, Tuple, List, Dict, Any
import numpy as np
import torch
import pandas as pd
import librosa
from soundhub_utils.constants import SUPPORTED_AUDIO_FORMATS
from soundhub_utils.audio import _shared
from soundhub_utils.audio._shared import AudioData
from cocina.config_handler import ConfigHandler
ch = ConfigHandler(__file__)


#
# DEFAULTS (from ConfigHandler for module-specific config)
#


#
# MAIN
#
def load_audio(
    src: Union[str, Path],
    sample_rate: Optional[int] = None,
    file_type: Optional[str] = None
) -> AudioData:
    """
    Load audio file using PyTorch and librosa.

    Note: Uses librosa for loading (same as tensorflow_audio) to avoid torchaudio
    hanging issues with large files.

    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

    Returns:
        AudioData containing tensor, sample rate, duration, and channels

    Raises:
        IOError: If file cannot be loaded
        ValueError: If invalid file type or sample rate specified
    """
    src_path = Path(src)

    # Auto-detect file type if not provided
    if file_type is None:
        file_type = src_path.suffix.lower()

    if file_type not in SUPPORTED_AUDIO_FORMATS:
        raise ValueError(f"Unsupported audio format: {file_type}")

    # Load audio using librosa (handles FLAC, WAV, MP3, etc.)
    # librosa automatically handles resampling when sr parameter is provided
    audio_np, loaded_sr = librosa.load(str(src_path), sr=sample_rate, mono=True)

    # Convert to PyTorch tensor
    audio_tensor = torch.from_numpy(audio_np).float()

    # Calculate duration
    duration = len(audio_np) / loaded_sr

    return AudioData(
        tensor=audio_tensor,
        sample_rate=loaded_sr,
        duration=duration,
        channels=1  # librosa loads as mono
    )



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'
) -> Union[Dict[str, Any], pd.DataFrame]:
    """
    TODO: inference for pytorch models
    """
    raise NotImplementedError
