import gradio as gr
import requests
from huggingface_hub import InferenceClient

# Initialize Hugging Face client
HF_MODEL = "mistralai/Mixtral-8x7B-Instruct-v0.1"
HF_TOKEN = "your_hugging_face_api_token"  # Replace with your token
client = InferenceClient(model=HF_MODEL, token=HF_TOKEN)

# Persistent bot knowledge state
bot_knowledge = {"dataset": None}


# Train chatbot by setting the dataset
def train_chatbot(dataset):
    bot_knowledge["dataset"] = dataset
    return "Chatbot trained successfully!"


# Chat function to process user input and generate bot responses
def chat_with_bot(history, user_input):
    if not bot_knowledge["dataset"]:
        return history + [{"role": "bot", "content": "No dataset loaded. Please train the bot first."}]

    # Append user input to the chat history
    history.append({"role": "user", "content": user_input})

    # Generate bot response
    prompt = f"{bot_knowledge['dataset']} {user_input}"
    try:
        response = client.text_generation(prompt=prompt, max_new_tokens=128)
        bot_response = response.get("generated_text", "Sorry, I couldn't generate a response.")
    except Exception as e:
        bot_response = f"Error generating response: {e}"

    # Append bot response to the history
    history.append({"role": "bot", "content": bot_response})
    return history


# Gradio Interface
with gr.Blocks(theme="default") as app:
    gr.Markdown("# **Intelligent Chatbot with Knowledge Training**")
    gr.Markdown(
        """
        Train a chatbot with custom datasets and interact with it dynamically. 
        The bot will persist knowledge from the dataset and answer questions accordingly.
        """
    )

    # Train chatbot section
    with gr.Row():
        chat_dataset = gr.Textbox(
            label="Dataset for Training",
            placeholder="Paste a dataset here to train the chatbot.",
            lines=5,
        )
        train_button = gr.Button("Train Chatbot")

    train_status = gr.Textbox(label="Training Status", interactive=False)

    # Chat section
    with gr.Row():
        chatbot = gr.Chatbot(
            label="Chat with Trained Bot",
            type="messages",
            avatar_user="https://example.com/user-avatar.png",
            avatar_bot="https://example.com/bot-avatar.png",
        )
        user_input = gr.Textbox(
            label="Your Message",
            placeholder="Type your message and press Enter...",
            lines=1,
        )

    # Train chatbot logic
    train_button.click(train_chatbot, inputs=[chat_dataset], outputs=[train_status])

    # Chat interaction logic
    user_input.submit(chat_with_bot, inputs=[chatbot, user_input], outputs=chatbot)

# Launch app
if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860)