# Gradio UI for local chatbot testing with MCP client
import logging
import asyncio
import gradio as gr
from environs import Env

from mcp_client import MCPClient

logger = logging.getLogger(__name__)

env = Env()
env.read_env()

AUTH_USERNAME = env.str("AUTH_USERNAME", "admin")
AUTH_PASSWORD = env.str("AUTH_PASSWORD", "admin")

async def initialize_client():
    logger.info("Initializing MCP client...")
    client = MCPClient()
    logger.info("MCP client created")
    await client.initialize()
    result = client
    logger.info("MCP client initialized")
    return result


def launch_ui():
    with gr.Blocks(title="MCP Chatbot UI", fill_height=True, fill_width=True) as demo:
        gr.Markdown("## Dnext Product Catalog AI Assistant")
        
        # Authentication section
        with gr.Row() as auth_section:
            with gr.Column(scale=1):
                gr.Markdown("### Authentication")
                username = gr.Textbox(label="Username", placeholder="Enter your username")
                password = gr.Textbox(label="Password", placeholder="Enter your password", type="password")
                login_btn = gr.Button("Login", variant="primary")

        # Main content section (initially hidden)
        with gr.Row(visible=False) as main_content:
            with gr.Column(scale=2):
                gr.Markdown("### Chat Interface")
                chatbot = gr.Chatbot(height=500, label="Chatbot", type="messages")
                with gr.Row():
                    msg = gr.Textbox(
                        label="Enter your request",
                        placeholder="Type your message here...",
                        lines=2,
                        scale=4
                    )
                    submit_btn = gr.Button("Submit", scale=1)
                clear_btn = gr.Button("Clear Chat", variant="secondary")

            with gr.Column(scale=1):
                gr.Markdown("### Example Use Cases")
                with gr.Accordion("Product Creation Flow", open=False):
                    gr.Markdown("""
                    **Step 1:** Define a product characteristic
                    ```
                    define a product characteristic named 'characteristic1' with various attributes like 'Simple Value', 'MBPS' as the unit of measure, and a value of '1000'
                    ```
                    
                    **Step 2:** Define price structure
                    ```
                    define the price structure for the product 'price1'. The price is set to a one-time charge of 100 USD, with a unit of measure 'Day' and no recurring charges
                    ```
                    
                    **Step 3:** Create product specification
                    ```
                    create a product specification for 'specification1', which links the previously created characteristic (characteristic1) to the product specification
                    ```
                    
                    **Step 4:** Define complete product offering
                    ```
                    define a complete product offering. Tie the product specification (specification1), pricing (price1), and characteristics (characteristic1) that you previously created
                    ```
                    """)

        def authenticate(username, password):
            if username == AUTH_USERNAME and password == AUTH_PASSWORD:
                return {
                    auth_section: gr.Row.update(visible=False),
                    main_content: gr.Row.update(visible=True)
                }
            return {
                auth_section: gr.Row.update(visible=True),
                main_content: gr.Row.update(visible=False)
            }

        async def respond(user_message):
            bot_response = await client.invoke(user_message)
            return [
                {"role": "user", "content": user_message},
                {"role": "assistant", "content": bot_response},
            ], ""  # Return empty string to clear the input field


        async def clear_chat_and_memory():
            await client.clear_memory()
            return [], ""

        login_btn.click(
            fn=authenticate,
            inputs=[username, password],
            outputs=[auth_section, main_content]
        )
        submit_btn.click(fn=respond, inputs=[msg], outputs=[chatbot, msg])
        clear_btn.click(fn=clear_chat_and_memory, inputs=[], outputs=[chatbot, msg])

    return demo


async def main():
    try:
        logger.info("Starting application...")
        global client
        logger.info("Initializing MCP client...")
        client = await initialize_client()
        logger.info("MCP client initialized successfully")
        
        logger.info("Launching Gradio UI...")
        demo = launch_ui()
        logger.info("Gradio UI created successfully")
        
        logger.info(f"Starting Gradio server on {env.str('GRADIO_SERVER_HOST', '0.0.0.0')}:{env.int('GRADIO_SERVER_PORT', 7860)}")
        demo.launch(
            server_name=env.str("GRADIO_SERVER_HOST", "0.0.0.0"),
            server_port=env.int("GRADIO_SERVER_PORT", 7860),
            share=False,
            show_error=True,
            show_api=False
        )
        logger.info("Gradio server started successfully")
    except Exception as e:
        logger.error(f"Error starting application: {str(e)}")
        raise


if __name__ == "__main__":
    try:
        logger.info("Starting main application...")
        asyncio.run(main())
    except Exception as e:
        logger.error(f"Fatal error in main: {str(e)}")
        raise

# uv run app.py