---

from pathlib import Path
import os
import sys
import inspect
import subprocess
import textwrap
import numpy as np
import torch
import soundfile as sf


# ============================================================
# CONFIG
# ============================================================

COSYVOICE_ROOT = Path(r"C:\Users\Geodezik\image_generation\CosyVoice_official").resolve()
MATCHA_ROOT = COSYVOICE_ROOT / "third_party" / "Matcha-TTS"
MODEL_DIR = COSYVOICE_ROOT / "pretrained_models" / "Fun-CosyVoice3-0.5B"
PROMPT_WAV = COSYVOICE_ROOT / "asset" / "zero_shot_prompt.wav"

OUT_DIR = COSYVOICE_ROOT / "_deep_diag_outputs"
OUT_DIR.mkdir(parents=True, exist_ok=True)

sys.path.insert(0, str(MATCHA_ROOT))
sys.path.insert(0, str(COSYVOICE_ROOT))


# ============================================================
# PRINT HELPERS
# ============================================================

def header(s):
    print()
    print("=" * 100)
    print(s)
    print("=" * 100)


def section(s):
    print()
    print("-" * 100)
    print(s)
    print("-" * 100)


# ============================================================
# ENV INFO
# ============================================================

header("ENV")

print("python:", sys.executable)
print("cwd:", os.getcwd())
print("COSYVOICE_ROOT:", COSYVOICE_ROOT)
print("MODEL_DIR:", MODEL_DIR)
print("PROMPT_WAV:", PROMPT_WAV)

print("torch:", torch.__version__)
print("torch file:", torch.__file__)
print("cuda available:", torch.cuda.is_available())

if torch.cuda.is_available():
    print("cuda device:", torch.cuda.get_device_name(0))
    print("cuda capability:", torch.cuda.get_device_capability(0))
    print("torch.version.cuda:", torch.version.cuda)

try:
    import torchaudio
    print("torchaudio:", torchaudio.__version__)
    print("torchaudio file:", torchaudio.__file__)
except Exception as e:
    print("torchaudio import error:", repr(e))

try:
    import onnxruntime as ort
    print("onnxruntime:", ort.__version__)
    print("onnxruntime providers:", ort.get_available_providers())
except Exception as e:
    print("onnxruntime import error:", repr(e))

try:
    import soxr
    print("soxr:", soxr.__version__)
except Exception as e:
    print("soxr not available:", repr(e))

try:
    import scipy
    print("scipy:", scipy.__version__)
except Exception as e:
    print("scipy not available:", repr(e))


# ============================================================
# IMPORT CHECK
# ============================================================

header("IMPORT CHECK")

import cosyvoice
import cosyvoice.cli.cosyvoice as cosyvoice_cli
import cosyvoice.cli.frontend as frontend
import cosyvoice.utils.file_utils as file_utils

print("cosyvoice:", cosyvoice.__file__)
print("cosyvoice cli:", cosyvoice_cli.__file__)
print("frontend:", frontend.__file__)
print("file_utils:", file_utils.__file__)

expected = str(COSYVOICE_ROOT / "cosyvoice").lower()
actual = str(Path(cosyvoice.__file__).parent).lower()

print("expected package dir:", expected)
print("actual package dir:", actual)

if not actual.startswith(expected):
    raise RuntimeError("WRONG COSYVOICE IMPORT PATH")


# ============================================================
# SIGNATURES
# ============================================================

header("SIGNATURES")

from cosyvoice.cli.cosyvoice import AutoModel

print("AutoModel:", AutoModel)
print("AutoModel signature:")
try:
    print(inspect.signature(AutoModel))
except Exception as e:
    print("signature error:", repr(e))

for name in [
    "CosyVoice",
    "CosyVoice2",
    "CosyVoice3",
]:
    if hasattr(cosyvoice_cli, name):
        cls = getattr(cosyvoice_cli, name)
        print()
        print(name, cls)
        try:
            print("__init__:", inspect.signature(cls.__init__))
        except Exception as e:
            print("__init__ signature error:", repr(e))


