from __future__ import annotations

import wave
from pathlib import Path

import numpy as np

from app.errors import VaaniError

_MEL: np.ndarray | None = None


def assign_guest_names(labels: list[int]) -> list[str]:
    mapping: dict[int, str] = {}
    names: list[str] = []
    for label in labels:
        if label not in mapping:
            mapping[label] = f"Guest {len(mapping) + 1}"
        names.append(mapping[label])
    return names


def _mel_bank(sample_rate: int = 16000, n_fft: int = 512, n_mels: int = 26) -> np.ndarray:
    global _MEL
    if _MEL is not None:
        return _MEL
    n_freq = n_fft // 2 + 1

    def hz_to_mel(hz: float) -> float:
        return 2595.0 * np.log10(1.0 + hz / 700.0)

    def mel_to_hz(mel: float) -> float:
        return 700.0 * (10 ** (mel / 2595.0) - 1.0)

    points = np.linspace(hz_to_mel(0), hz_to_mel(sample_rate / 2), n_mels + 2)
    hz = np.vectorize(mel_to_hz)(points)
    bins = np.floor((n_fft + 1) * hz / sample_rate).astype(int)
    bins = np.clip(bins, 0, n_freq - 1)
    bank = np.zeros((n_mels, n_freq), dtype=np.float64)
    for index in range(n_mels):
        left, center, right = int(bins[index]), int(bins[index + 1]), int(bins[index + 2])
        if center <= left:
            center = left + 1
        if right <= center:
            right = center + 1
        for bin_index in range(left, min(center, n_freq)):
            bank[index, bin_index] = (bin_index - left) / (center - left)
        for bin_index in range(center, min(right, n_freq)):
            bank[index, bin_index] = (right - bin_index) / (right - center)
    _MEL = bank
    return bank


