import spaces
import os
import uuid
import hashlib
import numpy as np
import librosa
import librosa.display
import soundfile as sf
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import gradio as gr

TARGET_SR = 16000
N_FFT = 2048
HOP = 512
BLOCK_SECONDS = 2.0
PAYLOAD_BITS = 32
ALPHA = 0.018
LOW_HZ = 700.0
HIGH_HZ = 7000.0
DETECT_THRESHOLD = 0.22
MIN_BLOCKS = 3
EPS = 1e-8


def _seed_from_key(key):
    h = hashlib.sha256(str(int(key)).encode()).digest()
    return int.from_bytes(h[:8], "little") % (2**32 - 1)


def _payload_from_key(key):
    digest = hashlib.sha256(
        f"AudioShield-v3:{int(key)}".encode()
    ).digest()
    return np.unpackbits(
        np.frombuffer(digest[:4], dtype=np.uint8)
    ).astype(np.int8)


def _freq_bins(sr):
    freqs = librosa.fft_frequencies(sr=sr, n_fft=N_FFT)
    upper = min(HIGH_HZ, sr / 2 - 300)
    idx = np.where((freqs >= LOW_HZ) & (freqs <= upper))[0]
    if len(idx) < 20:
        raise ValueError("Fréquence d'échantillonnage trop faible.")
    return freqs, idx


def _make_keyed_pattern(n_freq, n_frames, key, bit_index, block_index):
    seed = (
        _seed_from_key(key)
        ^ ((bit_index + 1) * 0x9E3779B1)
        ^ ((block_index + 1) * 0x85EBCA77)
    ) & 0xFFFFFFFF

    rng = np.random.default_rng(seed)
    p = rng.choice([-1.0, 1.0], size=(n_freq, n_frames))

    if n_frames >= 5:
        kernel = np.array([1, 2, 3, 2, 1], dtype=np.float32)
        kernel /= kernel.sum()
        p = np.apply_along_axis(
            lambda row: np.convolve(row, kernel, mode="same"),
            1, p
        )

    p /= np.sqrt(np.mean(p * p) + EPS)
    return p


