"""
Python/Librosa Audio Processing Module

Librosa/numpy-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 matplotlib.pyplot as plt
import pandas as pd
import librosa
from soundhub_utils import names
from soundhub_utils.constants import SUPPORTED_AUDIO_FORMATS
from soundhub_utils.audio import _shared
from soundhub_utils.audio._shared import AudioData
from cocina.printer import Printer
from cocina.config_handler import ConfigHandler
printer = Printer()
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 librosa.

    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 array, 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.)
    audio_np, loaded_sr = librosa.load(str(src_path), sr=sample_rate, mono=True)

    # Calculate duration
    duration = len(audio_np) / loaded_sr

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


def to_array(audio_data: AudioData, to_numpy: bool = False) -> np.ndarray:
    """
    Extract audio array from AudioData.

    Args:
        audio_data: AudioData container
        to_numpy: Ignored (always returns numpy array)

    Returns:
        Audio array as numpy array
    """
    return audio_data.tensor




def pysox_spectrogram(
    audio_array: Union[np.ndarray, 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
) -> Union[np.ndarray, str]:
    """
    Create Sox-matching spectrogram using librosa/numpy with absolute dB normalization.

    Uses absolute dB normalization (Method 2) for stable, consistent Sox matching.
    Achieves ~91.8% within ±5 pixels, stable across different sample sizes.

    Args:
        audio_array: Audio data as numpy array or AudioData
        shape: Output shape of spectrogram (height, width)
        dest: Optional file path to save spectrogram as PNG
        segment: Segment number for extracting specific audio portion
        segment_duration: Duration of segment in seconds
        segment_overlap: Overlap between segments in seconds
        sample_rate: Audio sample rate in Hz (required for array inputs)
        create_parents: Create parent directories if needed
        overwrite: Overwrite existing file if it exists

    Returns:
        Spectrogram array (uint8) if dest is None, otherwise path to saved PNG file

    Raises:
        ValueError: If invalid parameters provided
        FileExistsError: If file exists and overwrite=False

    Notes:
        - Uses exact Sox algorithm: n_fft=512, hop_length=20
        - Block averaging matches Sox's implementation
        - dB_range=80.0, brightness_correction=-96 for Sox compatibility
        - Output is uint8 grayscale (4-231 range like Sox)
        - Absolute dB normalization (Method 2) is stable across sample sizes
        - See /workspace/py-sox-testing/sh/NORMALIZATION_FIX.md for details
    """
    # 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 = _shared.DEFAULT_SAMPLE_RATE

    # Ensure shape is a tuple or list of integers
    if isinstance(shape, str):
        shape = ast.literal_eval(shape)

    target_height, target_width = shape[0], shape[1]

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

    # Get Sox parameters from config
    n_fft = ch.get('SOX_N_FFT', 512)
    hop_length = ch.get('SOX_HOP_LENGTH', 20)
    dB_range = ch.get('SOX_DB_RANGE', 80.0)

    # METHOD 2: Always use absolute dB normalization (more stable across sample sizes)
    # NOTE: Ignoring reference_db parameter - using absolute dB regardless
    brightness_correction = -96  # Optimized for Method 2 (absolute dB)

    # STFT parameters
    window = np.hanning(n_fft)

    # Compute STFT with center padding (matching Sox)
    D = librosa.stft(
        y,
        n_fft=n_fft,
        hop_length=hop_length,
        window=window,
        center=True
    )

    # Compute POWER spectrum (not magnitude)
    S_power = np.abs(D) ** 2

    # Sox's block averaging algorithm
    # Accumulate power across blocks, THEN convert to dB
    n_frames = S_power.shape[0]
    n_output_cols = target_width
    frames_per_block = S_power.shape[1] / n_output_cols

    S_power_accumulated = np.zeros((n_frames, n_output_cols))

    for col_idx in range(n_output_cols):
        # Calculate frame range for this output column
        start_frame = col_idx * frames_per_block
        end_frame = (col_idx + 1) * frames_per_block

        start_idx = int(start_frame)
        end_idx = int(np.ceil(end_frame))

        if end_idx > S_power.shape[1]:
            end_idx = S_power.shape[1]

        # Accumulate power (Sox's += operation)
        accumulated = np.sum(S_power[:, start_idx:end_idx], axis=1)

        # Average (Sox's block_norm = 1.0 / block_steps)
        num_frames = end_idx - start_idx
        S_power_accumulated[:, col_idx] = accumulated / num_frames

    # NOW convert to dB (after power accumulation)
    S_db = 10 * np.log10(S_power_accumulated + 1e-10)

    # METHOD 2: Absolute dB normalization (no reference subtraction)
    # - More stable across different sample sizes
    # - No first-pass needed to calculate global_peak_db
    # - Achieved 91.8% within ±5px in testing
    S_db_norm = S_db  # Use absolute dB values directly

    # Map to pixel values (Sox formula)
    S_pixels = 4 + np.clip((S_db_norm + dB_range) / dB_range * 227, 0, 227)

    # Apply brightness correction
    S_pixels = S_pixels + brightness_correction
    S_pixels = np.clip(S_pixels, 4, 231)

    # Flip vertically (Sox has high freq at top)
    S_pixels = np.flipud(S_pixels)

    # Convert to uint8
    spectrogram = S_pixels.astype(np.uint8)

    # Verify dimensions
    assert spectrogram.shape == (target_height, target_width), \
        f"Output shape {spectrogram.shape} != expected ({target_height}, {target_width})"

    if dest is None:
        return spectrogram

    # Save spectrogram as grayscale PNG (raw data, no colorbar)
    _shared.save_spectrogram_png(spectrogram, dest, create_parents, overwrite, colormap=None)
    return str(dest)


def spectrograms(
    src: Union[str, Path, AudioData, np.ndarray],
    dest_folder: str,
    spectrogram_shape: Tuple[int, int],
    basename: Optional[str] = None,
    directory: Optional[str] = None,
    segment_duration: Optional[float] = None,
    segment_overlap: Optional[float] = None,
    sample_rate: Optional[int] = None,
    spectrogram_type: Optional[str] = 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[Union[np.ndarray, str]], List[str]]:
    """
    Generate multiple spectrograms from audio source with segmentation.

    Args:
        src: Audio source - file path, AudioData, or numpy array
        dest_folder: Folder to save spectrogram files
        spectrogram_shape: Output shape for each spectrogram (height, width)
        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 (required for array inputs)
        spectrogram_type: Type of spectrogram ('sox' only for python backend)
        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 arrays or file paths
        If output_format='tfrecord': List of TFRecord file paths

    Raises:
        ValueError: If invalid parameters provided
        IOError: If audio loading or file operations fail
        FileExistsError: If files exist and overwrite=False
    """
    # Validate spectrogram_type (Python only supports pysox and sox for now)
    if spectrogram_type not in [None, 'sox', 'pysox']:
        raise ValueError(f"Python backend only supports spectrogram_type=None, 'sox', or 'pysox', got: {spectrogram_type}")

    # Get output format from config if not specified
    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

    # Handle different source types
    _audio_path = None
    _native_sr = None
    _total_frames = None

    if isinstance(src, (str, Path)):
        # Streaming mode: read only file header, don't load audio into memory
        audio_info = _shared.get_audio_info(src)
        actual_sample_rate = sample_rate if sample_rate is not None else audio_info.sample_rate
        audio_array = None
        _audio_path = str(src)
        _native_sr = audio_info.sample_rate
        _total_frames = audio_info.frames
        audio_length_seconds = audio_info.duration
        if basename is None:
            basename = Path(src).stem
    elif isinstance(src, AudioData):
        # Use AudioData directly
        audio_array = src.tensor
        actual_sample_rate = src.sample_rate
        audio_length_seconds = len(audio_array) / actual_sample_rate
        if basename is None:
            basename = "audio_data"
    elif isinstance(src, np.ndarray):
        # Use array directly
        audio_array = src
        if sample_rate is None:
            raise ValueError("sample_rate must be provided for numpy array input")
        actual_sample_rate = sample_rate
        audio_length_seconds = len(audio_array) / actual_sample_rate
        if basename is None:
            basename = "audio_array"
    else:
        raise ValueError(f"Unsupported source type: {type(src)}")

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

    # Select spectrogram function based on type
    if spectrogram_type == 'pysox':
        spec_func = pysox_spectrogram
    elif spectrogram_type == 'sox':
        spec_func = lambda *args, **kwargs: _shared.sox_spectrogram(*args, **kwargs, backend='librosa')
    else:
        spec_func = pysox_spectrogram  # Default to pysox

    # Use shared spectrogram generation (handles both PNG and TFRecord)
    return _shared.generate_spectrograms(
        spec_func=spec_func,
        audio_data=audio_array,
        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=_audio_path,
        native_sample_rate=_native_sr,
        total_frames=_total_frames,
    )
