import spaces  # imported first, as ZeroGPU requires

import os
import sys
import uuid
from pathlib import Path
from contextlib import contextmanager

import ruptures as rpt

import numpy as np
import soundfile as sf
import torch
import gradio as gr
import librosa

from pyharp.core import ModelCard, build_endpoint
from pyharp.media.audio import save_audio
from pyharp import LabelList, AudioLabel, OutputLabel, SampleRate, Channels
from pyharp.tags import Category
from audiotools import AudioSignal
from audioseal import AudioSeal

LOUDNESS_DB = -16.
SAMPLE_RATE = 48_000
ENCODEC_SAMPLE_RATE = 16_000
AUDIOSEAL_SAMPLE_RATE = 16_000

model_card = ModelCard(
    name="AudioSeal",
    description=(
        "Meta AudioSeal watermarking generation and detection model The "
        "watermark is applied under 16kHz. The model is fixed-rate and mono, so "
        "the result is resampled back to the input's rate and copied across its "
        "channels: the output has the format it came in with, but no stereo "
        "image. Repository: https://github.com/facebookresearch/audioseal "
        "Paper: https://arxiv.org/abs/2401.17264"
    ),
    author="Robin San Roman, Pierre Fernandez, Alexandre Défossez, Teddy Furon, Tuan Tran, Hady Elsahar",
    tags=[Category.UTILITY, SampleRate(16000), Channels(1), 'domain:speech']
)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# A lossless input keeps its container. A lossy one comes back as WAV, since encoding the result
# lossily again would expose the artifacts of the first encode and add a second set.
LOSSLESS = {".wav", ".flac"}

# pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the
# server, so both models are loaded there alone
if __name__ != "__main__":
    print("Initializing AudioSeal model...")
    generator = AudioSeal.load_generator("audioseal_wm_16bits")
    detector = AudioSeal.load_detector("audioseal_detector_16bits")
    generator.eval().to(DEVICE)
    detector.eval().to(DEVICE)

def load_audio(audio_path):
    try:
        wav, sr = librosa.load(audio_path, mono=True)
        return wav, sr
    
    except Exception as e:
        print(f"Audio preprocessing failed: {e}")
        raise ValueError(f"Failed to load audio: {str(e)}")
    
@torch.no_grad()
def split_bands(signal: AudioSignal, sample_rate: float = ENCODEC_SAMPLE_RATE):
    nyq = sample_rate // 2
    high = signal.clone().high_pass(cutoffs=int(nyq * 0.95), zeros=51)
    low  = signal.clone().low_pass(cutoffs=int(nyq * 1.05), zeros=51)
    loud_db = low.loudness()
    low = low.resample(sample_rate)
    return low, high, loud_db

@torch.no_grad()
def merge_bands(low, high, loud_db):
    low = low.clone().to(high.device).resample(high.sample_rate)
    low.audio_data = low.audio_data[..., :high.signal_length]
    low.audio_data = torch.nn.functional.pad(
        low.audio_data, (0, max(0, high.signal_length - low.signal_length))
    )
    return low.normalize(loud_db) + high
    
@torch.no_grad()
def embed(signal: AudioSignal, embedder: torch.nn.Module):
    orig_ch, orig_sr = signal.num_channels, signal.sample_rate
    sig = signal.clone().resample(SAMPLE_RATE)
    if orig_ch > 1:
        b, c, n = sig.audio_data.shape
        sig.audio_data = sig.audio_data.reshape(b * c, 1, n)
    low, high, loud = split_bands(sig.clone(), AUDIOSEAL_SAMPLE_RATE)
    wm = embedder.get_watermark(low.audio_data, AUDIOSEAL_SAMPLE_RATE)
    low.audio_data = low.audio_data + wm
    merged = merge_bands(low, high, loud)
    if orig_ch > 1:
        b2, c2, n2 = merged.audio_data.shape
        merged.audio_data = merged.audio_data.reshape(-1, orig_ch * c2, n2)
    return merged.resample(orig_sr)

@torch.no_grad()
def detect(signal: AudioSignal, detector: torch.nn.Module):
    sig = signal.clone().to_mono().resample(AUDIOSEAL_SAMPLE_RATE)
    result, _ = detector.forward(sig.audio_data, sample_rate=AUDIOSEAL_SAMPLE_RATE)
    return result[0, 1, :].detach().cpu().numpy()

def match_input_format(output_path: str, input_path: str) -> str:
    """
    Rewrite an output file at the sample rate and channel count it came in with.

    The models here are fixed-rate and mono, and hand the writing to vendored
    upstream code, so the format is corrected afterwards rather than threaded
    through inference. Duplicated channels are copies, not a recovered stereo
    image; the model card says so.

    Args:
        output_path (str): File the model wrote. Overwritten in place.
        input_path (str): File the request supplied.

    Returns:
        output_path (str): The same path, for use in a return statement.
    """

    import librosa
    import numpy as np
    import soundfile

    want = soundfile.info(input_path)
    have = soundfile.info(output_path)

    if have.samplerate == want.samplerate and have.channels == want.channels:
        return output_path

    data, rate = soundfile.read(output_path, always_2d=True)

    if rate != want.samplerate:
        data = librosa.resample(
            data.T, orig_sr=rate, target_sr=want.samplerate
        ).T

    if data.shape[1] == 1 and want.channels > 1:
        data = np.tile(data, (1, want.channels))
    elif data.shape[1] > 1 and want.channels == 1:
        data = data.mean(axis=1, keepdims=True)

    soundfile.write(output_path, data, want.samplerate)

    return output_path


