import gradio as gr
from openai import OpenAI
import torch
from models import build_model
from kokoro import generate

# Fungsi untuk menginisialisasi model Kokoro-82M
def init_kokoro_model():
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model = build_model('kokoro-v0_19.pth', device)
    voice_name = 'af'  # Default voice
    voicepack = torch.load(f'voices/{voice_name}.pt', weights_only=True).to(device)
    return model, voicepack, voice_name

# Inisialisasi model Kokoro-82M
MODEL, VOICEPACK, VOICE_NAME = init_kokoro_model()

# Fungsi untuk menghasilkan respons dari Deepseek
def get_deepseek_response(api_key, user_input):
    client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com")
    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=[
            {"role": "system", "content": "You are a helpful assistant"},
            {"role": "user", "content": user_input},
        ],
        max_tokens=1024,
        temperature=0.7,
        stream=False
    )
    return response.choices[0].message.content

# Fungsi untuk mengubah teks menjadi suara menggunakan Kokoro-82M
def text_to_speech(text):
    audio, _ = generate(MODEL, text, VOICEPACK, lang=VOICE_NAME[0])
    return (24000, audio)

# Fungsi utama yang akan dipanggil oleh Gradio
def chat_with_ai(api_key, user_input):
    # Mendapatkan respons dari Deepseek
    response_text = get_deepseek_response(api_key, user_input)
    
    # Mengubah respons teks menjadi suara
    audio_output = text_to_speech(response_text)
    
    return response_text, audio_output

# Membuat antarmuka Gradio
with gr.Blocks() as demo:
    gr.Markdown("# AI Chatbot with Deepseek and Kokoro-82M")
    
    with gr.Row():
        api_key_input = gr.Textbox(label="Deepseek API Key", placeholder="Masukkan API Key Anda di sini")
        user_input = gr.Textbox(label="Input Anda", placeholder="Ketik pesan Anda di sini")
    
    with gr.Row():
        text_output = gr.Textbox(label="Respons AI", interactive=False)
        audio_output = gr.Audio(label="Suara AI")
    
    submit_button = gr.Button("Kirim")
    
    submit_button.click(
        fn=chat_with_ai,
        inputs=[api_key_input, user_input],
        outputs=[text_output, audio_output]
    )

# Menjalankan aplikasi Gradio
demo.launch()