---

from pathlib import Path
import sys
import numpy as np
import torch
import soundfile as sf

# ============================================================
# PATHS
# ============================================================

COSYVOICE_ROOT = Path(r"C:\Users\Geodezik\image_generation\CosyVoice_official").resolve()
MATCHA_ROOT = COSYVOICE_ROOT / "third_party" / "Matcha-TTS"
OLD_BAD_ROOT = Path(r"C:\Users\Geodezik\image_generation\cosyvoice").resolve()

MODEL_DIR = COSYVOICE_ROOT / "pretrained_models" / "Fun-CosyVoice3-0.5B"
QWEN_DIR = MODEL_DIR / "CosyVoice-BlankEN"
PROMPT_WAV = COSYVOICE_ROOT / "asset" / "zero_shot_prompt.wav"

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

# ============================================================
# CLEAN IMPORTS
# ============================================================

sys.path = [
    p for p in sys.path
    if str(OLD_BAD_ROOT).lower() not in str(Path(p).resolve()).lower()
]

for name in list(sys.modules.keys()):
    if (
        name == "cosyvoice"
        or name.startswith("cosyvoice.")
        or name == "wetext"
        or name.startswith("wetext.")
    ):
        del sys.modules[name]

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

# ============================================================
# BASIC CHECKS
# ============================================================

