import gradio as gr
import os
import json
import uuid
import asyncio
from datetime import datetime
from typing import List, Dict, Any, Optional, Generator
import logging

# Import required libraries
from huggingface_hub import InferenceClient
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.docstore.document import Document

# Import document parsers
import PyPDF2
from pptx import Presentation
import pandas as pd
from docx import Document as DocxDocument
import io

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Get HuggingFace token from environment
HF_TOKEN = os.getenv("hf_token")
if not HF_TOKEN:
    raise ValueError("HuggingFace token not found in environment variables")

# Initialize HuggingFace Inference Client
client = InferenceClient(model="meta-llama/Llama-3.1-8B-Instruct", token=HF_TOKEN)

# Initialize embeddings
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

class MCPMessage:
    """Model Context Protocol Message Structure"""
    def __init__(self, sender: str, receiver: str, msg_type: str,
                 trace_id: str = None, payload: Dict = None):
        self.sender = sender
        self.receiver = receiver
        self.type = msg_type
        self.trace_id = trace_id or str(uuid.uuid4())
        self.payload = payload or {}
        self.timestamp = datetime.now().isoformat()
    
    def to_dict(self):
        return {
            "sender": self.sender,
            "receiver": self.receiver,
            "type": self.type,
            "trace_id": self.trace_id,
            "payload": self.payload,
            "timestamp": self.timestamp
        }

class MessageBus:
    """In-memory message bus for MCP communication"""
    def __init__(self):
        self.messages = []
        self.subscribers = {}
    
    def publish(self, message: MCPMessage):
        """Publish message to the bus"""
        self.messages.append(message)
        logger.info(f"Message published: {message.sender} -> {message.receiver} [{message.type}]")
        
        # Notify subscribers
        if message.receiver in self.subscribers:
            for callback in self.subscribers[message.receiver]:
                callback(message)
    
    def subscribe(self, agent_name: str, callback):
        """Subscribe agent to receive messages"""
        if agent_name not in self.subscribers:
            self.subscribers[agent_name] = []
        self.subscribers[agent_name].append(callback)

# Global message bus
message_bus = MessageBus()

class IngestionAgent:
    """Agent responsible for document parsing and preprocessing"""
    
    def __init__(self, message_bus: MessageBus):
        self.name = "IngestionAgent"
        self.message_bus = message_bus
        self.message_bus.subscribe(self.name, self.handle_message)
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200
        )
    
    def handle_message(self, message: MCPMessage):
        """Handle incoming MCP messages"""
        if message.type == "INGESTION_REQUEST":
            self.process_documents(message)
    
    def parse_pdf(self, file_path: str) -> str:
        """Parse PDF document"""
        try:
            with open(file_path, 'rb') as file:
                pdf_reader = PyPDF2.PdfReader(file)
                text = ""
                for page in pdf_reader.pages:
                    text += page.extract_text()
                return text
        except Exception as e:
            logger.error(f"Error parsing PDF: {e}")
            return ""
    
    def parse_pptx(self, file_path: str) -> str:
        """Parse PPTX document"""
        try:
            prs = Presentation(file_path)
            text = ""
            for slide in prs.slides:
                for shape in slide.shapes:
                    if hasattr(shape, "text"):
                        text += shape.text + "\n"
            return text
        except Exception as e:
            logger.error(f"Error parsing PPTX: {e}")
            return ""
    
    def parse_csv(self, file_path: str) -> str:
        """Parse CSV document"""
        try:
            df = pd.read_csv(file_path)
            return df.to_string()
        except Exception as e:
            logger.error(f"Error parsing CSV: {e}")
            return ""
    
    def parse_docx(self, file_path: str) -> str:
        """Parse DOCX document"""
        try:
            doc = DocxDocument(file_path)
            text = ""
            for paragraph in doc.paragraphs:
                text += paragraph.text + "\n"
            return text
        except Exception as e:
            logger.error(f"Error parsing DOCX: {e}")
            return ""
    
    def parse_txt(self, file_path: str) -> str:
        """Parse TXT/Markdown document"""
        try:
            with open(file_path, 'r', encoding='utf-8') as file:
                return file.read()
        except Exception as e:
            logger.error(f"Error parsing TXT: {e}")
            return ""
    
    def process_documents(self, message: MCPMessage):
        """Process uploaded documents"""
        files = message.payload.get("files", [])
        processed_docs = []
        
        for file_path in files:
            file_ext = os.path.splitext(file_path)[1].lower()
            
            # Parse document based on file type
            if file_ext == '.pdf':
                text = self.parse_pdf(file_path)
            elif file_ext == '.pptx':
                text = self.parse_pptx(file_path)
            elif file_ext == '.csv':
                text = self.parse_csv(file_path)
            elif file_ext == '.docx':
                text = self.parse_docx(file_path)
            elif file_ext in ['.txt', '.md']:
                text = self.parse_txt(file_path)
            else:
                logger.warning(f"Unsupported file type: {file_ext}")
                continue
            
            if text:
                # Split text into chunks
                chunks = self.text_splitter.split_text(text)
                docs = [Document(page_content=chunk, metadata={"source": file_path}) 
                        for chunk in chunks]
                processed_docs.extend(docs)
        
        # Send processed documents to RetrievalAgent
        response = MCPMessage(
            sender=self.name,
            receiver="RetrievalAgent",
            msg_type="INGESTION_COMPLETE",
            trace_id=message.trace_id,
            payload={"documents": processed_docs}
        )
        self.message_bus.publish(response)

