import os
import sys
import tempfile
import subprocess
import logging

import spaces
import gradio as gr
import torch
from huggingface_hub import hf_hub_download
from scipy.io.wavfile import write
import numpy as np
from tqdm import tqdm

# ---------------------------------------------------------
# Лагіраванне
# ---------------------------------------------------------
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ---------------------------------------------------------
# 1. Клануем і падключаем coqui-ai-TTS (fork з падтрымкай BE)
# ---------------------------------------------------------
REPO_URL = "https://github.com/tuteishygpt/coqui-ai-TTS.git"
REPO_DIR = "coqui-ai-TTS"

if not os.path.exists(REPO_DIR):
    # Клануем fork з беларускай падтрымкай
    subprocess.run(
        ["git", "clone", REPO_URL, REPO_DIR],
        check=True,
    )

# Дадаём корань рэпазіторыя ў sys.path, каб "import TTS" бачыў пакет
repo_root = os.path.abspath(REPO_DIR)
if repo_root not in sys.path:
    sys.path.insert(0, repo_root)

from TTS.tts.configs.xtts_config import XttsConfig
from TTS.tts.models.xtts import Xtts
# ✅ Выкарыстоўваем токенайзер з coqui TTS замест underthesea
from TTS.tts.layers.xtts.tokenizer import (
    split_sentence,
    VoiceBpeTokenizer,
)

# ---------------------------------------------------------
# 2. Шляхі да файлаў мадэлі
# ---------------------------------------------------------
repo_id = "archivartaunik/BE_XTTS_V2_10ep250k"
model_dir = "./model"
os.makedirs(model_dir, exist_ok=True)

checkpoint_file = os.path.join(model_dir, "model.pth")
config_file = os.path.join(model_dir, "config.json")
vocab_file = os.path.join(model_dir, "vocab.json")
default_voice_file = os.path.join(model_dir, "voice.wav")
speakers_file = os.path.join(model_dir, "speakers_xtts.pth")

# Спампоўваем асноўныя файлы мадэлі
if not os.path.exists(checkpoint_file):
    hf_hub_download(repo_id, filename="model.pth", local_dir=model_dir)
if not os.path.exists(config_file):
    hf_hub_download(repo_id, filename="config.json", local_dir=model_dir)
if not os.path.exists(vocab_file):
    hf_hub_download(repo_id, filename="vocab.json", local_dir=model_dir)
if not os.path.exists(default_voice_file):
    hf_hub_download(repo_id, filename="voice.wav", local_dir=model_dir)

# ---------------------------------------------------------
# 2.1. Спампоўка speakers_xtts.pth (калі няма)
# ---------------------------------------------------------
if not os.path.exists(speakers_file):
    try:
        hf_hub_download(repo_id, filename="speakers_xtts.pth", local_dir=model_dir)
        logger.info("speakers_xtts.pth паспяхова спампаваны.")
    except Exception as e:
        logger.warning("Не атрымалася спампаваць speakers_xtts.pth: %s", e)

# ---------------------------------------------------------
# 3. Загрузка мадэлі і токенайзера
# ---------------------------------------------------------
config = XttsConfig()
config.load_json(config_file)
XTTS_MODEL = Xtts.init_from_config(config)
XTTS_MODEL.load_checkpoint(
    config,
    checkpoint_path=checkpoint_file,
    vocab_path=vocab_file,
    use_deepspeed=False,
)

device = "cuda:0" if torch.cuda.is_available() else "cpu"
XTTS_MODEL.to(device)
sampling_rate = int(XTTS_MODEL.config.audio["sample_rate"])

# Ініцыялізуем VoiceBpeTokenizer і падкладаем у мадэль
tokenizer = VoiceBpeTokenizer(vocab_file=vocab_file)
XTTS_MODEL.tokenizer = tokenizer

# ---------------------------------------------------------
# 4. Загрузка speakers_xtts.pth (гатовыя галасы)
# ---------------------------------------------------------
SPEAKERS_DB: dict[str, dict] = {}
SPEAKER_CHOICES: list[str] = ["— з аўдыё (reference) —"]

