# Copyright 2026 Umbra Tools (umbratools.com)
# SPDX-License-Identifier: Apache-2.0

"""Probe: can the narrator vmatrix be morphed (gender / creature) via instruct?

Read-only against the project; writes WAVs + measurements to an off-volume temp dir.
Uses the cached narrator vmatrix (prompt_slug) so voice identity is pinned.
"""
from __future__ import annotations

import io
import json
import time
from pathlib import Path

import librosa
import numpy as np
import soundfile as sf

import httpx

URL = "http://127.0.0.1:8881"
OUT = Path("/private/var/folders/y8/cbvhds297918lrxt2wd3btsr0000gn/T/opencode/probe_out")
OUT.mkdir(parents=True, exist_ok=True)

TEXT = (
    "The lantern guttered as she pushed the door open, and the cold air of the "
    "cavern rolled over her like the breath of something waiting."
)


def measure(wav: np.ndarray, sr: int) -> dict:
    y = wav.astype(np.float32)
    f0, _, _ = librosa.pyin(y, fmin=60.0, fmax=400.0, sr=sr)
    v = f0[~np.isnan(f0)]
    S = np.abs(librosa.stft(y, n_fft=2048, hop_length=512))
    cen = float(librosa.feature.spectral_centroid(S=S, sr=sr).mean())
    flat = float(librosa.feature.spectral_flatness(S=S).mean())
    # long-term average spectrum slope: brightness/tilth
    rms = float(20 * np.log10(np.sqrt((y**2).mean()) + 1e-12))
    if len(v):
        out = {
            "f0_median": round(float(np.median(v)), 1),
            "f0_sd": round(float(v.std()), 1),
            "f0_range": [round(float(v.min())), round(float(v.max()))],
            "semitones_vs_132hz": round(12 * np.log2(np.median(v) / 132.4), 1),
        }
    else:
        out = {"f0_median": None, "f0_sd": None, "f0_range": None, "semitones_vs_132hz": None}
    out.update(
        centroid_hz=round(cen),
        flatness=round(flat, 5),
        rms_db=round(rms, 1),
        dur_s=round(len(y) / sr, 2),
    )
    return out


REF = Path("/Volumes/1TB_APL_Storage/UmbraAudio/.umbra/voices/narrator_ref.wav")


def gen(label: str, *, instruct: str | None, seed: int = 1234) -> dict:
    data = {"text": TEXT, "sample_rate": 24000, "temperature": 0.7, "do_sample": True,
            "seed": seed}
    files = {"ref_audio": ("ref.wav", REF.read_bytes(), "audio/wav")}
    if instruct:
        data["voice_description"] = instruct
    t0 = time.time()
    r = httpx.post(f"{URL}/v1/tts", data=data, files=files, timeout=900)
    dt = time.time() - t0
    if r.status_code != 200:
        return {"label": label, "error": f"HTTP {r.status_code}: {r.text[:200]}"}
    wav, sr = sf.read(io.BytesIO(r.content), dtype="float32")
    if wav.ndim > 1:
        wav = wav.mean(axis=1)
    sf.write(OUT / f"{label}.wav", wav, sr)
    m = measure(wav, sr)
    m.update(label=label, seconds=round(dt, 1), instruct=instruct)
    print(json.dumps(m), flush=True)
    return m


CASES = [
    ("A_baseline", None),
    ("B_female", "A young woman's voice, bright and light, with a higher and clearer tone."),
    ("C_elderly_m", "An elderly man's voice, very deep, gravelly and slow."),
    ("D_dragon", "A huge ancient dragon's voice: a deep, rumbling, growling monster with "
                 "a metallic edge and a resonant, echoing chest tone."),
]

if __name__ == "__main__":
    results = [gen(label, instruct=instr) for label, instr in CASES]
    (OUT / "probe_results.json").write_text(json.dumps(results, indent=2))
    print("WROTE", OUT / "probe_results.json")