import gradio as gr
import os
import tempfile
import uuid
from datetime import datetime
from typing import List, Dict, Any, Optional
import json
import asyncio
from dataclasses import dataclass, asdict
import logging
import sys # Import sys for exiting if token is missing

# Document processing imports
import PyPDF2
import pandas as pd
from docx import Document
from pptx import Presentation
import markdown

# ML/AI imports
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.schema import Document as LCDocument
from huggingface_hub import InferenceClient

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

# --- Get HF token from environment and perform a crucial check ---
HF_TOKEN = os.getenv('hf_token')

if HF_TOKEN is None:
    logger.error("FATAL ERROR: HuggingFace token (HF_TOKEN) environment variable is not set.")
    logger.error("Please set the HF_TOKEN environment variable before running the application.")
    sys.exit("HuggingFace token (HF_TOKEN) is not set. Exiting.")


# MCP Message Structure
@dataclass
class MCPMessage:
    sender: str
    receiver: str
    type: str
    trace_id: str
    payload: Dict[str, Any]
    timestamp: str = None

    def __post_init__(self):
        if self.timestamp is None:
            self.timestamp = datetime.now().isoformat()

    def to_dict(self):
        return asdict(self)

# MCP Communication Layer
class MCPCommunicator:
    def __init__(self):
        self.message_queue = asyncio.Queue()
        self.subscribers = {}

    async def send_message(self, message: MCPMessage):
        logger.info(f"MCP: {message.sender} -> {message.receiver}: {message.type}")
        await self.message_queue.put(message)

    async def receive_message(self, agent_name: str) -> MCPMessage:
        while True:
            message = await self.message_queue.get()
            if message.receiver == agent_name:
                return message
            # Re-queue if not for this agent
            await self.message_queue.put(message)

# Global MCP instance
mcp = MCPCommunicator()

# Base Agent Class
class BaseAgent:
    def __init__(self, name: str):
        self.name = name
        self.mcp = mcp

    async def send_mcp_message(self, receiver: str, msg_type: str, payload: Dict[str, Any], trace_id: str):
        message = MCPMessage(
            sender=self.name,
            receiver=receiver,
            type=msg_type,
            trace_id=trace_id,
            payload=payload
        )
        await self.mcp.send_message(message)

    async def receive_mcp_message(self) -> MCPMessage:
        return await self.mcp.receive_message(self.name)

