"""
TensorFlow Audio Processing Module

TensorFlow-based audio processing and spectrogram generation utilities for soundhub

License: BSD 3-Clause
"""
def safe_load_tensorflow(legacy_keras):
    if legacy_keras:
        os.environ['TF_USE_LEGACY_KERAS'] = str(int(legacy_keras))
    import tensorflow as tf
    return tf


import os
import ast
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 sox_tensorflow import processor
from cocina.printer import Printer
from cocina.config_handler import ConfigHandler
from importlib import reload
printer = Printer()
ch = ConfigHandler(__file__)
tf = safe_load_tensorflow(ch.get('TF_USE_LEGACY_KERAS'))


#
# MAIN
#
def load_audio(
    src: Union[str, Path],
    sample_rate: Optional[int] = None,
    file_type: Optional[str] = None
) -> AudioData:
    """
    Load audio file using TensorFlow and 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 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 TensorFlow tensor
    audio_tensor = tf.convert_to_tensor(audio_np, dtype=tf.float32)

    # 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 to_array(audio_data: AudioData, to_numpy: bool = False) -> Union['tf.Tensor', np.ndarray]:
    """
    Extract audio tensor from AudioData.

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

    Returns:
        Audio array as numpy array or TensorFlow tensor
    """
    if to_numpy:
        return audio_data.tensor.numpy()
    return audio_data.tensor




def spectrogram(
    audio_array: Union['tf.Tensor', np.ndarray],
    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,
    mel: bool = False,
    frame_length: Optional[int] = None,
    frame_step: Optional[int] = None,
    colormap: Optional[str] = 'viridis',
    create_parents: bool = True,
    overwrite: bool = True
) -> Union['tf.Tensor', str]:
    """
    Convert audio array to spectrogram using TensorFlow.

    Args:
        audio_array: Audio data as TensorFlow tensor or numpy array
        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
        mel: If True, compute mel-spectrogram instead of regular spectrogram
        frame_length: STFT frame length (defaults to config value or 512)
        frame_step: STFT frame step (defaults to config value or 256)
        colormap: Matplotlib colormap ('gray', 'viridis', etc.) or None for raw grayscale
        create_parents: Create parent directories if needed
        overwrite: Overwrite existing file if it exists

    Returns:
        Spectrogram tensor if dest is None, otherwise path to saved PNG file

    Raises:
        ValueError: If invalid parameters provided
        FileExistsError: If file exists and overwrite=False
    """
    if sample_rate is None:
        sample_rate = _shared.DEFAULT_SAMPLE_RATE

    # Convert to TensorFlow tensor if needed
    if isinstance(audio_array, np.ndarray):
        audio_tensor = tf.convert_to_tensor(audio_array, dtype=tf.float32)
    else:
        audio_tensor = audio_array

    # Extract segment if specified
    if segment is not None:
        start_idx, end_idx = _shared.calculate_segment_indices(
            len(audio_tensor), sample_rate, segment, segment_duration, segment_overlap
        )
        audio_tensor = audio_tensor[start_idx:end_idx]
    if mel:
        return mel_spectrogram(
            audio_tensor, shape, dest, None, None, None, sample_rate,
            frame_length=frame_length, frame_step=frame_step,
            create_parents=create_parents, overwrite=overwrite
        )

    # Get STFT parameters from function args, config, or defaults
    if frame_length is None:
        frame_length = ch.get('STFT_FRAME_LENGTH', 512)
    if frame_step is None:
        frame_step = ch.get('STFT_FRAME_STEP', 96)

    stft = tf.signal.stft(
        audio_tensor,
        frame_length=frame_length,
        frame_step=frame_step,
        window_fn=tf.signal.hann_window
    )

    # Convert to power spectrogram (magnitude squared)
    power_spectrogram = tf.square(tf.abs(stft))

    # Convert to dB scale with 90 dB dynamic range (matching SoX -z 90)
    # SoX formula: 10 * log10(power) with 90 dB range
    spectrogram_db = 10.0 * tf.math.log(tf.maximum(power_spectrogram, 1e-9)) / tf.math.log(10.0)

    # Apply 90 dB dynamic range clipping (like SoX -z 90)
    db_max = tf.reduce_max(spectrogram_db)
    db_min = db_max - 90.0  # 90 dB dynamic range
    spectrogram_tensor = tf.clip_by_value(spectrogram_db, db_min, db_max)


    # Resize to target shape if needed
    if tuple(spectrogram_tensor.shape) != shape:
        # Ensure shape is a tuple or list of integers
        if isinstance(shape, str):
            # Handle string representation like "(257, 1000)"
            shape = ast.literal_eval(shape)

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

        # Add batch and channel dimensions for resize
        spec_resized = tf.image.resize(
            tf.expand_dims(tf.expand_dims(spectrogram_tensor, 0), -1),
            tf.convert_to_tensor([target_height, target_width], dtype=tf.int32)
        )
        spectrogram_tensor = tf.squeeze(spec_resized)

    if dest is None:
        return spectrogram_tensor

    # Save spectrogram as PNG
    _shared.save_spectrogram_png(spectrogram_tensor, dest, create_parents, overwrite, colormap)
    return str(dest)


def mel_spectrogram(
    audio_array: Union['tf.Tensor', np.ndarray],
    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,
    num_mel_bins: int = 128,
    frame_length: Optional[int] = None,
    frame_step: Optional[int] = None,
    colormap: Optional[str] = 'viridis',
    create_parents: bool = True,
    overwrite: bool = True
) -> Union['tf.Tensor', str]:
    """
    Convert audio array to mel-spectrogram using tf.keras.utils.audio.

    Args:
        audio_array: Audio data as TensorFlow tensor or numpy array
        shape: Output shape of mel-spectrogram (height, width)
        dest: Optional file path to save mel-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
        num_mel_bins: Number of mel frequency bins
        frame_length: STFT frame length (defaults to config value or 512)
        frame_step: STFT frame step (defaults to config value or 256)
        colormap: Matplotlib colormap ('gray', 'viridis', etc.) or None for raw grayscale
        create_parents: Create parent directories if needed
        overwrite: Overwrite existing file if it exists

    Returns:
        Mel-spectrogram tensor if dest is None, otherwise path to saved PNG file

    Raises:
        ValueError: If invalid parameters provided
        FileExistsError: If file exists and overwrite=False
    """
    if sample_rate is None:
        sample_rate = _shared.DEFAULT_SAMPLE_RATE

    # Convert to TensorFlow tensor if needed
    if isinstance(audio_array, np.ndarray):
        audio_tensor = tf.convert_to_tensor(audio_array, dtype=tf.float32)
    else:
        audio_tensor = audio_array

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

    # Use TensorFlow's built-in mel-spectrogram computation
    # First compute STFT with configurable parameters
    if frame_length is None:
        frame_length = ch.get('STFT_FRAME_LENGTH', 512)
    if frame_step is None:
        frame_step = ch.get('STFT_FRAME_STEP', 96)

    stft = tf.signal.stft(
        audio_tensor,
        frame_length=frame_length,
        frame_step=frame_step,
        window_fn=tf.signal.hann_window
    )

    # Convert to power spectrogram
    power_spectrogram = tf.real(stft * tf.math.conj(stft))

    # Create mel-scale filterbank
    linear_to_mel_weight_matrix = tf.signal.linear_to_mel_weight_matrix(
        num_mel_bins=num_mel_bins,
        num_spectrogram_bins=power_spectrogram.shape[-2],
        sample_rate=sample_rate,
        lower_edge_hertz=0.0,
        upper_edge_hertz=sample_rate / 2.0
    )

    # Apply mel-scale filterbank
    mel_spec = tf.tensordot(power_spectrogram, linear_to_mel_weight_matrix, 1)

    # Resize to target shape if needed
    if tuple(mel_spec.shape) != shape:
        # Ensure shape is a tuple or list of integers
        if isinstance(shape, str):
            # Handle string representation like "(257, 1000)"
            shape = ast.literal_eval(shape)

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

        # Add batch and channel dimensions for resize
        mel_resized = tf.image.resize(
            tf.expand_dims(tf.expand_dims(mel_spec, 0), -1),
            tf.convert_to_tensor([target_height, target_width], dtype=tf.int32)
        )
        mel_spec = tf.squeeze(mel_resized)

    if dest is None:
        return mel_spec

    # Save mel-spectrogram as PNG
    _shared.save_spectrogram_png(mel_spec, dest, create_parents, overwrite, colormap)
    return str(dest)

def pysox_spectrogram(
    audio_array: Union['tf.Tensor', 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['tf.Tensor', str]:
    # """
    # Create Sox-matching spectrogram using native TensorFlow 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 TensorFlow tensor, 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 tensor/array inputs)
    #     create_parents: Create parent directories if needed
    #     overwrite: Overwrite existing file if it exists

    # Returns:
    #     Spectrogram tensor (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
    # """
    return processor.spectrogram(
        audio_array=audio_array,
        shape=shape,
        dest=dest,
        segment=segment,
        segment_duration=segment_duration,
        segment_overlap=segment_overlap,
        sample_rate=sample_rate,
        create_parents=create_parents,
        overwrite=overwrite
     )