if os.path.exists(speakers_file):
    try:
        raw = torch.load(speakers_file, map_location="cpu")
        # магчымыя фарматы:
        # 1) {"speakers": {name: {...}}}
        # 2) {name: {...}}
        if isinstance(raw, dict) and "speakers" in raw and isinstance(raw["speakers"], dict):
            speakers_dict = raw["speakers"]
        else:
            speakers_dict = raw

        valid_count = 0
        if isinstance(speakers_dict, dict):
            for name, val in speakers_dict.items():
                if (
                    isinstance(val, dict)
                    and "gpt_cond_latent" in val
                    and "speaker_embedding" in val
                ):
                    SPEAKERS_DB[str(name)] = {
                        "gpt_cond_latent": val["gpt_cond_latent"],
                        "speaker_embedding": val["speaker_embedding"],
                    }
                    valid_count += 1

        if valid_count > 0:
            s_names = sorted(SPEAKERS_DB.keys())
            SPEAKER_CHOICES.extend(s_names)
            logger.info("Загружана %d галасоў з speakers_xtts.pth", valid_count)
        else:
            logger.warning(
                "speakers_xtts.pth загружаны, але не знойдзена ніводнага "
                "галасу з ключамі 'gpt_cond_latent' і 'speaker_embedding'."
            )
    except Exception as e:
        logger.exception("Памылка пры загрузцы speakers_xtts.pth: %s", e)
else:
    logger.warning("speakers_xtts.pth не знойдзены па шляху: %s", speakers_file)

# ---------------------------------------------------------
# 5. Функцыя TTS (з токенайзерам і падтрымкай speakers_xtts.pth)
# ---------------------------------------------------------
@spaces.GPU(duration=60)
def text_to_speech(
    belarusian_story: str,
    speaker_audio_file=None,
    preset_speaker: str = "— з аўдыё (reference) —",
):
    if not belarusian_story or belarusian_story.strip() == "":
        raise gr.Error("Увядзі хоць нейкі тэкст 🙂")

    # 3) атрыманне латэнтаў голасу:
    #    альбо з speakers_xtts.pth, альбо з reference audio
    use_preset = (
        isinstance(preset_speaker, str)
        and preset_speaker in SPEAKERS_DB
    )

    if use_preset:
        # 🔊 выкарыстоўваем загадзя падрыхтаваны голас
        try:
            sp = SPEAKERS_DB[preset_speaker]
            gpt_cond_latent = sp["gpt_cond_latent"].to(device)
            speaker_embedding = sp["speaker_embedding"].to(device)
        except Exception as e:
            logger.exception(
                "Памылка пры выкарыстанні галасу '%s' з speakers_xtts.pth", preset_speaker
            )
            raise gr.Error(
                f"Памылка пры выкарыстанні падрыхтаванага галасу '{preset_speaker}': {e}"
            )
    else:
        # калі аўдыё не перададзена — бярэм голас па змаўчанні
        ref_path = speaker_audio_file
        if not ref_path or (not isinstance(ref_path, str) and getattr(ref_path, "name", "") == ""):
            ref_path = default_voice_file

        try:
            gpt_cond_latent, speaker_embedding = XTTS_MODEL.get_conditioning_latents(
                audio_path=ref_path,
                gpt_cond_len=XTTS_MODEL.config.gpt_cond_len,
                max_ref_length=XTTS_MODEL.config.max_ref_len,
                sound_norm_refs=XTTS_MODEL.config.sound_norm_refs,
            )
        except Exception as e:
            logger.exception("Памылка пры атрыманні латэнтаў голасу з аўдыё")
            raise gr.Error(f"Памылка пры атрыманні латэнтаў голасу: {e}")

    # ✅ Замяняем sent_tokenize на split_sentence з токенайзера
    try:
        lang = "be"
        chunk_limit = tokenizer.char_limits.get(lang, 250)
        tts_texts = split_sentence(
            belarusian_story.strip(),
            lang=lang,
            text_split_length=chunk_limit,
        )
        tts_texts = [s.strip() for s in tts_texts if s and s.strip()]
        if not tts_texts:
            raise gr.Error("Не атрымалася падзяліць тэкст на сказы/чанкі.")
    except Exception as e:
        raise gr.Error(f"Памылка пры падзеле тэксту на сказы: {e}")

    all_wavs = []
    for text in tqdm(tts_texts):
        try:
            with torch.no_grad():
                wav_chunk = XTTS_MODEL.inference(
                    text=text,
                    language="be",
                    gpt_cond_latent=gpt_cond_latent,
                    speaker_embedding=speaker_embedding,
                    temperature=0.1,
                    length_penalty=1.0,
                    repetition_penalty=10.0,
                    top_k=10,
                    top_p=0.3,
                )
            all_wavs.append(wav_chunk["wav"])
        except Exception as e:
            raise gr.Error(f"Памылка пры генерырацыі аўдыя: {e}")

    try:
        out_wav = np.concatenate(all_wavs).astype(np.float32)
    except ValueError:
        raise gr.Error(
            "Немагчыма згенераваць аўдыё. Праверце ўваходны тэкст і аўдыёфайл."
        )
    except Exception as e:
        raise gr.Error(f"Памылка пры аб'яднанні аўдыя: {e}")

    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
    write(temp_file.name, sampling_rate, out_wav)

    return temp_file.name