def _embed_mono(y, sr, key, alpha):
    y16 = librosa.resample(
        y.astype(np.float32),
        orig_sr=sr,
        target_sr=TARGET_SR
    )

    stft = librosa.stft(
        y16, n_fft=N_FFT, hop_length=HOP,
        win_length=N_FFT, window="hann"
    )
    mag = np.abs(stft)
    phase = np.angle(stft)

    _, fidx = _freq_bins(TARGET_SR)
    block_frames = max(1, int(BLOCK_SECONDS * TARGET_SR / HOP))
    n_blocks = max(1, int(np.ceil(mag.shape[1] / block_frames)))
    payload = _payload_from_key(key)

    wm_mag = mag.copy()
    used_blocks = 0

    for b in range(n_blocks):
        a = b * block_frames
        z = min((b + 1) * block_frames, mag.shape[1])

        if z - a < max(12, block_frames // 3):
            continue

        local = mag[fidx, a:z]
        ref = np.median(local, axis=1, keepdims=True)
        ref = np.maximum(
            ref,
            np.percentile(local, 25, axis=1, keepdims=True)
        )
        strength = np.clip(
            ref / (np.median(ref) + EPS), 0.25, 2.5
        )

        for bit_i, bit in enumerate(payload):
            p = _make_keyed_pattern(
                len(fidx), z - a, key, bit_i, b
            )
            symbol = 1.0 if bit else -1.0
            delta = alpha * symbol * p * strength
            wm_mag[fidx, a:z] *= np.exp(delta)

        used_blocks += 1

    out = librosa.istft(
        wm_mag * np.exp(1j * phase),
        hop_length=HOP,
        win_length=N_FFT,
        window="hann",
        length=len(y16)
    )

    out = np.clip(out, -0.999, 0.999)

    if sr != TARGET_SR:
        out = librosa.resample(
            out, orig_sr=TARGET_SR, target_sr=sr
        )
        out = out[:len(y)]
        if len(out) < len(y):
            out = np.pad(out, (0, len(y) - len(out)))

    return out, used_blocks


def _detect_mono(y, sr, key):
    y16 = librosa.resample(
        y.astype(np.float32),
        orig_sr=sr,
        target_sr=TARGET_SR
    )

    stft = librosa.stft(
        y16, n_fft=N_FFT, hop_length=HOP,
        win_length=N_FFT, window="hann"
    )
    mag = np.abs(stft)

    _, fidx = _freq_bins(TARGET_SR)
    block_frames = max(1, int(BLOCK_SECONDS * TARGET_SR / HOP))
    n_blocks = max(1, int(np.ceil(mag.shape[1] / block_frames)))
    payload = _payload_from_key(key)

    bit_scores = [[] for _ in range(PAYLOAD_BITS)]

    for b in range(n_blocks):
        a = b * block_frames
        z = min((b + 1) * block_frames, mag.shape[1])

        if z - a < max(12, block_frames // 3):
            continue

        x = mag[fidx, a:z]
        med = np.median(x, axis=1, keepdims=True)
        mad = np.median(np.abs(x - med), axis=1, keepdims=True) + EPS
        x = np.clip((x - med) / (4.0 * mad), -3.0, 3.0)

        for bit_i, bit in enumerate(payload):
            p = _make_keyed_pattern(
                len(fidx), z - a, key, bit_i, b
            )

            xx = x - np.mean(x)
            pp = p - np.mean(p)

            denom = (
                np.linalg.norm(xx) *
                np.linalg.norm(pp)
            ) + EPS

            corr = float(np.sum(xx * pp) / denom)
            bit_scores[bit_i].append(corr)

    if not all(bit_scores):
        return 0.0, 0, "Pas assez de blocs exploitables."

    scores = []

    for vals in bit_scores:
        vals = np.asarray(vals, dtype=np.float32)
        k = max(1, len(vals) // 2)
        strongest = vals[
            np.argsort(np.abs(vals))[-k:]
        ]
        scores.append(float(np.mean(strongest)))

    expected = np.where(payload > 0, 1.0, -1.0)
    aligned = np.asarray(scores) * expected

    confidence = float(np.mean(aligned))
    positive_bits = int(np.sum(aligned > 0))

    detected = (
        len(bit_scores[0]) >= MIN_BLOCKS
        and confidence >= DETECT_THRESHOLD
        and positive_bits >= int(PAYLOAD_BITS * 0.75)
    )

    return (
        confidence,
        positive_bits,
        "WATERMARK DÉTECTÉ"
        if detected else
        "WATERMARK NON CONFIRMÉ"
    )


def _load_audio(path):
    y, sr = librosa.load(
        path, sr=None, mono=False, duration=300
    )
    return y.astype(np.float32), sr


# ------------------------------------------------------------
# ZeroGPU functions
# ------------------------------------------------------------
# The @spaces.GPU decorator is required by Hugging Face
# ZeroGPU. Keep it on the OUTER processing functions.
# ------------------------------------------------------------

@spaces.GPU(duration=120)
def embed_watermark(audio_path, watermark_key=42, alpha=ALPHA):
    if not audio_path:
        return None, None, "Veuillez fournir un fichier audio."

    try:
        y, sr = _load_audio(audio_path)
        key = int(watermark_key)
        alpha = float(alpha)

        if y.ndim == 1:
            out, blocks = _embed_mono(y, sr, key, alpha)
            out_sf = out
            original = y
        else:
            channels = []
            blocks = 0

            for ch in range(y.shape[0]):
                wm, b = _embed_mono(
                    y[ch], sr, key, alpha
                )
                channels.append(wm)
                blocks = max(blocks, b)

            out_sf = np.vstack(channels).T
            original = y[0]

        uid = uuid.uuid4().hex[:8]
        output_path = (
            f"audio_watermarked_v3_{uid}.wav"
        )

        sf.write(
            output_path,
            out_sf,
            sr,
            subtype="PCM_24"
        )

        wm_plot = (
            out_sf if out_sf.ndim == 1
            else out_sf[:, 0]
        )

        D0 = librosa.amplitude_to_db(
            np.abs(librosa.stft(
                original,
                n_fft=N_FFT,
                hop_length=HOP
            )),
            ref=np.max
        )

        D1 = librosa.amplitude_to_db(
            np.abs(librosa.stft(
                wm_plot,
                n_fft=N_FFT,
                hop_length=HOP
            )),
            ref=np.max
        )

        fig, ax = plt.subplots(
            2, 1, figsize=(11, 7), sharex=True
        )

        librosa.display.specshow(
            D0, sr=sr, hop_length=HOP,
            x_axis="time", y_axis="hz",
            ax=ax[0]
        )
        ax[0].set_title(
            "Original — spectrogramme"
        )

        librosa.display.specshow(
            D1 - D0, sr=sr, hop_length=HOP,
            x_axis="time", y_axis="hz",
            ax=ax[1]
        )
        ax[1].set_title(
            "Différence spectrale — Watermark v3"
        )

        plt.tight_layout()

        plot_path = f"spectrogram_v3_{uid}.png"
        plt.savefig(plot_path, dpi=140)
        plt.close(fig)

        return (
            output_path,
            plot_path,
            "✅ Watermark v3 injecté.\n"
            f"Blocs utilisés : {blocks}\n"
            f"Clé : {key}\n"
            f"Alpha : {alpha:.3f}\n\n"
            "Watermark réparti dans le spectre "
            "sans porteuse ultrasonique fixe."
        )

    except Exception as e:
        return None, None, f"❌ Erreur : {e}"


@spaces.GPU(duration=120)
def detect_watermark(audio_path, watermark_key=42):
    if not audio_path:
        return "Veuillez fournir un fichier audio."

    try:
        y, sr = _load_audio(audio_path)
        key = int(watermark_key)

        if y.ndim == 1:
            conf, bits, status = _detect_mono(
                y, sr, key
            )
        else:
            results = [
                _detect_mono(y[ch], sr, key)
                for ch in range(y.shape[0])
            ]

            conf = float(
                np.mean([r[0] for r in results])
            )
            bits = int(
                np.mean([r[1] for r in results])
            )
            status = (
                "WATERMARK DÉTECTÉ"
                if all(
                    r[2] == "WATERMARK DÉTECTÉ"
                    for r in results
                )
                else
                "WATERMARK NON CONFIRMÉ"
            )

        icon = (
            "🟢"
            if status == "WATERMARK DÉTECTÉ"
            else "🔴"
        )

        return (
            f"{icon} {status}\n\n"
            f"Confiance : {conf:.3f}\n"
            f"Bits cohérents : {bits}/{PAYLOAD_BITS}\n"
            f"Seuil : {DETECT_THRESHOLD:.3f}\n"
            f"Clé testée : {key}\n\n"
            "⚠️ Résultat statistique du prototype. "
            "Ce résultat n'est pas une preuve cryptographique "
            "de provenance."
        )

    except Exception as e:
        return f"❌ Erreur : {e}"


# ------------------------------------------------------------
# Interface
# ------------------------------------------------------------

with gr.Blocks(
    title="AudioShield v3 — Robust Watermark"
) as demo:

    gr.Markdown(
        """
# 🛡️ AudioShield v3 — Watermarking audio robuste

Watermark invisible **à spectre étalé et clé secrète**,
sans tonalité ultrasonique fixe.

- MP3 / WAV / FLAC / OGG / M4A / AAC / AIFF
- Mono et stéréo
- Payload déterministe de 32 bits
- Répétition par blocs
- Détection multi-blocs
- Score de confiance
- Analyse spectrale Original / Watermark
"""
    )

    with gr.Tab("1. Injecter le Watermark"):

        with gr.Row():

            with gr.Column():

                audio_in = gr.Audio(
                    type="filepath",
                    label="Audio source"
                )

                key_in = gr.Number(
                    value=42,
                    label="Clé secrète",
                    precision=0
                )

                alpha_in = gr.Slider(
                    minimum=0.006,
                    maximum=0.030,
                    value=ALPHA,
                    step=0.001,
                    label="Force d'injection"
                )

                btn_embed = gr.Button(
                    "Appliquer le Watermark v3",
                    variant="primary"
                )

            with gr.Column():

                audio_out = gr.Audio(
                    label="Audio watermarké — WAV PCM 24-bit"
                )

                plot_out = gr.Image(
                    label="Analyse spectrale"
                )

                text_out = gr.Textbox(
                    label="Statut",
                    lines=6
                )

        btn_embed.click(
            embed_watermark,
            inputs=[
                audio_in,
                key_in,
                alpha_in
            ],
            outputs=[
                audio_out,
                plot_out,
                text_out
            ]
        )

    with gr.Tab("2. Vérifier / Détecter"):

        with gr.Row():

            with gr.Column():

                audio_verify = gr.Audio(
                    type="filepath",
                    label="Audio à vérifier"
                )

                key_verify = gr.Number(
                    value=42,
                    label="Clé secrète",
                    precision=0
                )

                btn_detect = gr.Button(
                    "Vérifier le Watermark",
                    variant="secondary"
                )

            with gr.Column():

                detect_out = gr.Textbox(
                    label="Résultat de détection",
                    lines=9
                )

        btn_detect.click(
            detect_watermark,
            inputs=[
                audio_verify,
                key_verify
            ],
            outputs=[detect_out]
        )

    gr.Markdown(
        """
### ⚠️ Validation

Tester séparément :

**ORIGINAL → doit rester NON CONFIRMÉ**

**WATERMARKÉ → doit être DÉTECTÉ**

Puis tester MP3, bruit, resampling, variation de volume,
low-pass/high-pass, time-stretch et pitch-shift.
"""
    )


# Required for ZeroGPU request handling.
demo.queue().launch()