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

load_dotenv()
PASSWORD = os.getenv("ARMAAN_PASS")

# Load fast model
chatbot = pipeline("conversational", model="facebook/blenderbot-1B-distill")
system_prompt = "You are a helpful chatbot."

def chat_fn(message, history):
    convo = Conversation(f"{message}")
    result = chatbot(convo)
    return result.generated_responses[-1]

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

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

with gr.Blocks() as demo:
    gr.Markdown("# Chatbot (Fast CPU-Friendly)")
    
    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()