import gradio as gr
import numpy as np
from sentence_transformers import SentenceTransformer
import faiss
import os
import PyPDF2
import docx
import pandas as pd

class PureRAGBot:
    def __init__(self):
        # Embedding model for document search
        self.embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
        self.documents = []
        self.index = None
        self.is_ready = False

    def load_file_content(self, file):
        """Load content from various file types"""
        try:
            if file is None:
                return "Please select a file first."
            
            file_path = file.name
            file_extension = os.path.splitext(file_path)[1].lower()
            
            if file_extension == '.txt':
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read()
                chunks = self.split_text_into_chunks(content)
                
            elif file_extension == '.csv':
                df = pd.read_csv(file_path)
                chunks = self.dataframe_to_chunks(df)
                
            elif file_extension == '.pdf':
                chunks = self.pdf_to_chunks(file_path)
                
            elif file_extension in ['.docx', '.doc']:
                chunks = self.docx_to_chunks(file_path)
                
            else:
                return "Unsupported file format. Please upload TXT, CSV, PDF, or DOCX files."
            
            # Create FAISS index
            if chunks:
                embeddings = self.embedder.encode(chunks)
                self.index = faiss.IndexFlatIP(embeddings.shape[1])
                self.index.add(embeddings.astype('float32'))
                self.documents = chunks
                self.is_ready = True
                return f"✅ Document successfully loaded! {len(chunks)} content sections indexed. You can now ask questions about this document."
            else:
                return "❌ No readable content found in the file."
                
        except Exception as e:
            return f"❌ Error processing file: {str(e)}"
    
    def split_text_into_chunks(self, text, chunk_size=300):
        """Split text into manageable chunks"""
        sentences = text.split('.')
        chunks = []
        current_chunk = ""
        
        for sentence in sentences:
            sentence = sentence.strip()
            if not sentence:
                continue
                
            if len(current_chunk) + len(sentence) < chunk_size:
                current_chunk += sentence + '. '
            else:
                if current_chunk:
                    chunks.append(current_chunk.strip())
                current_chunk = sentence + '. '
        
        if current_chunk:
            chunks.append(current_chunk.strip())
        
        return chunks if chunks else [text[:500]]
    
    def dataframe_to_chunks(self, df):
        """Convert DataFrame to text chunks"""
        chunks = []
        for idx, row in df.iterrows():
            row_text = " | ".join([str(cell) for cell in row if pd.notna(cell)])
            if len(row_text) > 500:
                row_text = row_text[:500] + "..."
            chunks.append(f"Row {idx+1}: {row_text}")
        return chunks
    
    def pdf_to_chunks(self, file_path):
        """Extract text from PDF"""
        try:
            with open(file_path, 'rb') as file:
                reader = PyPDF2.PdfReader(file)
                text = ""
                for page in reader.pages:
                    text += page.extract_text() + "\n"
                return self.split_text_into_chunks(text)
        except Exception as e:
            return [f"Error reading PDF: {str(e)}"]
    
    def docx_to_chunks(self, file_path):
        """Extract text from DOCX"""
        try:
            doc = docx.Document(file_path)
            text = ""
            for paragraph in doc.paragraphs:
                if paragraph.text.strip():
                    text += paragraph.text + "\n"
            return self.split_text_into_chunks(text)
        except Exception as e:
            return [f"Error reading DOCX: {str(e)}"]
    
    def search_documents(self, query, k=3):
        """Search for relevant documents"""
        if not self.is_ready:
            return []
        
        try:
            query_embedding = self.embedder.encode([query])
            distances, indices = self.index.search(query_embedding.astype('float32'), k)
            
            results = []
            for i, idx in enumerate(indices[0]):
                if idx < len(self.documents):
                    results.append({
                        'content': self.documents[idx],
                        'score': float(distances[0][i])
                    })
            return results
        except Exception as e:
            print(f"Search error: {e}")
            return []
    
    def generate_rag_response(self, query):
        """Generate response purely from document content"""
        if not self.is_ready:
            return "❌ Please upload a document file first to ask questions about its content."
        
        # Search for relevant content
        results = self.search_documents(query)
        
        if not results:
            return f"❌ I couldn't find any information about '{query}' in the uploaded document. Please try rephrasing your question or ask about different content from the document."
        
        # Filter relevant results
        relevant_results = [r for r in results if r['score'] > 0.3]
        
        if not relevant_results:
            return f"❌ The document contains some text, but nothing specifically relevant to '{query}'. Please ask about content that might be in the document."
        
        # Build response from document content
        response = "📚 **Based on your document:**\n\n"
        
        for i, result in enumerate(relevant_results[:3]):  # Show top 3 results
            response += f"**• Section {i+1}:** {result['content']}\n\n"
        
        # Add suggestions
        response += "💡 **Tip:** You can ask about:\n- Key topics in the document\n- Specific information you're looking for\n- Summaries of sections\n- Explanations of concepts mentioned"
        
        return response

