"""Measure generated Kabyle content, speaker identity and predicted speech quality."""
from __future__ import annotations
import re
import unicodedata
from pathlib import Path
import numpy as np
import onnxruntime as ort
import requests
import torch
from huggingface_hub import hf_hub_download
from numpy.typing import NDArray
from pydantic import Field
from torch import Tensor, nn
from torch.nn import functional
from torchaudio.compliance.kaldi import fbank
from torchaudio.models import squim_subjective_base
from transformers import Wav2Vec2CTCTokenizer, Wav2Vec2FeatureExtractor, Wav2Vec2ForCTC
from .artifacts import atomic_binary, digest_file, run_lock
from .configuration import Specification
from .evaluation import ASR_ID, ASR_REVISION, edit_count, extract_input
from .waveforms import Waveform, measure_signal, resample_waveform
VOICE_ID = 'Wespeaker/wespeaker-voxceleb-resnet34-LM'
VOICE_REVISION = 'f0c48c298fd835726c27956a5d617bad7115627e'
VOICE_SHA256 = '7bb2f06e9df17cdf1ef14ee8a15ab08ed28e8d0ef5054ee135741560df2ec068'
MINIMUM_SPEAKER_SAMPLES = 400
QUALITY_SHA256 = 'dd36e1e4a5e3b36a5e7c32c936097b57b3221663a0529704171a3bb792699718'
INSTRUMENTS = {'asr': {'id': ASR_ID, 'revision': ASR_REVISION}, 'speaker': {'id': VOICE_ID, 'revision': VOICE_REVISION, 'sha256': VOICE_SHA256}, 'speech_quality': {'id': 'torchaudio/SQUIM_SUBJECTIVE', 'sha256': QUALITY_SHA256}}

def quality_checkpoint() -> Path:
    path = Path(torch.hub.get_dir()) / 'torchaudio/models/squim_subjective_bvcc_daps.pth'
    with run_lock(path.parent):
        if not path.exists():
            url = 'https://download.pytorch.org/torchaudio/models/squim_subjective_bvcc_daps.pth'
            with requests.get(url, stream=True, timeout=(10, 60)) as response:
                response.raise_for_status()
                with atomic_binary(path) as stream:
                    for block in response.iter_content(chunk_size=1024 * 1024):
                        stream.write(block)
        if digest_file(path) != QUALITY_SHA256:
            message = f'speech-quality checkpoint checksum mismatch at {path}'
            raise ValueError(message)
    return path

class Quality(Specification):
    hypothesis: str
    character_errors: int = Field(ge=0)
    characters: int = Field(ge=1)
    voice_cosine: float = Field(ge=-1, le=1)
    predicted_mos: float = Field(ge=1, le=5)
    duration_s: float = Field(gt=0)
    peak: float = Field(ge=0)
    rms_dbfs: float

    @property
    def cer(self) -> float:
        return self.character_errors / self.characters

def spoken_text(text: str) -> str:
    plain = unicodedata.normalize('NFC', re.sub('<\\|[^<>]*\\|>', '', text)).lower()
    result = ' '.join(''.join((char if char.isalnum() else ' ' for char in plain)).split())
    if not result:
        message = f'no spoken text in {text!r}'
        raise ValueError(message)
    return result

def speaker_features(waveform: Waveform) -> NDArray[np.float32]:
    audio = resample_waveform(waveform, 24000, 16000)
    if audio.size < MINIMUM_SPEAKER_SAMPLES:
        message = f'speaker verification needs 400 samples at 16 kHz, got {audio.size}'
        raise ValueError(message)
    features: Tensor = fbank(torch.from_numpy(audio).reshape(1, -1) * (1 << 15), num_mel_bins=80, frame_length=25, frame_shift=10, dither=0, sample_frequency=16000, window_type='hamming', use_energy=False)
    return np.asarray((features - features.mean(0)).unsqueeze(0).numpy(), dtype=np.float32)

class SpeechVerifier:

    def __init__(self, device: torch.device) -> None:
        self.device = device
        path = Path(hf_hub_download(VOICE_ID, 'voxceleb_resnet34_LM.onnx', revision=VOICE_REVISION, etag_timeout=30))
        if digest_file(path) != VOICE_SHA256:
            message = f'speaker-model checksum mismatch at {path}'
            raise ValueError(message)
        options = ort.SessionOptions()
        options.inter_op_num_threads, options.intra_op_num_threads = (1, 4)
        self.speaker = ort.InferenceSession(str(path), sess_options=options, providers=['CPUExecutionProvider'])
        self.recognizer: nn.Module = Wav2Vec2ForCTC.from_pretrained(ASR_ID, revision=ASR_REVISION)
        self.quality: nn.Module = squim_subjective_base()
        self.quality.load_state_dict(torch.load(quality_checkpoint(), weights_only=True, map_location='cpu'))
        for model in (self.recognizer, self.quality):
            nn.Module.to(model, device=device)
            nn.Module.train(model, False)
        self.extractor = Wav2Vec2FeatureExtractor.from_pretrained(ASR_ID, revision=ASR_REVISION)
        self.tokenizer = Wav2Vec2CTCTokenizer.from_pretrained(ASR_ID, revision=ASR_REVISION)

    def embed(self, waveform: Waveform) -> Tensor:
        output: object = self.speaker.run(['embs'], {'feats': speaker_features(waveform)})[0]
        if not isinstance(output, np.ndarray) or output.shape != (1, 256) or (not np.isfinite(output).all()):
            message = f'invalid speaker embedding from {VOICE_ID}'
            raise ValueError(message)
        tensor = torch.from_numpy(np.asarray(output, dtype=np.float32))
        if float(tensor.norm()) == 0:
            message = 'speaker encoder returned a zero vector'
            raise ValueError(message)
        return functional.normalize(tensor, dim=-1)

    @torch.inference_mode()
    def assess(self, waveform: Waveform, text: str, reference: Waveform) -> Quality:
        signal = measure_signal(waveform, 24000)
        if signal.rms_dbfs is None:
            message = 'generated speech is silent'
            raise ValueError(message)
        audio = resample_waveform(waveform, 24000, 16000)
        output = self.recognizer(input_values=extract_input(self.extractor, audio, self.device))
        logits = getattr(output, 'logits', None)
        if not isinstance(logits, Tensor):
            message = 'ASR returned no logits'
            raise TypeError(message)
        raw = str(self.tokenizer.decode(logits[0].argmax(-1).tolist()))
        hypothesis = spoken_text(raw) if raw.strip() else ''
        expected = spoken_text(text)
        score = self.quality(torch.from_numpy(audio).reshape(1, -1).to(self.device), torch.from_numpy(resample_waveform(reference, 24000, 16000)).reshape(1, -1).to(self.device))
        if not isinstance(score, Tensor):
            message = 'speech-quality model returned no MOS tensor'
            raise TypeError(message)
        return Quality(hypothesis=hypothesis, character_errors=edit_count(expected, hypothesis), characters=len(expected), voice_cosine=float((self.embed(waveform) * self.embed(reference)).sum().clamp(-1, 1)), predicted_mos=float(score[0]), duration_s=signal.samples / signal.rate_hz, peak=signal.peak, rms_dbfs=signal.rms_dbfs)