def spectrograms(
    src: Union[str, Path, AudioData, 'tf.Tensor', 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['tf.Tensor', str]], str]:
    """
    Generate multiple spectrograms from audio source with segmentation.

    Args:
        src: Audio source - file path, AudioData, TensorFlow tensor, or numpy array
        dest_folder: Folder to save spectrogram files (or TFRecord file if output_format='tfrecord')
        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 tensor/array inputs)
        spectrogram_type: Type of spectrogram ('mel', 'sox', or None for regular)
        create_parents: Create parent directories if needed
        overwrite: Overwrite existing files if they exist
        output_format: Output format ('png' or 'tfrecord'). If None, uses config default
        tfrecord_compression: TFRecord compression ('GZIP', 'ZLIB', or None). If None, uses config default
        tfrecord_chunk_size: Spectrograms per TFRecord file (None = single file). If None, uses config default
        metadata: Optional metadata dict to include in TFRecord (e.g., {'audio_uri': 's3://...', 'start_time': 500})

    Returns:
        If output_format='png': List of PNG file paths or spectrogram tensors
        If output_format='tfrecord': Single TFRecord file path (or list of paths if chunked)

    Raises:
        ValueError: If invalid parameters provided
        IOError: If audio loading or file operations fail
        FileExistsError: If files exist and overwrite=False
    """
    # Validate spectrogram_type
    if spectrogram_type not in [None, 'mel', 'sox', 'pysox']:
        raise ValueError(f"spectrogram_type must be None, 'mel', '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)

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

    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_tensor = 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_tensor = src.tensor
        actual_sample_rate = src.sample_rate
        audio_length_seconds = len(audio_tensor) / actual_sample_rate
        if basename is None:
            basename = "audio_data"
    elif isinstance(src, np.ndarray):
        # Convert numpy array to tensor
        audio_tensor = tf.convert_to_tensor(src, dtype=tf.float32)
        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_tensor) / actual_sample_rate
        if basename is None:
            basename = "audio_array"
    elif isinstance(src, 'tf.Tensor'):
        # Use tensor directly
        audio_tensor = src
        if sample_rate is None:
            raise ValueError("sample_rate must be provided for tensor input")
        actual_sample_rate = sample_rate
        audio_length_seconds = len(audio_tensor) / actual_sample_rate
        if basename is None:
            basename = "audio_tensor"
    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 == 'mel':
        spec_func = mel_spectrogram
    elif spectrogram_type == 'pysox':
        spec_func = pysox_spectrogram
    elif spectrogram_type == 'sox':
        spec_func = lambda *args, **kwargs: _shared.sox_spectrogram(*args, **kwargs, backend='tensorflow')
    else:
        spec_func = spectrogram

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