# ============================================================
# MODEL CONFIG
# ============================================================

header("MODEL CONFIG FILES")

for name in [
    "cosyvoice.yaml",
    "cosyvoice2.yaml",
    "cosyvoice3.yaml",
    "configuration.json",
    "config.json",
]:
    p = MODEL_DIR / name
    print()
    print(name, "exists:", p.exists())
    if p.exists():
        txt = p.read_text(encoding="utf-8", errors="ignore")
        print(txt[:4000])


# ============================================================
# PROMPT AUDIO INSPECT
# ============================================================

header("PROMPT AUDIO")

x, sr = sf.read(str(PROMPT_WAV), dtype="float32", always_2d=True)
print("sr:", sr)
print("shape:", x.shape)
print("duration:", x.shape[0] / sr)
print("min/max/std:", float(x.min()), float(x.max()), float(x.std()))


# ============================================================
# RESAMPLERS
# ============================================================

def resample_with_soxr_or_scipy(audio_1d_np: np.ndarray, sr: int, target_sr: int):
    if sr == target_sr:
        return audio_1d_np.astype(np.float32, copy=False)

    try:
        import soxr
        y = soxr.resample(audio_1d_np, sr, target_sr)
        return y.astype(np.float32, copy=False)
    except Exception as e:
        print("[resample] soxr failed:", repr(e))

    try:
        from scipy.signal import resample_poly
        from math import gcd
        g = gcd(sr, target_sr)
        up = target_sr // g
        down = sr // g
        y = resample_poly(audio_1d_np, up, down)
        return y.astype(np.float32, copy=False)
    except Exception as e:
        print("[resample] scipy failed:", repr(e))

    # final fallback: torchaudio
    import torchaudio.functional as AF
    t = torch.from_numpy(audio_1d_np[None, :].astype(np.float32))
    y = AF.resample(t, sr, target_sr).squeeze(0).numpy()
    return y.astype(np.float32, copy=False)


def make_loader(loader_name: str):
    def load_wav_no_torchcodec(wav, target_sr, min_sr=16000):
        wav = Path(wav)
        audio, sr = sf.read(str(wav), dtype="float32", always_2d=True)

        # [samples, channels] -> mono np
        mono = audio.mean(axis=1).astype(np.float32, copy=False)

        print()
        print(f"[{loader_name}]")
        print("  file:", wav)
        print("  target_sr:", target_sr)
        print("  original sr:", sr)
        print("  original shape:", audio.shape)
        print("  duration:", len(mono) / sr)
        print("  np min/max/std:", float(mono.min()), float(mono.max()), float(mono.std()))

        if sr < min_sr:
            raise ValueError(f"sample rate {sr} is lower than min_sr {min_sr}")

        if loader_name == "soundfile_soxr_scipy":
            mono = resample_with_soxr_or_scipy(mono, sr, target_sr)
            speech = torch.from_numpy(mono).unsqueeze(0)
        else:
            import torchaudio.functional as AF
            speech = torch.from_numpy(mono).unsqueeze(0)
            if sr != target_sr:
                speech = AF.resample(speech, sr, target_sr)

        print("  out shape:", tuple(speech.shape))
        print("  out duration:", speech.shape[1] / target_sr)
        print("  out min/max/std:", speech.min().item(), speech.max().item(), speech.std().item())

        return speech

    return load_wav_no_torchcodec


# ============================================================
# SAVE
# ============================================================

def save_outs(outs, output_path, sample_rate):
    outs = list(outs)
    print("chunks:", len(outs))

    if not outs:
        raise RuntimeError("No outputs")

    for i, o in enumerate(outs):
        t = o["tts_speech"]
        print(f"chunk {i}: shape={tuple(t.shape)} min={t.min().item():.6f} max={t.max().item():.6f} std={t.std().item():.6f}")

    wav = torch.cat([o["tts_speech"] for o in outs], dim=1)
    print("concat:", tuple(wav.shape))
    print("sr:", sample_rate)
    print("duration:", wav.shape[1] / sample_rate)
    print("min/max/std:", wav.min().item(), wav.max().item(), wav.std().item())

    audio = wav.detach().float().cpu().squeeze(0).numpy()
    audio = np.clip(audio, -1.0, 1.0)
    sf.write(str(output_path), audio, sample_rate)
    print("saved:", output_path)