# Document Ingestion Agent
class IngestionAgent(BaseAgent):
    def __init__(self):
        super().__init__("IngestionAgent")
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            length_function=len,
        )

    def parse_pdf(self, file_path: str) -> str:
        """Parse PDF file and extract text"""
        try:
            with open(file_path, 'rb') as file:
                pdf_reader = PyPDF2.PdfReader(file)
                text = ""
                for page in pdf_reader.pages:
                    text += page.extract_text() + "\n"
                return text
        except Exception as e:
            logger.error(f"Error parsing PDF: {e}")
            return ""

    def parse_docx(self, file_path: str) -> str:
        """Parse DOCX file and extract text"""
        try:
            doc = Document(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_pptx(self, file_path: str) -> str:
        """Parse PPTX file and extract text"""
        try:
            prs = Presentation(file_path)
            text = ""
            for slide_num, slide in enumerate(prs.slides, 1):
                text += f"Slide {slide_num}:\n"
                for shape in slide.shapes:
                    if hasattr(shape, "text"):
                        text += shape.text + "\n"
                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 file and convert to text"""
        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_txt_md(self, file_path: str) -> str:
        """Parse TXT or MD file"""
        try:
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
                # If markdown, convert to plain text
                if file_path.lower().endswith('.md'):
                    content = markdown.markdown(content)
                return content
        except Exception as e:
            logger.error(f"Error parsing TXT/MD: {e}")
            return ""

    async def process_documents(self, files: List[str], trace_id: str) -> List[LCDocument]:
        """Process uploaded documents and return chunked documents"""
        all_documents = []

        for file_path in files:
            file_ext = os.path.splitext(file_path)[1].lower()
            filename = os.path.basename(file_path)

            # Parse based on file extension
            if file_ext == '.pdf':
                content = self.parse_pdf(file_path)
            elif file_ext == '.docx':
                content = self.parse_docx(file_path)
            elif file_ext == '.pptx':
                content = self.parse_pptx(file_path)
            elif file_ext == '.csv':
                content = self.parse_csv(file_path)
            elif file_ext in ['.txt', '.md']:
                content = self.parse_txt_md(file_path)
            else:
                logger.warning(f"Unsupported file type: {file_ext}")
                continue

            if content.strip():
                # Split content into chunks
                chunks = self.text_splitter.split_text(content)

                # Create LangChain documents
                for i, chunk in enumerate(chunks):
                    doc = LCDocument(
                        page_content=chunk,
                        metadata={
                            "source": filename,
                            "chunk_id": i,
                            "file_type": file_ext
                        }
                    )
                    all_documents.append(doc)

        return all_documents

# Retrieval Agent
class RetrievalAgent(BaseAgent):
    def __init__(self):
        super().__init__("RetrievalAgent")
        self.embeddings = HuggingFaceEmbeddings(
            model_name="sentence-transformers/all-MiniLM-L6-v2"
        )
        self.vector_store = None

    async def create_vector_store(self, documents: List[LCDocument], trace_id: str):
        """Create vector store from documents"""
        try:
            if documents:
                self.vector_store = FAISS.from_documents(documents, self.embeddings)
                logger.info(f"Created vector store with {len(documents)} documents")
            else:
                logger.warning("No documents to create vector store")
        except Exception as e:
            logger.error(f"Error creating vector store: {e}")

    async def retrieve_relevant_chunks(self, query: str, k: int = 5, trace_id: str = None) -> List[Dict]:
        """Retrieve relevant chunks for a query"""
        if not self.vector_store:
            return []

        try:
            # Similarity search
            docs = self.vector_store.similarity_search(query, k=k)

            # Format results
            results = []
            for doc in docs:
                results.append({
                    "content": doc.page_content,
                    "source": doc.metadata.get("source", "Unknown"),
                    "chunk_id": doc.metadata.get("chunk_id", 0),
                    "file_type": doc.metadata.get("file_type", "Unknown")
                })

            return results
        except Exception as e:
            logger.error(f"Error retrieving chunks: {e}")
            return []

# LLM Response Agent
class LLMResponseAgent(BaseAgent):
    def __init__(self):
        super().__init__("LLMResponseAgent")
        # Use the global HF_TOKEN which is validated at script start
        self.client = InferenceClient(
            model="meta-llama/Llama-3.1-8B-Instruct",
            token=HF_TOKEN # Pass the token here
        )

    def format_prompt_for_conversational(self, query: str, context_chunks: List[Dict]) -> str:
        """
        Format prompt with context and query as a single 'user' input
        suitable for a conversational model.
        """
        context_text = "\n\n".join([
            f"Source: {chunk['source']}\nContent: {chunk['content']}"
            for chunk in context_chunks
        ])

        # We are putting the RAG prompt into the 'user' input for the conversational model.
        # This is a common way to use a conversational model for RAG if text_generation isn't available.
        prompt_as_user_input = f"""Based on the following context from uploaded documents, please answer the user's question.

Context:
{context_text}

Question: {query}

Please provide a comprehensive answer based on the context above. If the context doesn't contain enough information to fully answer the question, please mention what information is available and what might be missing.

Answer:""" # Keeping "Answer:" to guide the model to start generating the answer directly.
        return prompt_as_user_input

    async def generate_response(self, query: str, context_chunks: List[Dict], trace_id: str) -> str:
        """Generate response using LLM via the conversational task."""
        try:
            # Format the RAG prompt as the user's input for the conversational model
            formatted_input = self.format_prompt_for_conversational(query, context_chunks)

            # Use the conversational task
            response = self.client.conversational(
                inputs=formatted_input, # This is the current user turn
                # No past_user_inputs or generated_responses are provided initially
                # to keep it stateless per query, akin to text_generation.
                parameters={
                    "temperature": 0.7,
                    "max_new_tokens": 512,
                    # Add other parameters if needed, e.g., do_sample, top_p, top_k
                    # "do_sample": True,
                    # "top_p": 0.95,
                    # "top_k": 50,
                }
            )

            # The conversational response has a list of generated responses.
            # We assume the first one is the primary answer.
            if response.generated_responses:
                return response.generated_responses[0]
            else:
                logger.warning("LLM generated an empty response via conversational API.")
                return "I apologize, the model did not generate a response."

        except Exception as e:
            logger.error(f"Error generating response with conversational LLM: {e}")
            return f"I apologize, but I encountered an error while generating the response: {str(e)}"

# Coordinator Agent
class CoordinatorAgent(BaseAgent):
    def __init__(self):
        super().__init__("CoordinatorAgent")
        self.ingestion_agent = IngestionAgent()
        self.retrieval_agent = RetrievalAgent()
        self.llm_agent = LLMResponseAgent() # LLMResponseAgent will use the global HF_TOKEN
        self.documents_processed = False

    async def process_documents(self, files: List[str]) -> str:
        """Orchestrate document processing"""
        trace_id = str(uuid.uuid4())

        try:
            # Step 1: Ingestion
            await self.send_mcp_message(
                "IngestionAgent",
                "DOCUMENT_INGESTION_REQUEST",
                {"files": files},
                trace_id
            )

            documents = await self.ingestion_agent.process_documents(files, trace_id)

            await self.send_mcp_message(
                "RetrievalAgent",
                "VECTOR_STORE_CREATE_REQUEST",
                {"documents": len(documents)},
                trace_id
            )

            # Step 2: Create vector store
            await self.retrieval_agent.create_vector_store(documents, trace_id)

            self.documents_processed = True

            return f"Successfully processed {len(documents)} document chunks from {len(files)} files."

        except Exception as e:
            logger.error(f"Error in document processing: {e}")
            return f"Error processing documents: {str(e)}"

    async def answer_query(self, query: str) -> tuple[str, List[Dict]]:
        """Orchestrate query answering"""
        if not self.documents_processed:
            return "Please upload and process documents first.", []

        trace_id = str(uuid.uuid4())

        try:
            # Step 1: Retrieval
            await self.send_mcp_message(
                "RetrievalAgent",
                "RETRIEVAL_REQUEST",
                {"query": query},
                trace_id
            )

            context_chunks = await self.retrieval_agent.retrieve_relevant_chunks(query, k=5, trace_id=trace_id)

            # Step 2: LLM Response
            await self.send_mcp_message(
                "LLMResponseAgent",
                "LLM_GENERATION_REQUEST",
                {"query": query, "context_chunks": len(context_chunks)},
                trace_id
            )

            response = await self.llm_agent.generate_response(query, context_chunks, trace_id)

            return response, context_chunks

        except Exception as e:
            logger.error(f"Error in query processing: {e}")
            return f"Error processing query: {str(e)}", []

# Global coordinator instance
coordinator = CoordinatorAgent()

async def process_files(files):
    """Process uploaded files"""
    if not files:
        return "❌ Please upload at least one file."

    file_paths = []
    for file in files:
        temp_dir = tempfile.gettempdir()
        unique_filename = f"{uuid.uuid4()}_{os.path.basename(file.name)}"
        temp_path = os.path.join(temp_dir, unique_filename)
        try:
            file_content = file.read()
            with open(temp_path, 'wb') as f:
                f.write(file_content)
            file_paths.append(temp_path)
        except Exception as e:
            logger.error(f"Error saving uploaded file {file.name}: {e}")
            return f"❌ Error saving uploaded file {file.name}: {e}"

    result = await coordinator.process_documents(file_paths)

    for path in file_paths:
        try:
            os.remove(path)
        except Exception as e:
            logger.warning(f"Could not remove temporary file {path}: {e}")

    return result

async def answer_question(query, history):
    """Answer user question"""
    if not query.strip():
        return history, ""

    response, context_chunks = await coordinator.answer_query(query)

    if context_chunks:
        sources = "\n\n**Sources:**\n"
        for i, chunk in enumerate(context_chunks[:3], 1):
            sources += f"{i}. {chunk['source']} (Chunk {chunk['chunk_id']})\n"
        response += sources

    history.append((query, response))

    return history, ""

# Custom CSS (unchanged)
custom_css = """
/* Main container styling */
.gradio-container {
    max-width: 1200px !important;
    margin: 0 auto !important;
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif !important;
}

/* Header styling */
.header-container {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    color: white !important;
    padding: 2rem !important;
    border-radius: 15px !important;
    margin-bottom: 2rem !important;
    text-align: center !important;
    box-shadow: 0 8px 32px rgba(0,0,0,0.1) !important;
}

.header-title {
    font-size: 2.5rem !important;
    font-weight: 700 !important;
    margin-bottom: 0.5rem !important;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.3) !important;
}