# ---------------------------------------------------------
# 6. Прыклады
#    (3-яя калёнка — значэнне дропдауна: тут па змаўчанні
#     выкарыстоўваем reference-аудыё, таму ставім базавы пункт меню)
# ---------------------------------------------------------
examples = [
    [
        "Такім чынам, клуб стаў уладальнікам усіх існых на сёння міжнародных трафеяў паўднёваамерыканскага футболу.",
        "Nestarka.wav",
    ],
    [
        "Яму не ўдалося палепшыць фінансавае становішча каралеўства, а, наадварот, прыйшлося распрадаваць каштоўнасці чэшскай кароны.",
        "muzh.wav",
    ],
    [
        "Кампілятарамі называюць праграмы, якія пераўтвараюць код вышэйшага ўзроўню ў код ніжэйшага ўзроўню.",
        "chunk_100.wav",
    ],
    [
        "Акрамя таго, ліхачы аддаюць перавагу рэгі, хіп-хопу і класічнай музыцы.",
        "d1015.mp3",
    ],
    [
        "Позірк можа быць уважлівым, зацікаўленым, захопленым, але бывае і нахабным, задзірлівым, пагардлівым, напышлівым.",
        "donarka_ench.wav",
    ],
    [
        "Такі нават шчыры, ці што: родная мова народу – трасянка, а беларуская яму чужая!",
        "muzhcynski.wav",
    ],
]

analytics_script = """
<script async src="https://www.googletagmanager.com/gtag/js?id=G-TKDCRCQ7FK"></script>
<script>
  window.dataLayer = window.dataLayer || [];
  function gtag(){dataLayer.push(arguments);}
  gtag('js', new Date());

  gtag('config', 'G-TKDCRCQ7FK');
</script>
"""

# ---------------------------------------------------------
# 7. Графічны інтэрфейс Gradio з выбарам гатовых галасоў
# ---------------------------------------------------------
demo = gr.Blocks()
with demo:
    gr.HTML(analytics_script)
    gr.Interface(
        fn=text_to_speech,
        inputs=[
            gr.Textbox(lines=5, label="Тэкст на беларускай мове"),
            gr.Audio(
                type="filepath",
                label="Прыклад голасу (без іншых гукаў) не карацей 7 секунд",
                interactive=True,
            ),
            gr.Dropdown(
                label="Падрыхтаваныя галасы ",
                choices=SPEAKER_CHOICES,
                value=SPEAKER_CHOICES[0],
            ),
        ],
        outputs=gr.Audio(
            type="filepath",
            label="Згенераванае аўдыя",
        ),
        title="Belarusian TTS Demo",
        description="""
        <p>Увядзіце тэкст, і мадэль пераўтворыць яго ў аўдыя. Вы можаце выкарыстоўваць 
        голас па змаўчанні, абраць голас з прыкладаў унізе, абраць гатовы голас з 
        <code>спісу замежных галасоў</code> або загрузіць уласны файл / запісаць аўдыё.</p>

        <p><strong>Карысныя парады:</strong></p>
        <ul>
            <li>Выкарыстоўвайце прыклады з добрай якасцю, без іншых гукаў і з разнастайнай інтанацыяй — 
            ад яе моцна залежыць вынік.</li>
            <li>Інтанацыя таксама ўплывае на націскі.</li>
            <li>Прыклады галасоў могуць быць на любой мове.</li>
        </ul>

        <p>Каб палепшыць якасць мадэлі (націскі і дакладнасць кланавання галасоў), патрэбны дадатковыя датасэты. 
        Ахвяруйце свой голас праз <a href="https://Donar.by" target="_blank">Donar.by</a></p>

        <p>Далучайцеся да нашай беларускай суполкі ў ТГ, каб дапамагчы ці даведацца пра навіны ШІ: 
        <a href="https://t.me/SHibelChat" target="_blank">https://t.me/SHibelChat</a>.</p>

        <p><strong>Падтрымаць праект:</strong> <a href="https://buymeacoffee.com/tuteishygpt" target="_blank">Buy Me a Coffee</a></p>
        """,
        examples=examples,
        cache_examples=False,
    )

if __name__ == "__main__":
    demo.launch()