def create_interface():
    bot = PureRAGBot()
    
    with gr.Blocks(theme=gr.themes.Soft(), title="Document RAG Assistant") as demo:
        gr.Markdown("""
        # 📚 Document RAG Assistant
        
        **Pure Document-Based Question Answering**
        
        - **🔍 Semantic Search**: Find relevant content in your documents
        - **📖 Content-Based Answers**: All answers come directly from your uploaded files
        - **🎯 Precision**: Only answers based on document content
        
        **Note**: This bot only answers questions based on your uploaded documents.
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📁 Upload Document")
                file_input = gr.File(
                    label="Upload your document",
                    file_types=[".txt", ".csv", ".pdf", ".docx", ".doc"],
                    type="filepath"
                )
                upload_btn = gr.Button("Process Document", variant="primary")
                status = gr.Textbox(
                    label="Status",
                    value="Please upload a document to begin...",
                    interactive=False
                )
                
                gr.Markdown("""
                ### ℹ️ How It Works
                1. Upload any document (TXT, PDF, CSV, DOCX)
                2. Ask questions about the content
                3. Get answers directly from the document
                4. No general knowledge - only document content
                """)
                
            with gr.Column(scale=2):
                gr.Markdown("### 💬 Ask About Your Document")
                chatbot = gr.Chatbot(
                    height=400,
                    label="Document Q&A",
                    show_copy_button=True,
                    placeholder="Ask questions about your uploaded document content..."
                )
                with gr.Row():
                    question_input = gr.Textbox(
                        label="Your question about the document",
                        placeholder="What would you like to know about this document?",
                        scale=4
                    )
                    send_btn = gr.Button("Search Document", variant="primary", scale=1)
                
                clear_btn = gr.Button("Clear Conversation", variant="secondary")
                
                # gr.Markdown("""
                # ### 💡 Example Questions:
                # **After uploading a document, try:**
                # - "What is the main topic of this document?"
                # - "Summarize the key points"
                # - "What are the main findings?"
                # - "Explain the methodology used"
                # - "What solutions are proposed?"
                # - "List the key recommendations"
                # - "What data is presented in this report?"
                # """)
        
        def process_file(file):
            return bot.load_file_content(file)
        
        def respond(message, chat_history):
            if not message.strip():
                return "", chat_history
            
            response = bot.generate_rag_response(message)
            chat_history.append((message, response))
            return "", chat_history
        
        def clear_chat():
            return []
        
        # Event handlers
        upload_btn.click(
            process_file,
            inputs=[file_input],
            outputs=[status]
        )
        
        question_input.submit(
            respond,
            inputs=[question_input, chatbot],
            outputs=[question_input, chatbot]
        )
        
        send_btn.click(
            respond,
            inputs=[question_input, chatbot],
            outputs=[question_input, chatbot]
        )
        
        clear_btn.click(
            clear_chat,
            outputs=[chatbot]
        )
    
    return demo

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