.header-subtitle {
    font-size: 1.2rem !important;
    opacity: 0.9 !important;
    font-weight: 300 !important;
}

/* Tab styling */
.tab-nav {
    background: white !important;
    border-radius: 12px !important;
    box-shadow: 0 4px 20px rgba(0,0,0,0.08) !important;
    padding: 0.5rem !important;
    margin-bottom: 1rem !important;
}

/* Card styling */
.upload-card, .chat-card {
    background: white !important;
    border-radius: 15px !important;
    padding: 2rem !important;
    box-shadow: 0 4px 20px rgba(0,0,0,0.08) !important;
    border: 1px solid #e1e5e9 !important;
    margin-bottom: 1.5rem !important;
}

/* Button styling */
.primary-button {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    color: white !important;
    border: none !important;
    border-radius: 10px !important;
    padding: 0.75rem 2rem !important;
    font-weight: 600 !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 15px rgba(102, 126, 234, 0.3) !important;
}

.primary-button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(102, 126, 234, 0.4) !important;
}

/* Chat interface styling */
.chat-container {
    max-height: 600px !important;
    overflow-y: auto !important;
    background: #f8f9fa !important;
    border-radius: 15px !important;
    padding: 1rem !important;
    border: 1px solid #e1e5e9 !important;
}

