import os
import gradio as gr
from datasets import load_dataset
from huggingface_hub import InferenceClient

# 1. Memuat dataset Anda
try:
    dataset = load_dataset("fareldevelopers/gpt", split="train")
except Exception as e:
    dataset = None
    print(f"Gagal memuat dataset: {e}")

# 2. Mengambil token otomatis dari sistem Hugging Face Space
hf_token = os.environ.get("HF_TOKEN")

# Inisialisasi Klien AI dengan Token Otentikasi
# Menggunakan model open-source Zephyr (alternatif yang sangat stabil untuk inference)
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta", token=hf_token)

def respons_chat(pesan_user, history):
    # Pengecekan Tahap 1: Cari di dataset terlebih dahulu
    if dataset:
        kolom_input = 'instruction'
        kolom_output = 'response'
        
        for baris in dataset:
            if kolom_input in baris and pesan_user.lower().strip() == str(baris[kolom_input]).lower().strip():
                return baris.get(kolom_output, "Data ditemukan, tetapi kolom respon kosong.")

    # Pengecekan Tahap 2: Jika tidak ada di dataset, lempar ke Model AI
    try:
        messages = []
        for user_msg, ai_msg in history:
            messages.append({"role": "user", "content": user_msg})
            messages.append({"role": "assistant", "content": ai_msg})
        messages.append({"role": "user", "content": pesan_user})
        
        response = ""
        for message in client.chat_completion(
            messages,
            max_tokens=512,
            stream=True,
            temperature=0.7,
        ):
            token = message.choices[0].delta.content
            if token:
                response += token
        return response

    except Exception as e:
        return f"Maaf, terjadi kesalahan pada server AI: {str(e)}"

# 3. Membuat Antarmuka Gradio
app = gr.ChatInterface(
    fn=respons_chat,
    title="Custom Hybrid GPT",
    description="Asisten AI yang terintegrasi dengan dataset lokal dan model generatif publik."
)

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