class RetrievalAgent:
    """Agent responsible for embedding and semantic retrieval"""
    
    def __init__(self, message_bus: MessageBus):
        self.name = "RetrievalAgent"
        self.message_bus = message_bus
        self.message_bus.subscribe(self.name, self.handle_message)
        self.vector_store = None
    
    def handle_message(self, message: MCPMessage):
        """Handle incoming MCP messages"""
        if message.type == "INGESTION_COMPLETE":
            self.create_vector_store(message)
        elif message.type == "RETRIEVAL_REQUEST":
            self.retrieve_context(message)
    
    def create_vector_store(self, message: MCPMessage):
        """Create vector store from processed documents"""
        documents = message.payload.get("documents", [])
        
        if documents:
            try:
                self.vector_store = FAISS.from_documents(documents, embeddings)
                logger.info(f"Vector store created with {len(documents)} documents")
                
                # Notify completion
                response = MCPMessage(
                    sender=self.name,
                    receiver="CoordinatorAgent",
                    msg_type="VECTORSTORE_READY",
                    trace_id=message.trace_id,
                    payload={"status": "ready"}
                )
                self.message_bus.publish(response)
            except Exception as e:
                logger.error(f"Error creating vector store: {e}")
    
    def retrieve_context(self, message: MCPMessage):
        """Retrieve relevant context for a query"""
        query = message.payload.get("query", "")
        k = message.payload.get("k", 3)
        
        if self.vector_store and query:
            try:
                docs = self.vector_store.similarity_search(query, k=k)
                context = [{"content": doc.page_content, "source": doc.metadata.get("source", "")} 
                           for doc in docs]
                
                response = MCPMessage(
                    sender=self.name,
                    receiver="LLMResponseAgent",
                    msg_type="CONTEXT_RESPONSE",
                    trace_id=message.trace_id,
                    payload={
                        "query": query,
                        "retrieved_context": context,
                        "top_chunks": [doc.page_content for doc in docs]
                    }
                )
                self.message_bus.publish(response)
            except Exception as e:
                logger.error(f"Error retrieving context: {e}")

class LLMResponseAgent:
    """Agent responsible for generating LLM responses"""
    
    def __init__(self, message_bus: MessageBus):
        self.name = "LLMResponseAgent"
        self.message_bus = message_bus
        self.message_bus.subscribe(self.name, self.handle_message)
    
    def handle_message(self, message: MCPMessage):
        """Handle incoming MCP messages"""
        if message.type == "CONTEXT_RESPONSE":
            self.generate_response(message)
    
    def generate_response(self, message: MCPMessage):
        """Generate response using retrieved context"""
        query = message.payload.get("query", "")
        context = message.payload.get("retrieved_context", [])
        
        # Build context string
        context_text = "\n\n".join([f"Source: {ctx['source']}\nContent: {ctx['content']}" 
                                      for ctx in context])
        
        # Create messages for conversational format
        messages = [
            {
                "role": "system",
                "content": "You are a helpful assistant. Based on the provided context below, answer the user's question accurately and comprehensively. Cite the sources if possible.",
            },
            {
                "role": "user", 
                "content": f"Context:\n\n{context_text}\n\nQuestion: {query}"
            }
        ]

        try:
            # Use client.chat_completion for conversational models
            response_stream = client.chat_completion(
                messages=messages,
                max_tokens=512,
                temperature=0.7,
                stream=True
            )
            
            # Send streaming response
            response = MCPMessage(
                sender=self.name,
                receiver="CoordinatorAgent",
                msg_type="LLM_RESPONSE_STREAM",
                trace_id=message.trace_id,
                payload={
                    "query": query,
                    "response_stream": response_stream,
                    "context": context
                }
            )
            self.message_bus.publish(response)
            
        except Exception as e:
            logger.error(f"Error generating response: {e}")
            # Send an error stream back
            error_msg = f"Error from LLM: {e}"
            def error_generator():
                yield error_msg
            
            response = MCPMessage(
                sender=self.name,
                receiver="CoordinatorAgent",
                msg_type="LLM_RESPONSE_STREAM",
                trace_id=message.trace_id,
                payload={"response_stream": error_generator()}
            )
            self.message_bus.publish(response)