def _pitch(samples: np.ndarray, sample_rate: int) -> float:
    if samples.size < sample_rate // 50:
        return 0.0
    step = max(1, sample_rate // 8000)
    series = samples[::step]
    rate = sample_rate / step
    series = series - np.mean(series)
    energy = float(np.dot(series, series))
    if energy < 1e-6:
        return 0.0
    min_lag = max(1, int(rate / 320))
    max_lag = int(rate / 70)
    if max_lag >= series.size or max_lag <= min_lag:
        return 0.0
    width = 1 << int(np.ceil(np.log2(series.size * 2)))
    spectrum = np.fft.rfft(series, n=width)
    corr = np.fft.irfft(spectrum * np.conj(spectrum), n=width)[: series.size]
    window = corr[min_lag:max_lag]
    if window.size == 0 or corr[0] <= 0:
        return 0.0
    lag = int(np.argmax(window)) + min_lag
    if corr[lag] < 0.3 * corr[0]:
        return 0.0
    return float(rate / lag)


def embed(samples: np.ndarray, sample_rate: int = 16000) -> np.ndarray:
    if samples.size == 0:
        return np.zeros(28, dtype=np.float64)
    series = np.asarray(samples, dtype=np.float64)
    if series.size > sample_rate * 12:
        series = series[: sample_rate * 12]
    emphasised = np.append(series[0], series[1:] - 0.97 * series[:-1])
    frame = int(0.025 * sample_rate)
    hop = int(0.010 * sample_rate)
    if emphasised.size < frame:
        emphasised = np.pad(emphasised, (0, frame - emphasised.size))
    n_frames = 1 + (emphasised.size - frame) // hop
    window = np.hamming(frame)
    n_fft = 512
    spec = np.empty((n_frames, n_fft // 2 + 1), dtype=np.float64)
    for index in range(n_frames):
        start = index * hop
        spec[index] = np.abs(np.fft.rfft(emphasised[start : start + frame] * window, n=n_fft))
    mel = np.maximum(_mel_bank() @ spec.T, 1e-10)
    log_mel = np.log(mel)
    n_mels = log_mel.shape[0]
    n_mfcc = 13
    coeffs = np.empty((n_mfcc, log_mel.shape[1]), dtype=np.float64)
    slots = np.arange(n_mels)
    for k in range(n_mfcc):
        coeffs[k] = np.sum(log_mel * np.cos(np.pi * k * (slots[:, None] + 0.5) / n_mels), axis=0)
    rms = float(np.sqrt(np.mean(np.square(series))))
    features = np.concatenate([coeffs.mean(axis=1), coeffs.std(axis=1), [ _pitch(series, sample_rate), rms ]])
    if not np.isfinite(features).all():
        return np.nan_to_num(features)
    return features


def _cluster(vectors: np.ndarray, speakers: int) -> list[int]:
    count = len(vectors)
    if count == 0:
        return []
    if speakers <= 1 or count == 1 or np.allclose(vectors, vectors[0]):
        return [0] * count
    speakers = min(speakers, count)
    from sklearn.cluster import KMeans
    from sklearn.preprocessing import StandardScaler

    scaled = StandardScaler().fit_transform(vectors)
    scaled = np.nan_to_num(scaled)
    labels = KMeans(n_clusters=speakers, n_init=10, random_state=0).fit_predict(scaled)
    return [int(label) for label in labels]


def _smooth(labels: list[int], durations: list[float]) -> list[int]:
    smoothed = list(labels)
    for index, duration in enumerate(durations):
        if duration >= 1.0 or len(smoothed) < 2:
            continue
        previous = smoothed[index - 1] if index else None
        nxt = smoothed[index + 1] if index + 1 < len(smoothed) else None
        if previous is not None and previous == nxt:
            smoothed[index] = previous
        elif previous is not None and nxt is None:
            smoothed[index] = previous
    return smoothed


def label_speakers(
    chunks: list[dict],
    samples: np.ndarray,
    sample_rate: int,
    speakers: int,
) -> tuple[list[dict], str | None]:
    if not chunks:
        return [], None
    usable: list[int] = []
    vectors: list[np.ndarray] = []
    for index, chunk in enumerate(chunks):
        start = max(0, int(float(chunk["start"]) * sample_rate))
        end = min(samples.size, int(float(chunk["end"]) * sample_rate))
        piece = samples[start:end]
        if piece.size < int(sample_rate * 0.25):
            continue
        usable.append(index)
        vectors.append(embed(piece, sample_rate))
    labels = [-1] * len(chunks)
    if vectors:
        clustered = _cluster(np.vstack(vectors), speakers)
        for index, label in zip(usable, clustered):
            labels[index] = label
    last = 0
    for index, label in enumerate(labels):
        if label >= 0:
            last = label
        else:
            labels[index] = last
    durations = [max(0.0, float(chunk["end"]) - float(chunk["start"])) for chunk in chunks]
    labels = _smooth(labels, durations)
    names = assign_guest_names(labels)
    labeled = []
    for chunk, name in zip(chunks, names):
        labeled.append({**chunk, "speaker": name})
    found = len({name for name in names})
    warning = None
    if speakers > 1 and found < speakers:
        warning = "Only one voice stood out, so the transcript uses Guest 1."
    return labeled, warning


def read_wav(path: Path) -> tuple[np.ndarray, int]:
    try:
        with wave.open(str(path), "rb") as handle:
            rate = handle.getframerate()
            channels = handle.getnchannels()
            width = handle.getsampwidth()
            raw = handle.readframes(handle.getnframes())
    except wave.Error as exc:
        raise VaaniError(
            "The audio file could not be read for the transcript.",
            "transcript_failed",
            500,
            "retry",
        ) from exc
    if width != 2:
        raise VaaniError(
            "The audio conversion did not produce a 16-bit file.",
            "transcript_failed",
            500,
            "retry",
        )
    data = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0
    if channels > 1:
        data = data.reshape(-1, channels).mean(axis=1)
    return np.ascontiguousarray(data), rate