import os
import gradio as gr
from transformers import pipeline
from dotenv import load_dotenv

# Load environment variables from .env
load_dotenv()
PASSWORD = os.getenv("ARMAAN_PASS")

# Load chatbot pipeline
chatbot = pipeline("text-generation", model="HuggingFaceH4/zephyr-7b-beta")

# Global system prompt (admin-editable)
system_prompt = "You are a helpful assistant."

# Chat function
def chat_fn(message, history):
    full_prompt = f"{system_prompt}\nUser: {message}\nBot:"
    response = chatbot(full_prompt, max_new_tokens=200, do_sample=True, temperature=0.7)[0]["generated_text"]
    reply = response.split("Bot:")[-1].strip()
    return reply

# Password check
def check_password(pass_input):
    if pass_input == PASSWORD:
        return gr.update(visible=True), ""
    else:
        return gr.update(visible=False), "Incorrect password"

# Prompt update
def update_prompt(new_prompt):
    global system_prompt
    system_prompt = new_prompt
    return "Prompt updated successfully."

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# ChatGPT-Style Chatbot")

    with gr.Tab("Chat"):
        gr.ChatInterface(fn=chat_fn)

    with gr.Tab("Train"):
        with gr.Column():
            pass_input = gr.Textbox(label="Enter Admin Password", type="password")
            login_btn = gr.Button("Login")
            error_text = gr.Markdown("", visible=False)

            with gr.Group(visible=False) as admin_panel:
                new_prompt = gr.Textbox(label="New System Prompt", lines=4)
                update_btn = gr.Button("Update Prompt")
                success_msg = gr.Markdown("")

        login_btn.click(fn=check_password, inputs=pass_input, outputs=[admin_panel, error_text])
        update_btn.click(fn=update_prompt, inputs=new_prompt, outputs=success_msg)

demo.launch()