class CoordinatorAgent:
    """Coordinator agent that orchestrates the entire workflow"""
    
    def __init__(self, message_bus: MessageBus):
        self.name = "CoordinatorAgent"
        self.message_bus = message_bus
        self.message_bus.subscribe(self.name, self.handle_message)
        self.current_response_stream = None
        self.vector_store_ready = False
    
    def handle_message(self, message: MCPMessage):
        """Handle incoming MCP messages"""
        if message.type == "VECTORSTORE_READY":
            self.vector_store_ready = True
        elif message.type == "LLM_RESPONSE_STREAM":
            self.current_response_stream = message.payload.get("response_stream")
    
    def process_files(self, files):
        """Process uploaded files"""
        if not files:
            return "No files uploaded."
        
        file_paths = [file.name for file in files]
        
        # Send ingestion request
        message = MCPMessage(
            sender=self.name,
            receiver="IngestionAgent",
            msg_type="INGESTION_REQUEST",
            payload={"files": file_paths}
        )
        self.message_bus.publish(message)
        
        return f"Processing {len(files)} files: {', '.join([os.path.basename(fp) for fp in file_paths])}"
    
    def handle_query(self, query: str, history: List) -> Generator[str, None, None]:
        """Handle user query and return streaming response"""
        if not self.vector_store_ready:
            yield "Please upload and process documents first."
            return
        
        # Send retrieval request
        message = MCPMessage(
            sender=self.name,
            receiver="RetrievalAgent",
            msg_type="RETRIEVAL_REQUEST",
            payload={"query": query}
        )
        self.message_bus.publish(message)
        
        # Wait for response and stream
        import time
        timeout = 20  # seconds
        start_time = time.time()
        
        while not self.current_response_stream and (time.time() - start_time) < timeout:
            time.sleep(0.1)
        
        if self.current_response_stream:
            try:
                # Stream tokens directly
                for chunk in self.current_response_stream:
                    # The token is in chunk.choices[0].delta.content for chat_completion
                    token = chunk.choices[0].delta.content
                    if token:
                        yield token
            except Exception as e:
                yield f"Error streaming response: {e}"
            finally:
                self.current_response_stream = None # Reset for next query
        else:
            yield "Timeout: No response received from LLM agent."

# Initialize agents
ingestion_agent = IngestionAgent(message_bus)
retrieval_agent = RetrievalAgent(message_bus)
llm_response_agent = LLMResponseAgent(message_bus)
coordinator_agent = CoordinatorAgent(message_bus)

