from typing import Optional, Callable
import resampy
import numpy as np
import librosa
from PIL import Image
import io
import matplotlib.pyplot as plt
from constants import DEFAULT_SR, TARGET_LOUDNESS, TARGET_TP, STREAMER_CLASSES
import warnings
import pyloudnorm as pyln


def to_gradio_audio(x: np.ndarray, sr: int) -> tuple[int, np.ndarray]:
    """Return (sample_rate, int16 mono array) for Gradio Audio. Gradio expects int16;
    passing float32 triggers an internal conversion and a warning."""
    x = np.asarray(x)

    # Remove extra dims like (1, n, 1) etc.
    x = np.squeeze(x)

    # If it's (channels, samples), transpose to (samples, channels)
    if x.ndim == 2 and x.shape[0] in (1, 2) and x.shape[1] > x.shape[0]:
        x = x.T

    # Ensure mono is (n_samples,)
    if x.ndim == 2 and x.shape[1] == 1:
        x = x[:, 0]

    x = x.astype(np.float32)
    x = np.clip(x, -1.0, 1.0)
    # Gradio Audio expects int16; convert here so Gradio doesn't convert and warn
    x = (x * 32767).astype(np.int16)

    return (sr, x)


def spec_image(
    audio_array: np.ndarray,
    sr: int = DEFAULT_SR,
    n_fft: int = 2048,
    hop_length: int = 512,
    n_mels: int = 128,
    fmax: Optional[float] = None,
) -> Image.Image:
    """
    Generate a mel-spectrogram image from an audio array.
    """
    y = audio_array.flatten()  # Ensure it's 1D
    S = librosa.feature.melspectrogram(
        y=y,
        sr=sr,
        n_fft=n_fft,
        hop_length=hop_length,
        n_mels=n_mels,
        fmax=fmax or sr // 2,
    )
    S_db = librosa.power_to_db(S, ref=np.max(S))
    fig, ax = plt.subplots(figsize=(8, 3), dpi=150)
    img = librosa.display.specshow(
        S_db, sr=sr, hop_length=hop_length, x_axis="time", y_axis="mel", ax=ax
    )
    cbar = fig.colorbar(img, ax=ax, format="%+2.0f dB")
    cbar.set_label("dB")
    ax.set_title("Mel-spectrogram")
    ax.set_xlabel("Time in s")
    ax.set_ylabel("Frequency in Hz")
    fig.tight_layout(pad=0.2)
    buf = io.BytesIO()
    fig.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
    plt.close(fig)
    buf.seek(0)
    return Image.open(buf).convert("RGB")


def compute_wer(reference: str, hypothesis: str) -> float:
    """
    Compute Word Error Rate (WER) between reference and hypothesis transcripts.
    """
    ref_words = reference.split()
    hyp_words = hypothesis.split()
    d = np.zeros((len(ref_words) + 1, len(hyp_words) + 1), dtype=np.uint8)
    for i in range(len(ref_words) + 1):
        d[i][0] = i
    for j in range(len(hyp_words) + 1):
        d[0][j] = j
    for i in range(1, len(ref_words) + 1):
        for j in range(1, len(hyp_words) + 1):
            if ref_words[i - 1] == hyp_words[j - 1]:
                cost = 0
            else:
                cost = 1
            d[i][j] = min(
                d[i - 1][j] + 1,      # Deletion
                d[i][j - 1] + 1,      # Insertion
                d[i - 1][j - 1] + cost,  # Substitution
            )
    wer = d[len(ref_words)][len(hyp_words)] / max(len(ref_words), 1)
    return wer


def measure_loudness(x: np.ndarray, sr: int) -> float:
    meter = pyln.Meter(sr)
    return float(meter.integrated_loudness(x))


def true_peak_limiter(x: np.ndarray, sr: int, max_true_peak: float = TARGET_TP) -> np.ndarray:
    upsampled_sr = 192000
    x_upsampled = librosa.resample(x, orig_sr=sr, target_sr=upsampled_sr)
    true_peak = np.max(np.abs(x_upsampled))

    if true_peak > 0:
        true_peak_db = 20 * np.log10(true_peak)
        if true_peak_db > max_true_peak:
            gain_db = max_true_peak - true_peak_db
            gain = 10 ** (gain_db / 20)
            x_upsampled = x_upsampled * gain

    x_limited = librosa.resample(x_upsampled, orig_sr=upsampled_sr, target_sr=sr)
    x_limited = librosa.util.fix_length(x_limited, size=x.shape[-1])
    return x_limited.astype("float32")


def normalize_lufs(x: np.ndarray, sr: int) -> np.ndarray:
    """
    Normalize audio to a fixed integrated loudness target and limit true peak.
    """
    try:
        current_lufs = measure_loudness(x, sr)
        
        if not np.isfinite(current_lufs):
            return x.astype("float32")

        gain_db = TARGET_LOUDNESS - current_lufs
        gain = 10 ** (gain_db / 20)

        y = x * gain
        y = true_peak_limiter(y, sr, max_true_peak=TARGET_TP)

        return y.astype("float32")
    except Exception as e:
        warnings.warn(f"LUFS normalization failed, returning input unchanged: {e}")
        return x.astype("float32")




def transcribe_audio(
    audio_array: np.ndarray,
    sr: int,
    streamer_type: str = "Deepgram",
    stream_name: str = "RAW",
    on_update: Optional[Callable[[str], None]] = None,
):
    """
    Transcribe an audio array using the specified STT streamer.

    Args:
        audio_array (np.ndarray): Audio data array
        sr (int): Sample rate of the audio array
        streamer_type (str): "Soniox" or "Deepgram"
        stream_name (str): Optional label for streamer instance ("RAW", "ENHANCED", etc.)
        on_update: Optional callback(text: str) called with partial transcript as results stream in.

    Returns:
        str: Final transcript text
    """
    if sr != DEFAULT_SR:
        audio_array = resampy.resample(audio_array, sr, DEFAULT_SR)
        sr = DEFAULT_SR

    if streamer_type not in STREAMER_CLASSES:
        raise ValueError(
            f"Invalid streamer_type '{streamer_type}'. "
            f"Choose from: {', '.join(STREAMER_CLASSES.keys())}"
        )

    StreamerClass = STREAMER_CLASSES[streamer_type]
    streamer = StreamerClass(sr, stream_name, on_update=on_update)

    transcript = streamer.stream_array(audio_array)

    return transcript