/* Input styling */
.input-container input, .input-container textarea {
    border: 2px solid #e1e5e9 !important;
    border-radius: 10px !important;
    padding: 0.75rem 1rem !important;
    font-size: 1rem !important;
    transition: border-color 0.3s ease !important;
}

.input-container input:focus, .input-container textarea:focus {
    border-color: #667eea !important;
    box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
    outline: none !important;
}

/* Status indicators */
.status-success {
    color: #28a745 !important;
    background: #d4edda !important;
    padding: 0.75rem 1rem !important;
    border-radius: 8px !important;
    border: 1px solid #c3e6cb !important;
    margin: 1rem 0 !important;
}

.status-error {
    color: #dc3545 !important;
    background: #f8d7da !important;
    padding: 0.75rem 1rem !important;
    border-radius: 8px !important;
    border: 1px solid #f5c6cb !important;
    margin: 1rem 0 !important;
}

/* File upload styling */
.file-upload {
    border: 2px dashed #667eea !important;
    border-radius: 15px !important;
    padding: 2rem !important;
    text-align: center !important;
    background: #f8f9ff !important;
    transition: all 0.3s ease !important;
}

.file-upload:hover {
    border-color: #764ba2 !important;
    background: #f0f4ff !important;
}

/* Architecture diagram container */
.architecture-container {
    background: white !important;
    border-radius: 15px !important;
    padding: 2rem !important;
    margin: 1rem 0 !important;
    box-shadow: 0 4px 20px rgba(0,0,0,0.08) !important;
    text-align: center !important;
}

/* Responsive design */
@media (max-width: 768px) {
    .header-title {
        font-size: 2rem !important;
    }

    .upload-card, .chat-card {
        padding: 1.5rem !important;
    }
}

/* Animation for loading states */
@keyframes pulse {
    0% { opacity: 1; }
    50% { opacity: 0.5; }
    100% { opacity: 1; }
}