# Gradio Interface
def create_interface():
    """Create Gradio interface"""
    
    with gr.Blocks(
        theme=gr.themes.Base().set(
            body_background_color="#0b0f19",
            body_text_color="#ffffff",
            button_primary_background_fill="#6366f1",
            button_primary_background_fill_hover="#4f46e5",
            button_primary_text_color="#ffffff",
            block_background_fill="#1f2937",
            block_border_color="#374151",
            input_background_fill="#374151",
            input_border_color="#4b5563",
            input_placeholder_color="#9ca3af",
            checkbox_background_color="#374151",
            checkbox_border_color="#4b5563",
            panel_background_fill="#111827",
            panel_border_color="#374151",
        ),
        css="""
        .gradio-container {
            max-width: 100vw !important;
            margin: 0 !important;
            padding: 20px !important;
            background: linear-gradient(135deg, #0b0f19 0%, #1a1b3a 100%);
            min-height: 100vh;
        }
        .header-container {
            text-align: center;
            padding: 2rem 0;
            margin-bottom: 2rem;
            background: linear-gradient(135deg, #6366f1 0%, #8b5cf6 100%);
            -webkit-background-clip: text;
            -webkit-text-fill-color: transparent;
            background-clip: text;
        }
        .header-title {
            font-size: 3.5rem;
            font-weight: 800;
            margin: 0;
            text-shadow: 0 0 30px rgba(99, 102, 241, 0.5);
            font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
        }
        .header-subtitle {
            font-size: 1.25rem;
            margin-top: 0.5rem;
            color: #9ca3af;
            font-weight: 400;
        }
        .sidebar-container {
            background: rgba(31, 41, 55, 0.8) !important;
            border-radius: 20px !important;
            padding: 24px !important;
            backdrop-filter: blur(10px);
            border: 1px solid #374151;
            height: fit-content;
            position: sticky;
            top: 20px;
        }
        .chat-container {
            background: rgba(31, 41, 55, 0.6) !important;
            border-radius: 20px !important;
            padding: 24px !important;
            backdrop-filter: blur(10px);
            border: 1px solid #374151;
            height: calc(100vh - 200px);
            display: flex;
            flex-direction: column;
        }
        .upload-area {
            border: 2px dashed #6366f1 !important;
            border-radius: 15px !important;
            padding: 24px !important;
            margin: 16px 0 !important;
            background: rgba(99, 102, 241, 0.05) !important;
            transition: all 0.3s ease;
        }
        .upload-area:hover {
            border-color: #8b5cf6 !important;
            background: rgba(139, 92, 246, 0.1) !important;
            transform: translateY(-2px);
            box-shadow: 0 8px 25px rgba(99, 102, 241, 0.2);
        }
        .status-display {
            background: linear-gradient(135deg, #065f46, #047857) !important;
            border: none !important;
            border-radius: 12px !important;
            color: #ecfdf5 !important;
            padding: 12px 16px !important;
            font-family: 'JetBrains Mono', monospace !important;
            font-size: 0.9rem !important;
        }
        .process-btn {
            background: linear-gradient(135deg, #6366f1, #8b5cf6) !important;
            border: none !important;
            border-radius: 12px !important;
            padding: 12px 24px !important;
            font-weight: 600 !important;
            text-transform: uppercase !important;
            letter-spacing: 0.5px !important;
            transition: all 0.3s ease !important;
            box-shadow: 0 4px 15px rgba(99, 102, 241, 0.3) !important;
        }
        .process-btn:hover {
            transform: translateY(-2px) !important;
            box-shadow: 0 8px 25px rgba(99, 102, 241, 0.4) !important;
        }
        .chatbot-container {
            flex: 1 !important;
            background: rgba(17, 24, 39, 0.8) !important;
            border-radius: 15px !important;
            border: 1px solid #374151 !important;
            margin-bottom: 20px !important;
            overflow: hidden !important;
        }
        .input-row {
            background: rgba(55, 65, 81, 0.5) !important;
            border-radius: 15px !important;
            padding: 16px !important;
            margin-top: 16px !important;
            border: 1px solid #4b5563 !important;
        }
        .send-btn {
            background: linear-gradient(135deg, #10b981, #059669) !important;
            border: none !important;
            border-radius: 10px !important;
            padding: 12px 20px !important;
            font-weight: 600 !important;
            transition: all 0.3s ease !important;
            box-shadow: 0 4px 15px rgba(16, 185, 129, 0.3) !important;
        }
        .send-btn:hover {
            transform: translateY(-2px) !important;
            box-shadow: 0 8px 25px rgba(16, 185, 129, 0.4) !important;
        }
        .examples-container {
            margin-top: 16px !important;
        }
        .examples-container .gr-button {
            background: rgba(75, 85, 99, 0.5) !important;
            border: 1px solid #6b7280 !important;
            border-radius: 8px !important;
            margin: 4px !important;
            transition: all 0.3s ease !important;
        }
        .examples-container .gr-button:hover {
            background: rgba(99, 102, 241, 0.2) !important;
            border-color: #6366f1 !important;
            transform: translateY(-1px) !important;
        }
        .architecture-info {
            background: rgba(17, 24, 39, 0.8) !important;
            border-radius: 15px !important;
            padding: 20px !important;
            margin-top: 24px !important;
            border: 1px solid #374151 !important;
        }
        .architecture-info h3 {
            color: #6366f1 !important;
            margin-bottom: 16px !important;
            font-size: 1.1rem !important;
            font-weight: 600 !important;
        }
        .architecture-info p, .architecture-info ul {
            color: #d1d5db !important;
            line-height: 1.6 !important;
            font-size: 0.9rem !important;
        }
        .architecture-info ul {
            margin-left: 16px !important;
        }
        .architecture-info li {
            margin-bottom: 8px !important;
        }
        .gr-textbox {
            background: rgba(55, 65, 81, 0.8) !important;
            border: 1px solid #4b5563 !important;
            border-radius: 12px !important;
            color: #f9fafb !important;
        }
        .gr-textbox:focus {
            border-color: #6366f1 !important;
            box-shadow: 0 0 0 3px rgba(99, 102, 241, 0.1) !important;
        }
        """,
        title="Agentic RAG Chatbot"
    ) as iface:
        
        # Header
        gr.HTML("""
        <div class="header-container">
            <h1 class="header-title">Agentic RAG Chatbot</h1>
            <p class="header-subtitle">Multi-Format Document QA with Model Context Protocol (MCP)</p>
        </div>
        """)
        
        with gr.Row(equal_height=False):
            # Sidebar
            with gr.Column(scale=1, min_width=350):
                with gr.Group(elem_classes=["sidebar-container"]):
                    gr.Markdown("### Document Upload", elem_id="upload-title")
                    
                    file_upload = gr.File(
                        file_count="multiple",
                        file_types=[".pdf", ".pptx", ".csv", ".docx", ".txt", ".md"],
                        label="Drop files here or click to browse",
                        elem_classes=["upload-area"],
                        scale=4
                    )
                    
                    upload_status = gr.Textbox(
                        label="Processing Status",
                        interactive=False,
                        max_lines=4,
                        elem_classes=["status-display"],
                        value="Ready to process documents..."
                    )
                    
                    process_btn = gr.Button(
                        "Process Documents", 
                        variant="primary",
                        size="lg",
                        elem_classes=["process-btn"],
                        scale=1
                    )
                
                with gr.Group(elem_classes=["architecture-info"]):
                    gr.Markdown("""
                    ###  System Architecture
                    
                    ** Agents:**
                    - **IngestionAgent**: Document parsing & preprocessing
                    - **RetrievalAgent**: Semantic search & embeddings  
                    - **LLMResponseAgent**: AI response generation
                    - **CoordinatorAgent**: Workflow orchestration
                    
                    **📡 MCP Communication:** 
                    Structured message passing between agents for scalable, maintainable AI workflows.
                    
                    **🔧 Supported Formats:**
                    PDF, PPTX, CSV, DOCX, TXT, MD
                    """)
            
            # Main Chat Area
            with gr.Column(scale=2, min_width=600):
                with gr.Group(elem_classes=["chat-container"]):
                    gr.Markdown("### 💬 Chat Interface", elem_id="chat-title")
                    
                    chatbot = gr.Chatbot(
                        height=550,
                        elem_classes=["chatbot-container"],
                        show_copy_button=True,
                        type="messages",
                        avatar_images=("[( ͡° ͜ʖ ͡°)]", " [¯\_(ツ)_/¯]"),
                        bubble_full_width=False,
                        show_share_button=False
                    )
                    
                    with gr.Group(elem_classes=["input-row"]):
                        with gr.Row():
                            msg = gr.Textbox(
                                label="",
                                placeholder="💭 Ask a question about your documents...",
                                scale=4,
                                autofocus=True,
                                container=False,
                                show_label=False
                            )
                            submit_btn = gr.Button(
                                "Send", 
                                scale=1, 
                                variant="primary",
                                elem_classes=["send-btn"]
                            )
                    
                    with gr.Group(elem_classes=["examples-container"]):
                        gr.Examples(
                            examples=[
                                "What are the main topics discussed in the documents?",
                                "Can you summarize the key findings?",
                                "What metrics or KPIs are mentioned?",
                                "What recommendations are provided?",
                                "Are there any trends or patterns identified?"
                            ],
                            inputs=msg,
                            label="Example Questions"
                        )
        
        # Event handlers
        def process_files_handler(files):
            if files:
                return f" Processing {len(files)} files: {', '.join([os.path.basename(f.name) for f in files])}\n⏳ Please wait..."
            return "No files uploaded."
        
        def respond(message, history):
            if message.strip():
                # Add user message to history in the new format
                history.append({"role": "user", "content": message})
                # Add a placeholder for the assistant's response
                history.append({"role": "assistant", "content": ""})

                # Get streaming response
                for token in coordinator_agent.handle_query(message, history):
                    # Append each token to the assistant's message content
                    history[-1]["content"] += token
                    yield history, "" # Yield updated history and clear the textbox
            else:
                yield history, message
        
        process_btn.click(
            process_files_handler,
            inputs=[file_upload],
            outputs=[upload_status]
        )
        
        submit_btn.click(
            respond,
            inputs=[msg, chatbot],
            outputs=[chatbot, msg],
            show_progress="minimal"
        )
        
        msg.submit(
            respond,
            inputs=[msg, chatbot],
            outputs=[chatbot, msg],
            show_progress="minimal"
        )
    
    return iface

# Launch the application
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860
    )