app.py

import os
import gradio as gr
from src.auth.auth import handle_login
from src.auth.db import initialize_db
from dotenv import load_dotenv
from src.interface import Interface

# Load environment variables
initialize_db()
load_dotenv()

bot = None
bot_ready = False

def start_bot(userid, password, api_key_input):
    global bot_ready
    login_status = handle_login(userid, password, api_key_input)

    if "✅" in login_status:
        bot_ready = False
        return (
            "Login successful. Setting up your bot, please wait...",
            gr.update(visible=False),  # Hide login/registration section
            gr.update(visible=False)   # Still hide chat section for now
        )
    else:
        return (
            login_status,
            gr.update(visible=True),    # Keep login/register open
            gr.update(visible=False)    # Hide chat section
        )

def setup_bot():
    global bot, bot_ready
    try:
        bot = Interface()  # Interface will pick up API key from environment
        bot_ready = True
        return (
            gr.update(visible=False),  # Hide login/registration
            gr.update(visible=True)    # Show chat section
        )
    except Exception as e:
        return (
            gr.update(visible=True),
            gr.update(visible=False)
        )

def answer(message, history):
    if bot is None or not bot_ready:
        return "Please log in and wait for your bot to be ready."
    
    try:
        answer_md, tables_display, images_display, retrieved_display = bot.get_answer(message)

        combined_response = f"{answer_md}\n\n{tables_display}"

        if images_display:
            combined_response += "\n\n" + "\n\n".join(images_display)

        return combined_response
    except Exception as e:
        return f"An error occurred: {e}"

with gr.Blocks(fill_height=True, fill_width=True) as app:

    with gr.Column(visible=True) as login_register_section:
        gr.Markdown("# 🔐 MediBot Login & Registration")
        with gr.Tabs():
            with gr.TabItem("Login"):
                userid_login = gr.Textbox(label="UserID")
                password_login = gr.Textbox(label="Password", type="password")
                login_btn = gr.Button("Login")
                login_output = gr.Textbox(label="Login Status", interactive=False)

            with gr.TabItem("Register"):
                gr.Markdown("## 🔐 Enter Your Groq Cloud API Key")
                gr.Markdown(
                    "You can create an API key at [Groq Cloud Console](https://console.groq.com/keys)"
                )
                userid_register = gr.Textbox(label="UserID")
                password_register = gr.Textbox(label="Password", type="password")
                api_key_register = gr.Textbox(
                    label="Groq API Key",
                    type="password",
                    placeholder="sk-... (required)"
                )
                register_btn = gr.Button("Register")
                register_output = gr.Textbox(label="Registration Status", interactive=False)

    # Chat Section (Initially hidden)
    with gr.Column(visible=False) as chat_section:
        gr.ChatInterface(
            answer,
            title="🩺 Medico-Bot",
            examples=["briefly explain me about cancer", "types of skin diseases?"],
            flagging_options=['Like', 'Dislike']
        )

    # Login button connection
    login_btn.click(
        start_bot,
        inputs=[userid_login, password_login, api_key_register],
        outputs=[login_output, login_register_section, chat_section]
    ).then(
        setup_bot,
        inputs=[],
        outputs=[login_register_section, chat_section]
    )

    # Register button connection
    register_btn.click(
        start_bot,
        inputs=[userid_register, password_register, api_key_register],
        outputs=[register_output, login_register_section, chat_section]
    ).then(
        setup_bot,
        inputs=[],
        outputs=[login_register_section, chat_section]
    )

app.queue(concurrency_count=1, max_size=5)
app.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)