def inference(
    model: Union[str, tf.keras.Model],
    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]:
    """
    Run TensorFlow model inference on spectrogram files (PNG or TFRecord).

    Automatically detects input format based on file extension.

    Args:
        model: Either path to .h5 model file or loaded TensorFlow model
        files: List of file paths (PNG or TFRecord) 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')

    Returns:
        Dictionary or DataFrame containing inference results with predictions,
        confidences, and metadata for each input file

    Raises:
        FileNotFoundError: If model file not found
        ValueError: If invalid parameters provided
        IOError: If inference fails
    """
    # Set default spectrogram shape if not provided
    if spectrogram_shape is None:
        spectrogram_shape = _shared.DEFAULT_SPECTROGRAM_SHAPE

    # Load model if path provided
    if isinstance(model, str):
        if not os.path.exists(model):
            raise FileNotFoundError(f"Model file not found: {model}")
        tf_model = tf.keras.models.load_model(model)
    else:
        tf_model = model


    # Convert files input to list of paths
    if isinstance(files, pd.DataFrame):
        if 'Filename' in files.columns:
            file_paths = files['Filename'].tolist()
        elif len(files.columns) == 1:
            file_paths = files.iloc[:, 0].tolist()
        else:
            raise ValueError("DataFrame must have 'Filename' column or single column")
    else:
        file_paths = files

    # Validate file paths
    for file_path in file_paths:
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"Spectrogram file not found: {file_path}")

    # Auto-detect input format (TFRecord vs PNG)
    is_tfrecord = _shared.is_tfrecord_path(file_paths)

    if is_tfrecord:
        # Load from TFRecord file(s)
        dataset = _shared.load_spectrograms_from_tfrecord(
            file_paths,
            spectrogram_shape,
            batch_size=100,
            compression=ch.get('TFRECORD_COMPRESSION', 'GZIP')
        )

        # Run inference on batches and collect results
        all_predictions = []
        all_metadata = []

        for spectrograms_batch, metadata_batch in dataset:
            batch_predictions = tf_model.predict(spectrograms_batch, verbose=0)
            all_predictions.append(batch_predictions)
            all_metadata.append(metadata_batch)

        # Concatenate all predictions
        predictions = np.concatenate(all_predictions, axis=0)

        # Build file paths from metadata for results processing
        # (TFRecord stores segment_index and filename, reconstruct paths if needed)
        result_file_paths = []
        for metadata_batch in all_metadata:
            for i in range(len(metadata_batch['segment_index'])):
                filename = metadata_batch['filename'][i].numpy().decode('utf-8')
                segment_idx = metadata_batch['segment_index'][i].numpy()
                # Create pseudo-path for results tracking
                result_file_paths.append(f"{filename}_segment{segment_idx:04d}")

    else:
        # Load spectrograms from PNG files
        input_data = _load_spectrograms_tf(file_paths, spectrogram_shape)

        # Run inference
        all_predictions = []
        for batch in input_data:
            batch_predictions = tf_model.predict(batch, verbose=0)
            all_predictions.append(batch_predictions)

        predictions = np.concatenate(all_predictions, axis=0)
        result_file_paths = file_paths

    # Process results
    results = _shared.process_inference_results(
        predictions, result_file_paths, confidence_threshold, top_k, class_names
    )

    if return_format == 'dataframe':
        return pd.DataFrame(results)
    else:
        return results


def _load_spectrograms_tf(file_paths: List[str], spectrogram_shape: Tuple[int, int]) -> '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)

    Returns:
        TensorFlow tensor with shape (batch_size, height, width, channels)
    """
    return _shared.load_spectrograms_tf(file_paths, spectrogram_shape, tf)


