import spaces
from kokoro import KModel, KPipeline
import gradio as gr
import os
import random
import torch

# GPU Kullanımını KAPATIYORUZ
CUDA_AVAILABLE = False  # GPU'yu tamamen devre dışı bırak
CHAR_LIMIT = 5000  # Karakter limiti (isteğe bağlı)

# Modelleri Yükle (SADECE CPU Kullanacak)
models = {False: KModel().to('cpu').eval()}  
pipelines = {lang_code: KPipeline(lang_code=lang_code, model=False) for lang_code in 'ab'}
pipelines['a'].g2p.lexicon.golds['kokoro'] = 'kˈOkəɹO'
pipelines['b'].g2p.lexicon.golds['kokoro'] = 'kˈQkəɹQ'

# GPU Fonksiyonunu Kaldırıyoruz
def forward_cpu(ps, ref_s, speed):
    return models[False](ps, ref_s, speed)

# Ses Üretme Fonksiyonu (Tamamen CPU Kullanır)
def generate_first(text, voice='af_heart', speed=1):
    text = text.strip()[:CHAR_LIMIT]  # Karakter limitini uygula
    pipeline = pipelines[voice[0]]
    pack = pipeline.load_voice(voice)

    for _, ps, _ in pipeline(text, voice, speed):
        ref_s = pack[len(ps)-1]
        try:
            audio = forward_cpu(ps, ref_s, speed)  # SADECE CPU Kullan
        except Exception as e:
            raise gr.Error(e)
        return (24000, audio.numpy()), ps
    return None, ''

# Arena API
def predict(text, voice='af_heart', speed=1):
    return generate_first(text, voice, speed)[0]

# Tokenizasyon Fonksiyonu
def tokenize_first(text, voice='af_heart'):
    pipeline = pipelines[voice[0]]
    for _, ps, _ in pipeline(text, voice):
        return ps
    return ''

# Akış Modu İçin CPU Kullanımı
def generate_all(text, voice='af_heart', speed=1):
    text = text.strip()[:CHAR_LIMIT]
    pipeline = pipelines[voice[0]]
    pack = pipeline.load_voice(voice)

    for _, ps, _ in pipeline(text, voice, speed):
        ref_s = pack[len(ps)-1]
        try:
            audio = forward_cpu(ps, ref_s, speed)  # SADECE CPU Kullan
        except Exception as e:
            raise gr.Error(e)
        yield 24000, audio.numpy()

# Rastgele Metin Getirme
random_texts = {'en': ["Hello, this is a test.", "How are you?", "Welcome to Kokoro!"]}
def get_random_text(voice):
    return random.choice(random_texts['en'])

# Ses Seçenekleri
CHOICES = {
    '🇺🇸 🚺 Heart ❤️': 'af_heart',
    '🇺🇸 🚺 Bella 🔥': 'af_bella',
    '🇺🇸 🚺 Nicole 🎧': 'af_nicole',
    '🇺🇸 🚹 Michael': 'am_michael',
    '🇬🇧 🚺 Emma': 'bf_emma',
    '🇬🇧 🚹 George': 'bm_george',
}
for v in CHOICES.values():
    pipelines[v[0]].load_voice(v)

# Token Bilgilendirme
TOKEN_NOTE = '''
💡 Özelleştirilmiş telaffuz için /slashes/ kullanabilirsiniz.
'''

# Arayüz - Generate Sekmesi
with gr.Blocks() as generate_tab:
    out_audio = gr.Audio(label='Çıktı Sesi', interactive=False, autoplay=True)
    generate_btn = gr.Button('Oluştur', variant='primary')
    with gr.Accordion('Çıktı Tokenleri', open=True):
        out_ps = gr.Textbox(interactive=False, show_label=False, info='Ses üretiminde kullanılan tokenler.')
        tokenize_btn = gr.Button('Tokenize', variant='secondary')

# Arayüz - Akış Sekmesi
with gr.Blocks() as stream_tab:
    out_stream = gr.Audio(label='Canlı Akış Sesi', interactive=False, streaming=True, autoplay=True)
    with gr.Row():
        stream_btn = gr.Button('Akışı Başlat', variant='primary')
        stop_btn = gr.Button('Durdur', variant='stop')

# Ana Arayüz
BANNER_TEXT = '***Kokoro*** **82M parametreli açık kaynak TTS modelidir.**'
with gr.Blocks() as app:
    with gr.Row():
        gr.Markdown(BANNER_TEXT, container=True)
    with gr.Row():
        with gr.Column():
            text = gr.Textbox(label='Metin Girişi', info=f"En fazla {CHAR_LIMIT} karakter")
            with gr.Row():
                voice = gr.Dropdown(list(CHOICES.items()), value='af_heart', label='Ses Seçimi')
            speed = gr.Slider(minimum=0.5, maximum=2, value=1, step=0.1, label='Konuşma Hızı')
            random_btn = gr.Button('Rastgele Metin', variant='secondary')
        with gr.Column():
            gr.TabbedInterface([generate_tab, stream_tab], ['Oluştur', 'Akış'])
    
    # Buton Bağlantıları
    random_btn.click(fn=get_random_text, inputs=[voice], outputs=[text])
    generate_btn.click(fn=generate_first, inputs=[text, voice, speed], outputs=[out_audio, out_ps])
    tokenize_btn.click(fn=tokenize_first, inputs=[text, voice], outputs=[out_ps])
    stream_event = stream_btn.click(fn=generate_all, inputs=[text, voice, speed], outputs=[out_stream])
    stop_btn.click(fn=None, cancels=stream_event)

# Uygulamayı Başlat
if __name__ == '__main__':
    app.queue().launch(show_api=True, ssr_mode=True)