import os
import gradio as gr
import tempfile
import shutil
from typing import List, Tuple
import numpy as np
from datetime import datetime
from io import BytesIO

# PDF Processing
from PyPDF2 import PdfReader

# Text Processing - FIXED IMPORT
from langchain_core.text_splitter import RecursiveCharacterTextSplitter

# Embeddings
from sentence_transformers import SentenceTransformer

# Vector Database
import faiss

# Groq LLM
from groq import Groq

# Document Generation
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, PageBreak
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import inch
from docx import Document
from docx.shared import Pt, Inches
from docx.enum.text import WD_ALIGN_PARAGRAPH

# Initialize Groq client
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

class RAGApplication:
    def __init__(self):
        self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
        self.dimension = 384
        self.index = None
        self.chunks = []
        self.current_pdf_name = None
        self.chat_history = []
        
    def extract_text_from_pdf(self, pdf_path: str) -> str:
        reader = PdfReader(pdf_path)
        text = ""
        for page in reader.pages:
            text += page.extract_text() + "\n"
        return text
    
    def create_chunks(self, text: str, chunk_size: int = 500, chunk_overlap: int = 50) -> List[str]:
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            length_function=len,
            separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""]
        )
        return text_splitter.split_text(text)
    
    def create_embeddings(self, chunks: List[str]) -> np.ndarray:
        return self.embedding_model.encode(chunks, show_progress_bar=True)
    
    def create_faiss_index(self, embeddings: np.ndarray):
        faiss.normalize_L2(embeddings)
        self.index = faiss.IndexFlatIP(self.dimension)
        self.index.add(embeddings)
    
    def process_pdf(self, pdf_file) -> str:
        if pdf_file is None:
            return "Please upload a PDF file."
        
        try:
            if isinstance(pdf_file, str):
                pdf_path = pdf_file
            else:
                with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
                    shutil.copyfileobj(pdf_file, tmp_file)
                    pdf_path = tmp_file.name
            
            text = self.extract_text_from_pdf(pdf_path)
            
            if not text.strip():
                return "Could not extract text from PDF."
            
            self.chunks = self.create_chunks(text)
            embeddings = self.create_embeddings(self.chunks)
            self.create_faiss_index(embeddings)
            
            if not isinstance(pdf_file, str) and os.path.exists(pdf_path):
                os.remove(pdf_path)
            
            self.current_pdf_name = os.path.basename(pdf_path) if isinstance(pdf_file, str) else "uploaded.pdf"
            self.chat_history = []
            
            return f"✅ Successfully processed PDF!\n📄 Document: {self.current_pdf_name}\n📝 Total chunks: {len(self.chunks)}"
            
        except Exception as e:
            return f"❌ Error: {str(e)}"
    
    def search_similar_chunks(self, query: str, k: int = 5) -> List[str]:
        if self.index is None or len(self.chunks) == 0:
            return []
        
        query_embedding = self.embedding_model.encode([query])
        faiss.normalize_L2(query_embedding)
        scores, indices = self.index.search(query_embedding, k)
        return [self.chunks[idx] for idx in indices[0] if idx < len(self.chunks)]
    
    def generate_response(self, query: str, context: List[str]) -> str:
        context_text = "\n\n".join([f"Chunk {i+1}:\n{chunk}" for i, chunk in enumerate(context)])
        
        prompt = f"""You are a helpful assistant that answers questions based on the provided document context.

Context from the document:
{context_text}

User Question: {query}

Please provide a comprehensive answer based on the context above. If the context doesn't contain enough information to answer the question, say so clearly.

Answer:"""
        
        chat_completion = client.chat.completions.create(
            messages=[
                {"role": "system", "content": "You are a helpful assistant that answers questions based on provided document context."},
                {"role": "user", "content": prompt}
            ],
            model="llama-3.3-70b-versatile",
            temperature=0.7,
            max_tokens=1024
        )
        
        return chat_completion.choices[0].message.content
    
    def chat(self, message: str, history: List[Tuple[str, str]]) -> Tuple[str, List[Tuple[str, str]]]:
        if self.index is None:
            return "Please upload a PDF document first!", history
        
        if not message.strip():
            return "Please enter a question.", history
        
        try:
            relevant_chunks = self.search_similar_chunks(message, k=5)
            
            if not relevant_chunks:
                response = "I couldn't find relevant information in the document to answer your question."
            else:
                response = self.generate_response(message, relevant_chunks)
            
            self.chat_history.append({
                "question": message,
                "answer": response,
                "context": relevant_chunks,
                "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            })
            
            history.append((message, response))
            return "", history
            
        except Exception as e:
            error_msg = f"Error generating response: {str(e)}"
            history.append((message, error_msg))
            return "", history
    
    def clear_chat(self):
        self.index = None
        self.chunks = []
        self.current_pdf_name = None
        self.chat_history = []
        return None, []
    
    def generate_pdf_report(self) -> str:
        if not self.chat_history:
            return None
        
        buffer = BytesIO()
        doc = SimpleDocTemplate(buffer, pagesize=letter, topMargin=1*inch, bottomMargin=1*inch)
        styles = getSampleStyleSheet()
        
        title_style = ParagraphStyle(
            'CustomTitle',
            parent=styles['Heading1'],
            fontSize=24,
            spaceAfter=30,
            textColor='#2C3E50'
        )
        
        question_style = ParagraphStyle(
            'QuestionStyle',
            parent=styles['Heading2'],
            fontSize=14,
            spaceAfter=12,
            textColor='#2980B9'
        )
        
        answer_style = ParagraphStyle(
            'AnswerStyle',
            parent=styles['BodyText'],
            fontSize=11,
            spaceAfter=20,
            leading=14
        )
        
        context_style = ParagraphStyle(
            'ContextStyle',
            parent=styles['BodyText'],
            fontSize=9,
            textColor='#7F8C8D',
            leftIndent=20
        )
        
        story = []
        story.append(Paragraph("RAG Chat Report", title_style))
        story.append(Paragraph(f"Document: {self.current_pdf_name or 'N/A'}", styles['Normal']))
        story.append(Paragraph(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", styles['Normal']))
        story.append(Spacer(1, 20))
        
        for i, qa in enumerate(self.chat_history, 1):
            story.append(Paragraph(f"Q{i}: {qa['question']}", question_style))
            answer_text = qa['answer'].replace('\n', '<br/>')
            story.append(Paragraph(f"<b>Answer:</b> {answer_text}", answer_style))
            story.append(Paragraph("<b>Source Context:</b>", styles['Heading3']))
            
            for j, ctx in enumerate(qa['context'][:2], 1):
                ctx_text = ctx[:300] + "..." if len(ctx) > 300 else ctx
                ctx_text = ctx_text.replace('\n', '<br/>')
                story.append(Paragraph(f"Context {j}: {ctx_text}", context_style))
            
            story.append(Spacer(1, 20))
            
            if i % 3 == 0 and i < len(self.chat_history):
                story.append(PageBreak())
        
        doc.build(story)
        buffer.seek(0)
        
        temp_path = f"/tmp/rag_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf"
        with open(temp_path, 'wb') as f:
            f.write(buffer.getvalue())
        
        return temp_path
    
    def generate_word_report(self) -> str:
        if not self.chat_history:
            return None
        
        doc = Document()
        title = doc.add_heading('RAG Chat Report', 0)
        title.alignment = WD_ALIGN_PARAGRAPH.CENTER
        
        doc.add_paragraph(f"Document: {self.current_pdf_name or 'N/A'}")
        doc.add_paragraph(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        doc.add_paragraph()
        
        for i, qa in enumerate(self.chat_history, 1):
            question_para = doc.add_paragraph()
            question_run = question_para.add_run(f"Q{i}: {qa['question']}")
            question_run.bold = True
            question_run.font.size = Pt(14)
            
            answer_para = doc.add_paragraph()
            answer_run = answer_para.add_run("Answer: ")
            answer_run.bold = True
            answer_para.add_run(qa['answer'])
            
            context_heading = doc.add_paragraph()
            context_run = context_heading.add_run("Source Context:")
            context_run.bold = True
            context_run.font.size = Pt(10)
            
            for j, ctx in enumerate(qa['context'][:2], 1):
                ctx_text = ctx[:300] + "..." if len(ctx) > 300 else ctx
                ctx_para = doc.add_paragraph(ctx_text, style='List Bullet')
                ctx_para.paragraph_format.left_indent = Inches(0.5)
            
            doc.add_paragraph()
        
        temp_path = f"/tmp/rag_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.docx"
        doc.save(temp_path)
        
        return temp_path

rag_app = RAGApplication()

def create_interface():
    with gr.Blocks(title="RAG PDF Chat with Export") as demo:
        gr.Markdown("""
        # 📚 RAG PDF Chat Application
        ### Upload a PDF, ask questions, and download your Q&A history!
        
        **Powered by:** 
        - 🦙 Llama 3.3 70B (via Groq)
        - 🔍 FAISS Vector Database
        - 📄 PDF & Word Export
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📤 Upload Document")
                pdf_input = gr.File(
                    label="Upload PDF",
                    file_types=[".pdf"],
                    type="filepath"
                )
                process_btn = gr.Button("🚀 Process PDF", variant="primary")
                status_output = gr.Textbox(
                    label="Status",
                    lines=4,
                    interactive=False
                )
                
                gr.Markdown("### 💾 Export Options")
                with gr.Row():
                    download_pdf_btn = gr.Button("📄 Download PDF", variant="secondary")
                    download_word_btn = gr.Button("📝 Download Word", variant="secondary")
                
                pdf_file_output = gr.File(label="PDF Report", visible=False)
                word_file_output = gr.File(label="Word Report", visible=False)
                
                clear_btn = gr.Button("🗑️ Clear All", variant="stop")
            
            with gr.Column(scale=2):
                gr.Markdown("### 💬 Chat with your Document")
                chatbot = gr.Chatbot(
                    height=500,
                    bubble_full_width=False,
                    show_copy_button=True
                )
                msg_input = gr.Textbox(
                    label="Your Question",
                    placeholder="Ask something about the uploaded document...",
                    lines=2
                )
                send_btn = gr.Button("Send", variant="primary")
        
        process_btn.click(
            fn=rag_app.process_pdf,
            inputs=pdf_input,
            outputs=status_output
        )
        
        msg_input.submit(
            fn=rag_app.chat,
            inputs=[msg_input, chatbot],
            outputs=[msg_input, chatbot]
        )
        
        send_btn.click(
            fn=rag_app.chat,
            inputs=[msg_input, chatbot],
            outputs=[msg_input, chatbot]
        )
        
        def handle_pdf_download():
            file_path = rag_app.generate_pdf_report()
            if file_path:
                return gr.update(value=file_path, visible=True)
            else:
                return gr.update(value=None, visible=False)
        
        def handle_word_download():
            file_path = rag_app.generate_word_report()
            if file_path:
                return gr.update(value=file_path, visible=True)
            else:
                return gr.update(value=None, visible=False)
        
        download_pdf_btn.click(
            fn=handle_pdf_download,
            inputs=None,
            outputs=pdf_file_output
        )
        
        download_word_btn.click(
            fn=handle_word_download,
            inputs=None,
            outputs=word_file_output
        )
        
        clear_btn.click(
            fn=rag_app.clear_chat,
            inputs=None,
            outputs=[pdf_input, chatbot]
        )
        
        gr.Markdown("""
        ---
        ### 📝 How to use:
        1. **Upload** your PDF document
        2. **Process** the PDF to create embeddings
        3. **Ask questions** in the chat
        4. **Download** your Q&A as PDF or Word document
        
        **Note:** Set your Groq API key in the Space secrets.
        """)
    
    return demo

demo = create_interface()

if __name__ == "__main__":
    demo.launch()