import gradio as gr
import requests

# GPT-2 대신 blenderbot-3B 모델 사용
API_URL = "https://api-inference.huggingface.co/models/facebook/blenderbot-3B"

# 전역 변수로 API 키 저장
api_key_store = {"key": None}

# API 키 저장 함수
def set_api_key(api_key):
    api_key_store["key"] = api_key
    return "✅ API Key Set Successfully!"

# LLM 질의 함수
def query_llm(message, history):
    api_key = api_key_store["key"]
    if not api_key:
        return "❌ Error: API key is missing. Please enter your Hugging Face API key first."

    headers = {"Authorization": f"Bearer {api_key}"}
    payload = {"inputs": message}

    response = requests.post(API_URL, headers=headers, json=payload)
    if response.status_code == 200:
        return response.json()[0]["generated_text"]
    else:
        return f"❌ Error: {response.status_code} - {response.json()}"

# Gradio UI (API 키 입력 + 채팅)
with gr.Blocks() as demo:
    with gr.Row():
        api_key_input = gr.Textbox(label="Hugging Face API Key", placeholder="Enter your API Key here", type="password")
        set_key_button = gr.Button("Set API Key")

    api_key_output = gr.Textbox(label="Status", interactive=False)  # API 키 상태 표시
    set_key_button.click(set_api_key, inputs=[api_key_input], outputs=[api_key_output])

    chatbot = gr.ChatInterface(fn=query_llm)

demo.launch()