# ============================================================
# MODEL BUILD PARAMS
# ============================================================

def build_model_with_best_effort(force_cpu=False, fp16_value=None):
    kwargs = {}
    sig = None

    try:
        sig = inspect.signature(AutoModel)
    except Exception:
        sig = None

    if sig is not None:
        params = sig.parameters

        if "fp16" in params and fp16_value is not None:
            kwargs["fp16"] = fp16_value

        if "device" in params and force_cpu:
            kwargs["device"] = "cpu"

        if "load_jit" in params:
            kwargs["load_jit"] = False

        if "load_trt" in params:
            kwargs["load_trt"] = False

        if "load_vllm" in params:
            kwargs["load_vllm"] = False

    print("AutoModel kwargs:", kwargs)

    if force_cpu:
        # Дополнительная страховка: прячем CUDA уже поздно, но некоторые компоненты смотрят на torch.cuda.
        print("force_cpu requested")

    model = AutoModel(model_dir=str(MODEL_DIR), **kwargs)
    return model


# ============================================================
# SINGLE TEST
# ============================================================

def run_test(test_name, loader_name="soundfile_torchaudio", force_cpu=False, fp16_value=None):
    header(f"RUN {test_name}")

    # patch loader
    loader = make_loader(loader_name)
    file_utils.load_wav = loader
    frontend.load_wav = loader

    print("loader:", loader_name)
    print("force_cpu:", force_cpu)
    print("fp16_value:", fp16_value)

    model = build_model_with_best_effort(force_cpu=force_cpu, fp16_value=fp16_value)

    print("model object:", type(model))
    print("sample_rate:", model.sample_rate)

    # Если force_cpu и объект имеет .model, пробуем насильно перевести torch-модули на CPU.
    if force_cpu:
        try:
            if hasattr(model, "model"):
                model.model = model.model.cpu()
                print("model.model moved to CPU")
        except Exception as e:
            print("moving model.model to CPU failed:", repr(e))

    outs = model.inference_zero_shot(
        "八百标兵奔北坡，北坡炮兵并排跑，炮兵怕把标兵碰，标兵怕碰炮兵炮。",
        "You are a helpful assistant.<|endofprompt|>希望你以后能够做的比我还好呦。",
        str(PROMPT_WAV),
        stream=False,
        speed=1.0,
    )

    output_path = OUT_DIR / f"{test_name}.wav"
    save_outs(outs, output_path, model.sample_rate)

    del model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


# ============================================================
# RUN TESTS
# ============================================================

tests = [
    # Базовый тест как раньше, но с новым скриптом.
    ("01_cuda_default_torchaudio_resample", "soundfile_torchaudio", False, None),

    # Принудительно fp16=False, если AutoModel это поддерживает.
    ("02_cuda_fp16_false_torchaudio_resample", "soundfile_torchaudio", False, False),

    # Другой resampler без torchaudio.functional.
    ("03_cuda_fp16_false_soxr_scipy_resample", "soundfile_soxr_scipy", False, False),
]

for test_name, loader_name, force_cpu, fp16_value in tests:
    try:
        run_test(test_name, loader_name, force_cpu, fp16_value)
    except Exception as e:
        print()
        print("[TEST FAILED]", test_name)
        print(type(e).__name__, str(e))


print()
print("=" * 100)
print("DONE")
print("=" * 100)
print("Outputs:", OUT_DIR)
print()
print("Теперь прослушай:")
for p in sorted(OUT_DIR.glob("*.wav")):
    print(" ", p)