def process_fn(inp_audio, option_text):
    # Checked before the GPU is reserved, so a request the model cannot take spends none of the
    # caller's ZeroGPU allowance
    if not inp_audio:
        raise gr.Error("Provide the audio to watermark, or to test for a watermark.")
    try:
        info = sf.info(inp_audio)
    except Exception as exc:
        raise gr.Error(f"Could not read the audio file: {exc}") from exc
    if info.frames == 0:
        raise gr.Error("The audio file is empty.")

    output_audio, output_labels = watermark(inp_audio, option_text)
    if output_audio == inp_audio:
        return output_audio, output_labels

    # The format is matched outside the GPU call, which it does not need
    output_audio = match_input_format(output_audio, inp_audio)
    input_ext = Path(inp_audio).suffix.lower()
    if input_ext in LOSSLESS and input_ext != ".wav":
        data, rate = sf.read(output_audio, always_2d=True)
        output_audio = str(Path(output_audio).with_suffix(input_ext))
        sf.write(output_audio, data, rate)
    return output_audio, output_labels


@spaces.GPU(duration=120)
def watermark(inp_audio, option_text):
    audio_np, sr = load_audio(inp_audio)

    print(f"sr: {sr}, audio shape: {audio_np.shape}")
    if audio_np.ndim == 1:
        audio_np = audio_np[None, None, :]
    else:
        audio_np = np.transpose(audio_np, (1, 0))[None, ...]

    print(f"formatted audio: {audio_np.shape}")

    ori_sig = AudioSignal(torch.from_numpy(audio_np).float(), sample_rate=sr)
    orig_loud = ori_sig.loudness()
    sig = ori_sig.to_mono().resample(SAMPLE_RATE).normalize(LOUDNESS_DB).ensure_max_of_audio()
    sig = sig.to(DEVICE)

    output_labels = LabelList()
    
    if option_text == "Generate Watermark":
        with torch.no_grad():
            wm_sig = embed(sig.clone(), generator).normalize(orig_loud).ensure_max_of_audio()
            output_labels.labels.append(
                AudioLabel(
                    t = 0,
                    label = "watermark: 1.0",
                    duration = wm_sig.duration,
                    description = f"watermark confidence: 1.0, start: 0.0s, end: {wm_sig.duration:.2f}s",
                    color = OutputLabel.rgb_color_to_int(255, 0, 0),
                    amplitude = 1.0
                )
            )
            return save_audio(wm_sig.cpu()), output_labels
    else:
        with torch.no_grad():
            scores = detect(sig, detector) # AUDIOSEAL_SAMPLE_RATE
        N = len(scores)

        hop = int(0.01 * AUDIOSEAL_SAMPLE_RATE)
        avg_curve = []
        for i in range(0, N, hop):
            seg = scores[i:i+hop]
            value = np.mean(seg)
            avg_curve.append(value)

        avg_curve = np.array(avg_curve)
        print(avg_curve.shape)

        min_size = max(2, int(0.25 * AUDIOSEAL_SAMPLE_RATE))
        bkps = rpt.Pelt(model="l2", min_size=1).fit_predict(avg_curve, 1.0)

        t0 = 0
        for t1 in bkps:
            print(t0, t1)
            seg = avg_curve[t0:t1]
            value = seg.mean()
            output_labels.labels.append(
                AudioLabel(
                    t = (t0 / 100),
                    label = f"watermark: {value:.2f}",
                    duration = (t1 - t0) / 100,
                    description = f"watermark confidence: {value:.2f}, start: {(t0 / 100):.2f}s, end: {(t1 / 100):.2f}s",
                    color = OutputLabel.rgb_color_to_int(int(value * 255), int((1 - value) * 255), 0),
                    amplitude = value * 2 - 1
                )
            )
            t0 = t1

        return inp_audio, output_labels


if __name__ == "__main__":
    with gr.Blocks() as app:
        gr.Markdown("## Meta AudioSeal Watermarking")

        # Inputs
        input_audio = gr.Audio(
            label="Input Audio",
            type="filepath",
            sources=["upload", "microphone"]
        ).set_info("Speech to watermark, or an already-watermarked file to test.")

        option_dropdown = gr.Dropdown(
            ["Generate Watermark", "Detect Watermark"],
            value='Generate Watermark',
            label='Option',
            info='Model Options'
        )

        # Outputs
        output_wav = gr.Audio(
            type="filepath",
            label="Watermarked Speech"
        ).set_info("The input with the watermark embedded, which should be inaudible.")
        output_label = gr.JSON(label="Watermark Confidence").set_info(
            "How sure the detector is that a watermark is present: one label spanning the "
            "clip after watermarking, or one per segment after detection."
        )

        _ = build_endpoint(
            model_card=model_card,
            input_components=[
                input_audio,
                option_dropdown
            ],
            output_components=[
                output_wav,
                output_label
            ],
            process_fn=process_fn
        )

    app.queue().launch(pwa=True, show_error=True)