.loading {
    animation: pulse 1.5s ease-in-out infinite !important;
}
"""

# Create Gradio Interface
def create_interface():
    with gr.Blocks(css=custom_css, title="🤖 Agentic RAG Chatbot") as demo:
        gr.HTML("""
        <div class="header-container">
            <h1 class="header-title">🤖 Agentic RAG Chatbot</h1>
            <p class="header-subtitle">Multi-Format Document QA using Model Context Protocol (MCP)</p>
        </div>
        """)

        with gr.Tabs() as tabs:
            # Upload Tab (now the first tab)
            with gr.TabItem("📁 Upload Documents", elem_classes=["tab-nav"]):
                gr.HTML("""
                <div class="upload-card">
                    <h3>📄 Document Upload</h3>
                    <p>Upload your documents in any supported format: PDF, DOCX, PPTX, CSV, TXT, or Markdown.</p>
                </div>
                """)

                file_upload = gr.File(
                    label="Choose Files",
                    file_count="multiple",
                    file_types=[".pdf", ".docx", ".pptx", ".csv", ".txt", ".md"],
                    elem_classes=["file-upload"]
                )

                upload_button = gr.Button(
                    "Process Documents",
                    variant="primary",
                    elem_classes=["primary-button"]
                )

                upload_status = gr.Textbox(
                    label="Processing Status",
                    interactive=False,
                    elem_classes=["input-container"]
                )

            # Chat Tab
            with gr.TabItem("💬 Chat", elem_classes=["tab-nav"]):
                gr.HTML("""
                <div class="chat-card">
                    <h3>🗨️ Ask Questions</h3>
                    <p>Ask questions about your uploaded documents. The AI will provide answers based on the document content.</p>
                </div>
                """)

                chatbot = gr.Chatbot(
                    label="Conversation",
                    height=400,
                    elem_classes=["chat-container"]
                )

                with gr.Row():
                    query_input = gr.Textbox(
                        label="Your Question",
                        placeholder="What are the key findings in the document?",
                        elem_classes=["input-container"]
                    )
                    ask_button = gr.Button(
                        "Ask",
                        variant="primary",
                        elem_classes=["primary-button"]
                    )

                gr.Examples(
                    examples=[
                        "What are the main topics covered in the documents?",
                        "Can you summarize the key findings?",
                        "What are the important metrics mentioned?",
                        "What recommendations are provided?",
                    ],
                    inputs=query_input,
                    label="Example Questions"
                )

            # Architecture Tab
            with gr.TabItem("🏗️ Architecture", elem_classes=["tab-nav"]):
                gr.HTML("""
                <div class="architecture-container">
                    <h3>🏛️ System Architecture</h3>
                    <p>This system uses an agentic architecture with Model Context Protocol (MCP) for inter-agent communication.</p>
                </div>
                """)

                gr.Markdown("""
                ## 🔄 Agent Flow Diagram

                ```
                User Upload → CoordinatorAgent → IngestionAgent → RetrievalAgent → LLMResponseAgent
                     ↓              ↓                ↓               ↓              ↓
                 Documents    MCP Messages    Text Chunks    Vector Store    Final Response
                ```

                ## 🤖 Agent Descriptions

                - **CoordinatorAgent**: Orchestrates the entire workflow and manages MCP communication
                - **IngestionAgent**: Parses and preprocesses documents (PDF, DOCX, PPTX, CSV, TXT, MD)
                - **RetrievalAgent**: Handles embeddings and semantic retrieval using FAISS
                - **LLMResponseAgent**: Generates final responses using Llama-3.1-8B-Instruct

                ## 🔗 Tech Stack

                - **Frontend**: Gradio with custom CSS
                - **LLM**: Meta Llama-3.1-8B-Instruct (via HuggingFace Inference)
                - **Embeddings**: sentence-transformers/all-MiniLM-L6-v2
                - **Vector Store**: FAISS
                - **Document Processing**: PyPDF2, python-docx, python-pptx, pandas
                - **Framework**: LangChain for document handling

                ## 📨 MCP Message Example

                ```json
                {
                  "sender": "RetrievalAgent",
                  "receiver": "LLMResponseAgent",
                  "type": "RETRIEVAL_RESULT",
                  "trace_id": "rag-457",
                  "payload": {
                    "retrieved_context": ["Revenue increased by 25%", "Q1 KPIs exceeded targets"],
                    "query": "What were the Q1 KPIs?"
                  },
                  "timestamp": "2025-07-21T10:30:00Z"
                }
                ```
                """)

        # Event handlers
        upload_button.click(
            fn=process_files,
            inputs=[file_upload],
            outputs=[upload_status]
        )

        ask_button.click(
            fn=answer_question,
            inputs=[query_input, chatbot],
            outputs=[chatbot, query_input]
        )

        query_input.submit(
            fn=answer_question,
            inputs=[query_input, chatbot],
            outputs=[chatbot, query_input]
        )

    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860,
        show_api=False
    )