import gradio as gr
import uuid
import logging
from src.controller.agent_cacher import AgentManager

# Configure logging
logging.basicConfig(level=logging.INFO, format="-->%(asctime)s [%(levelname)s] %(message)s")

# User credentials
VALID_USERS = {
    "ankit": "mlpass123",
    "demo": "test123"
}

def create_gradio_interface():
    manager = AgentManager()

    with gr.Blocks(title="AI Learning Assistant") as demo:
        auth_state = gr.State(False)  # store login status

        with gr.Column(visible=True) as login_section:
            gr.Markdown("## 🔐 Login to Continue")
            username = gr.Textbox(label="Username")
            password = gr.Textbox(label="Password", type="password")
            login_btn = gr.Button("Login")
            login_error = gr.Markdown(visible=False)

        with gr.Column(visible=False) as app_section:
            gr.Markdown("# 🧠 Learn with LlamaIndex Tools")

            with gr.Row():
                session_id = gr.Textbox(label="Session ID", value=str(uuid.uuid4()), visible=True)
                llm_selector = gr.Dropdown(
                    label="LLM Type",
                    choices=["Google", "OpenAI", "HuggingFace", "Mistral"],
                    value="Google",
                    interactive=True
                )

            with gr.Row():
                with gr.Column(scale=3):
                    chatbot = gr.Chatbot(label="Learning Dialog", height=500, type="messages")
                    query_input = gr.Textbox(label="Your Learning Query", placeholder="Ask about ML/DL algorithms...")
                    submit_btn = gr.Button("Submit")

                with gr.Column(scale=1):
                    gr.Markdown("### Tools Preview")
                    tool_output = gr.Textbox(label="Selected Tools", interactive=False)
                    response_output = gr.Textbox(label="Full Response", interactive=False, lines=10)

            def process_query(session_id_val, query, chat_history, llm_val):
                agent = manager.get_agent(session_id_val, llm_type=llm_val)
                chat_history, tools_used, response = agent.process_query(query)
                manager.save_agent(session_id_val, agent)
                return chat_history, tools_used, response, ""

            submit_btn.click(
                process_query,
                inputs=[session_id, query_input, chatbot, llm_selector],
                outputs=[chatbot, tool_output, response_output, query_input]
            )

            def load_history(session_id_val, llm_val):
                agent = manager.get_agent(session_id_val, llm_val)
                return agent.chat_history

            session_id.change(
                load_history,
                inputs=[session_id, llm_selector],
                outputs=chatbot
            )

        def check_login(username, password):
            if username == "admin" and password == "123":
                return "", True
            else:
                return "Invalid credentials", False

        login_btn.click(
            check_login,
            inputs=[username, password],
            outputs=[login_error, auth_state]
        )

        def toggle_ui(is_logged_in):
            return (
                gr.update(visible=is_logged_in),  # app_section
                gr.update(visible=not is_logged_in)  # login_section
            )

        auth_state.change(
            toggle_ui,
            inputs=auth_state,
            outputs=[app_section, login_section]
        )

    return demo

if __name__ == "__main__":
    interface = create_gradio_interface()
    interface.launch()