print("=" * 80)
print("ENV")
print("=" * 80)
print("python:", sys.executable)
print("torch:", torch.__version__)
print("cuda:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("cuda device:", torch.cuda.get_device_name(0))
    print("torch cuda:", torch.version.cuda)

print("\nPATHS")
print("COSYVOICE_ROOT:", COSYVOICE_ROOT)
print("MODEL_DIR:", MODEL_DIR)
print("QWEN_DIR:", QWEN_DIR)
print("PROMPT_WAV:", PROMPT_WAV)

for p in [
    MODEL_DIR / "cosyvoice3.yaml",
    MODEL_DIR / "llm.pt",
    MODEL_DIR / "flow.pt",
    MODEL_DIR / "hift.pt",
    QWEN_DIR / "config.json",
    QWEN_DIR / "model.safetensors",
    QWEN_DIR / "tokenizer_config.json",
    QWEN_DIR / "vocab.json",
    QWEN_DIR / "merges.txt",
    PROMPT_WAV,
]:
    print("exists:", p, p.exists(), f"{p.stat().st_size / 1024 / 1024:.2f} MB" if p.exists() and p.is_file() else "")
    if not p.exists():
        raise FileNotFoundError(p)

# ============================================================
# TOKENIZER CHECK
# ============================================================

print("\n" + "=" * 80)
print("TOKENIZER CHECK")
print("=" * 80)

from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained(
    str(QWEN_DIR),
    trust_remote_code=True,
    local_files_only=True,
)

print("tokenizer:", tok.__class__)
print("vocab_size:", getattr(tok, "vocab_size", None))
print("len:", len(tok))
print("special_tokens_map:", tok.special_tokens_map)
print("all_special_tokens:", tok.all_special_tokens)

for s in [
    "<|endofprompt|>",
    "<|im_start|>",
    "<|im_end|>",
    "[breath]",
    "<|zh|>",
    "<|en|>",
    "<|ja|>",
    "<|ko|>",
    "<|yue|>",
]:
    tid = tok.convert_tokens_to_ids(s)
    print(repr(s), "->", tid)

ids = tok("<|endofprompt|>", add_special_tokens=False)["input_ids"]
tokens = tok.convert_ids_to_tokens(ids)

print("encode <|endofprompt|> ids:", ids)
print("encode <|endofprompt|> tokens:", tokens)

if len(ids) != 1:
    raise RuntimeError(
        "<|endofprompt|> is NOT a single token. "
        f"ids={ids}, tokens={tokens}. "
        "Модель/tokenizer всё ещё некорректны для CosyVoice3."
    )

# ============================================================
# NO FFMPEG / NO TORCHCODEC LOADER
# ============================================================

def resample_np(audio_1d: np.ndarray, sr: int, target_sr: int) -> np.ndarray:
    audio_1d = audio_1d.astype(np.float32, copy=False)

    if sr == target_sr:
        return audio_1d

    try:
        import soxr
        return soxr.resample(audio_1d, sr, target_sr).astype(np.float32, copy=False)
    except Exception:
        from scipy.signal import resample_poly
        from math import gcd
        g = gcd(sr, target_sr)
        return resample_poly(audio_1d, target_sr // g, sr // g).astype(np.float32, copy=False)


def load_wav_no_torchcodec(wav, target_sr, min_sr=16000):
    audio, sr = sf.read(str(wav), dtype="float32", always_2d=True)
    mono = audio.mean(axis=1).astype(np.float32, copy=False)

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

    mono = resample_np(mono, sr, target_sr)
    return torch.from_numpy(mono).unsqueeze(0)

# ============================================================
# COSYVOICE IMPORT + PATCH
# ============================================================

print("\n" + "=" * 80)
print("COSYVOICE IMPORT")
print("=" * 80)

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

file_utils.load_wav = load_wav_no_torchcodec
frontend.load_wav = load_wav_no_torchcodec

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

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

if not actual.startswith(expected):
    raise RuntimeError(f"Wrong cosyvoice import: {actual}")

# wetext check
try:
    import wetext
    from wetext import Normalizer
    print("wetext:", wetext.__file__)
    print("wetext Normalizer:", Normalizer)
except Exception as e:
    raise RuntimeError(f"wetext import failed: {e}")

# ============================================================
# MODEL LOAD
# ============================================================

print("\n" + "=" * 80)
print("LOAD MODEL")
print("=" * 80)

from cosyvoice.cli.cosyvoice import AutoModel

cosyvoice_model = AutoModel(
    model_dir=str(MODEL_DIR),
    fp16=False,
)

print("model:", type(cosyvoice_model))
print("sample_rate:", cosyvoice_model.sample_rate)
print("frontend:", type(cosyvoice_model.frontend))
print("frontend text_frontend:", getattr(cosyvoice_model.frontend, "text_frontend", None))

if getattr(cosyvoice_model.frontend, "text_frontend", None) != "wetext":
    raise RuntimeError("frontend.text_frontend is not 'wetext'")

# ============================================================
# SAVE HELPER
# ============================================================

def save_outs(name, outs):
    outs = list(outs)
    if not outs:
        raise RuntimeError("No output chunks")

    wav = torch.cat([o["tts_speech"] for o in outs], dim=1)
    audio = wav.detach().float().cpu().squeeze(0).numpy()
    audio = np.clip(audio, -1.0, 1.0)

    out_path = OUT_DIR / f"{name}.wav"
    sf.write(str(out_path), audio, cosyvoice_model.sample_rate, subtype="PCM_16")

    print("\nSAVED:", out_path)
    print("shape:", tuple(wav.shape))
    print("duration:", wav.shape[1] / cosyvoice_model.sample_rate)
    print("min/max/std:", wav.min().item(), wav.max().item(), wav.std().item())

    if wav.shape[1] / cosyvoice_model.sample_rate < 1.0:
        raise RuntimeError(f"{name} output too short")

    return out_path

# ============================================================
# TEST 1: OFFICIAL ZH ZERO-SHOT
# ============================================================

print("\n" + "=" * 80)
print("TEST 1: OFFICIAL ZH ZERO-SHOT")
print("=" * 80)

TEXT_ZH = "八百标兵奔北坡，北坡炮兵并排跑，炮兵怕把标兵碰，标兵怕碰炮兵炮。"
PROMPT_TEXT_ZH = "You are a helpful assistant.<|endofprompt|>希望你以后能够做的比我还好呦。"

out1 = save_outs(
    "01_official_zh_zero_shot",
    cosyvoice_model.inference_zero_shot(
        TEXT_ZH,
        PROMPT_TEXT_ZH,
        str(PROMPT_WAV),
        stream=False,
        speed=1.0,
    )
)

# ============================================================
# TEST 2: EN ZERO-SHOT WITH SAME CHINESE PROMPT TRANSCRIPT
# ============================================================

print("\n" + "=" * 80)
print("TEST 2: EN ZERO-SHOT")
print("=" * 80)

TEXT_EN = (
    "Hello. This is a simple English test. "
    "I am checking whether the speech is clear, continuous, and understandable."
)

out2 = save_outs(
    "02_en_zero_shot",
    cosyvoice_model.inference_zero_shot(
        TEXT_EN,
        PROMPT_TEXT_ZH,
        str(PROMPT_WAV),
        stream=False,
        speed=1.0,
    )
)

print("\n" + "=" * 80)
print("DONE")
print("=" * 80)
print("Listen files:")
print(out1)
print(out2)