import gradio as gr
from modules.code_assistant import CodeAssistant
from modules.docs_assistant import DocsAssistant
from modules.pdf_assistant import PDFAssistant
from modules.utils import load_css

def create_app():
    # Initialize assistants
    code_assistant = CodeAssistant()
    docs_assistant = DocsAssistant()
    pdf_assistant = PDFAssistant()
    
    with gr.Blocks(css=load_css()) as demo:
        gr.Markdown("# Enterprise RAG Assistant")
        
        with gr.Tabs() as tabs:
            # Code Assistant Tab
            with gr.Tab("Code Assistant", id=1):
                with gr.Row():
                    with gr.Column():
                        code_input = gr.Textbox(
                            label="Ask coding questions",
                            placeholder="Enter your coding question...",
                            lines=3
                        )
                        code_submit = gr.Button("Get Code Solution")
                    code_output = gr.Code(
                        label="Code Output",
                        language="python"
                    )
            
            # Documentation Assistant Tab
            with gr.Tab("Documentation Assistant", id=2):
                with gr.Row():
                    with gr.Column():
                        docs_input = gr.Textbox(
                            label="Documentation Query",
                            placeholder="Ask about technical documentation...",
                            lines=3
                        )
                        docs_file = gr.File(
                            label="Upload Documentation",
                            file_types=[".pdf", ".txt", ".md"]
                        )
                        docs_submit = gr.Button("Search Documentation")
                    docs_output = gr.Markdown()
            
            # PDF RAG Assistant Tab
            with gr.Tab("PDF Assistant", id=3):
                with gr.Row():
                    with gr.Column():
                        pdf_file = gr.File(
                            label="Upload PDF",
                            file_types=[".pdf"]
                        )
                        pdf_query = gr.Textbox(
                            label="Ask about the PDF",
                            placeholder="Enter your question about the PDF...",
                            lines=3
                        )
                        pdf_submit = gr.Button("Get Answer")
                    pdf_output = gr.Markdown()
        
        # Event handlers
        code_submit.click(
            code_assistant.generate_response,
            inputs=[code_input],
            outputs=[code_output]
        )
        
        docs_submit.click(
            docs_assistant.search_docs,
            inputs=[docs_input, docs_file],
            outputs=[docs_output]
        )
        
        pdf_submit.click(
            pdf_assistant.answer_query,
            inputs=[pdf_query, pdf_file],
            outputs=[pdf_output]
        )
    
    return demo

if __name__ == "__main__":
    app = create_app()
    app.launch()