# gradio_chatbot.py
import gradio as gr
import logging

logger = logging.getLogger(__name__)

from app import (
    add_user_message,
    add_loading_animation,
    generate_response,
    reset_conversation,
    load_conversation_state,
    remove_loading_animations,
    global_state_manager,
)


def load_custom_css():
    try:
        with open("styles.css", "r", encoding="utf-8") as css_file:
            css_content = css_file.read()
            logger.info(f"CSS loaded successfully, length: {len(css_content)} characters")
            return css_content
    except FileNotFoundError:
        logger.warning("styles.css file not found, using default styling")
        return ""
    except Exception as e:
        logger.warning(f"Error reading styles.css: {e}")
        return ""


def restore_state_on_page_access():
    """
    Restore conversation state when page loads or user navigates back.
    This ensures persistence across page navigation.
    """
    try:
        current_state = global_state_manager.get_conversation_state()
        chat_history = current_state.get('chat_history', [])
        conversation_state_data = current_state.get('conversation_state', [])
        
        logger.info(f"✓ Restored state: {len(chat_history)} messages in chat, {len(conversation_state_data)} in conversation")
        
        return chat_history, conversation_state_data
    except Exception as e:
        logger.error(f"Failed to restore state: {e}")
        return [], []


with gr.Blocks() as demo:
    custom_css = load_custom_css()
    if custom_css:
        gr.HTML(f'<style>{custom_css}</style>')

    conversation_state = gr.State([])
    
    gr.HTML('<div class="title-header"><h1>Mimir</h1></div>')
    
    with gr.Row():
        chatbot = gr.Chatbot(
            type="messages",
            show_copy_button=True,
            show_share_button=False,
            layout="bubble",
            autoscroll=True,
            avatar_images=None,
            elem_id="main-chatbot",
            scale=1,
            height="65vh",
            value=[], 
            latex_delimiters=[
                {"left": "$$", "right": "$$", "display": True},
                {"left": "$", "right": "$", "display": False},
            ]
        )
    
    with gr.Row(elem_classes=["input-controls"]):
        msg = gr.Textbox(
            placeholder="Ask me about math, research, study strategies, or any educational topic...",
            show_label=False,
            lines=6,
            max_lines=8,
            elem_classes=["input-textbox"],
            container=False,
            scale=4
        )
        with gr.Column(elem_classes=["button-column"], scale=1):
            send = gr.Button("Send", elem_classes=["send-button"], size="sm")
            clear = gr.Button("Clear", elem_classes=["clear-button"], size="sm")
    

    demo.load(
        fn=restore_state_on_page_access,
        outputs=[chatbot, conversation_state],  
        queue=False
    )
    
    msg.submit(
        add_user_message,
        inputs=[msg, chatbot, conversation_state],
        outputs=[msg, chatbot, conversation_state],
        show_progress="hidden",
        queue=True,
    ).then(
        add_loading_animation,
        inputs=[chatbot, conversation_state],
        outputs=[chatbot, conversation_state],
        show_progress="hidden",
        queue=True,
    ).then(
        generate_response, 
        inputs=[chatbot, conversation_state],
        outputs=[chatbot, conversation_state],
        show_progress="hidden",
        queue=True,
    )
    
    send.click(
        add_user_message,
        inputs=[msg, chatbot, conversation_state],
        outputs=[msg, chatbot, conversation_state],
        show_progress="hidden",
        queue=True,
    ).then(
        add_loading_animation,
        inputs=[chatbot, conversation_state],
        outputs=[chatbot, conversation_state],
        show_progress="hidden",
        queue=True,
    ).then(
        generate_response, 
        inputs=[chatbot, conversation_state],
        outputs=[chatbot, conversation_state],
        show_progress="hidden",
        queue=True,
    )
    
    clear.click(
        reset_conversation,
        outputs=[chatbot, conversation_state],
        show_progress="hidden"
    )


if __name__ == "__main__":
    logger.info("Running chatbot interface standalone")
    demo.queue(default_concurrency_limit=1)
    demo.launch(server_name="0.0.0.0", server_port=7860)