import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import gradio as gr

# Load model from Hugging Face Hub
model_name = "Sarthak279/chatbot"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()

def clean_text(text):
    return text.strip()

def chatbot(query):
    query = clean_text(query)
    input_ids = tokenizer(query, return_tensors="pt", max_length=250, truncation=True).to(device)
    with torch.no_grad():
        outputs = model.generate(
            input_ids["input_ids"],
            max_length=250,
            num_beams=5,
            early_stopping=True
        )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Gradio app with visible title + toggle button
with gr.Blocks(css="""
.chat-btn {
    position: fixed;
    bottom: 30px;
    right: 30px;
    background: #0d6efd;
    color: white;
    border: none;
    border-radius: 50%;
    width: 60px;
    height: 60px;
    font-size: 28px;
    box-shadow: 0 4px 10px rgba(0,0,0,0.3);
    cursor: pointer;
    z-index: 999;
}
.title-text {
    font-size: 1.8em;
    font-weight: bold;
    text-align: center;
    margin-top: 2rem;
    color: #0d6efd;
}
""") as demo:
    is_open = gr.State(False)

    gr.HTML('<div class="title-text">🩺 Sarthak Query Based Health Chatbot</div>')

    toggle_btn = gr.Button("💬", elem_classes="chat-btn")

    with gr.Column(visible=False) as chat_panel:
        with gr.Row():
            input_box = gr.Textbox(lines=2, placeholder="Type your message...", show_label=False, scale=4)
            send_btn = gr.Button("Send", scale=1)
        output_box = gr.Textbox(label="Bot Response", interactive=False)
        send_btn.click(chatbot, inputs=input_box, outputs=output_box)
        input_box.submit(chatbot, inputs=input_box, outputs=output_box)

    def toggle_chat(state):
        return gr.update(visible=not state), not state

    toggle_btn.click(toggle_chat, is_open, outputs=[chat_panel, is_open])

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