import gradio as gr
from core.rag_agent import RAGAgent
from core.document_manager import DocumentManager
import os

# Initialize components
doc_manager = DocumentManager()
rag_agent = None

def initialize_agent():
    """Initialize RAG agent lazily"""
    global rag_agent
    if rag_agent is None:
        rag_agent = RAGAgent()
    return rag_agent

def upload_files(files):
    """Handle file uploads"""
    if not files:
        return "No files selected", get_file_list()
    
    results = []
    for file in files:
        try:
            result = doc_manager.add_document(file.name)
            results.append(result)
        except Exception as e:
            results.append(f"Error processing {os.path.basename(file.name)}: {str(e)}")
    
    return "\n".join(results), get_file_list()

def get_file_list():
    """Get list of documents in the knowledge base"""
    try:
        files = doc_manager.list_documents()
        if not files:
            return "No documents in knowledge base"
        return "\n".join([f"• {f}" for f in files])
    except Exception as e:
        return f"Error listing files: {str(e)}"

def clear_database():
    """Clear all documents from the knowledge base"""
    try:
        result = doc_manager.clear_all()
        return result, get_file_list()
    except Exception as e:
        return f"Error clearing database: {str(e)}", get_file_list()

def chat_with_agent(message, history):
    """Handle chat interactions with the RAG agent"""
    if not message.strip():
        return history
    
    try:
        agent = initialize_agent()
        
        # Stream the agent's response
        response_text = ""
        for event in agent.agent_graph.stream(
            {"messages": [("user", message)]},
            agent.get_config(),
            stream_mode="values"
        ):
            if "messages" in event and len(event["messages"]) > 0:
                last_message = event["messages"][-1]
                if hasattr(last_message, "content"):
                    response_text = last_message.content
        
        if not response_text:
            response_text = "I apologize, but I couldn't generate a response. Please try again."
        
        return response_text
    
    except Exception as e:
        return f"Error: {str(e)}"

def reset_conversation():
    """Reset the conversation thread"""
    global rag_agent
    if rag_agent:
        rag_agent.reset_thread()
    return None  # Clear chat history

def create_gradio_ui():
    """Create the complete Gradio interface"""
    
    with gr.Blocks(title="RAG Agent with Agentic Memory", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 🤖 RAG Agent with Agentic Memory
        
        Upload documents and chat with an intelligent agent that uses:
        - 📚 **Local Knowledge Base** (ChromaDB)
        - 🔍 **Web Search** (Tavily)
        - 📖 **Wikipedia**
        - 🎓 **ArXiv** (Academic Papers)
        """)
        
        with gr.Tabs():
            # Documents Tab
            with gr.Tab("📄 Documents"):
                gr.Markdown("### Upload and Manage Documents")
                gr.Markdown("Upload PDF or Markdown files to add them to the knowledge base.")
                
                with gr.Row():
                    with gr.Column(scale=2):
                        file_upload = gr.File(
                            label="Upload Documents",
                            file_count="multiple",
                            file_types=[".pdf", ".md"]
                        )
                        upload_btn = gr.Button("📤 Add to Knowledge Base", variant="primary")
                        upload_status = gr.Textbox(label="Upload Status", lines=3)
                    
                    with gr.Column(scale=1):
                        file_list = gr.Textbox(
                            label="Documents in Knowledge Base",
                            lines=10,
                            value=get_file_list()
                        )
                        refresh_btn = gr.Button("🔄 Refresh List")
                        clear_btn = gr.Button("🗑️ Clear All Documents", variant="stop")
                
                # Connect document management buttons
                upload_btn.click(
                    fn=upload_files,
                    inputs=[file_upload],
                    outputs=[upload_status, file_list]
                )
                
                refresh_btn.click(
                    fn=get_file_list,
                    outputs=[file_list]
                )
                
                clear_btn.click(
                    fn=clear_database,
                    outputs=[upload_status, file_list]
                )
            
            # Chat Tab
            with gr.Tab("💬 Chat"):
                gr.Markdown("### Chat with Your Documents")
                gr.Markdown("Ask questions about your documents or any topic. The agent will search multiple sources.")
                
                chatbot = gr.Chatbot(
                    label="Conversation",
                    height=500,
                    show_label=True,
                    avatar_images=(None, "🤖")
                )
                
                with gr.Row():
                    msg = gr.Textbox(
                        label="Your Message",
                        placeholder="Ask me anything about your documents or general knowledge...",
                        scale=4
                    )
                    submit_btn = gr.Button("Send", variant="primary", scale=1)
                
                with gr.Row():
                    clear_chat_btn = gr.Button("🔄 Reset Conversation")
                    gr.Markdown("*Note: Resetting clears the conversation history*")
                
                # Chat interface
                chat_interface = gr.ChatInterface(
                    fn=chat_with_agent,
                    chatbot=chatbot,
                    textbox=msg,
                    submit_btn=submit_btn,
                    retry_btn=None,
                    undo_btn=None,
                    clear_btn=None
                )
                
                clear_chat_btn.click(
                    fn=reset_conversation,
                    outputs=[chatbot]
                )
        
        gr.Markdown("""
        ---
        ### 🔧 How it works:
        1. **Upload documents** in the Documents tab
        2. **Ask questions** in the Chat tab
        3. The agent will:
           - Analyze your query
           - Search relevant sources
           - Provide comprehensive answers with citations
        """)
    
    return demo

if __name__ == "__main__":
    demo = create_gradio_ui()
    demo.launch(share=False, server_name="127.0.0.1", server_port=7860)