import gradio as gr
import PyPDF2
import fitz  # PyMuPDF
import io
import os
from groq import Groq
import tempfile
import traceback
import base64
from io import BytesIO
from PIL import Image

def extract_text_from_pdf(pdf_file):
    """Extract text from uploaded PDF file"""
    try:
        if pdf_file is None:
            return None, "No PDF file provided"
        
        # Read the PDF file
        with open(pdf_file.name, 'rb') as file:
            pdf_reader = PyPDF2.PdfReader(file)
            text = ""
            
            # Extract text from all pages
            for page_num in range(len(pdf_reader.pages)):
                page = pdf_reader.pages[page_num]
                text += page.extract_text() + "\n"
            
            if not text.strip():
                return None, "Could not extract text from PDF. The PDF might be image-based or encrypted."
            
            return text, None
    
    except Exception as e:
        return None, f"Error reading PDF: {str(e)}"

def extract_images_from_pdf(pdf_file_path):
    """Extract all images from PDF and return them as PIL Images with page info"""
    images = []
    
    # Open the PDF
    pdf_document = fitz.open(pdf_file_path)
    
    for page_num in range(len(pdf_document)):
        page = pdf_document.load_page(page_num)
        
        # Get image list from the page
        image_list = page.get_images(full=True)
        
        for img_index, img in enumerate(image_list):
            # Get the XREF of the image
            xref = img[0]
            
            # Extract the image bytes
            base_image = pdf_document.extract_image(xref)
            image_bytes = base_image["image"]
            
            # Convert to PIL Image
            pil_image = Image.open(BytesIO(image_bytes))
            
            images.append({
                'image': pil_image,
                'page': page_num + 1,
                'index': img_index + 1,
                'format': base_image["ext"]
            })
    
    pdf_document.close()
    return images

def convert_pdf_pages_to_images(pdf_file_path, dpi=150):
    """Convert each PDF page to an image for comprehensive analysis"""
    images = []
    
    # Open the PDF
    pdf_document = fitz.open(pdf_file_path)
    
    for page_num in range(len(pdf_document)):
        page = pdf_document.load_page(page_num)
        
        # Convert page to image
        mat = fitz.Matrix(dpi/72, dpi/72)  # zoom factor
        pix = page.get_pixmap(matrix=mat)
        
        # Convert to PIL Image
        img_data = pix.tobytes("png")
        pil_image = Image.open(BytesIO(img_data))
        
        images.append({
            'image': pil_image,
            'page': page_num + 1,
            'type': 'full_page'
        })
    
    pdf_document.close()
    return images

def encode_image_to_base64(pil_image):
    """Convert PIL Image to base64 string"""
    buffered = BytesIO()
    pil_image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    return f"data:image/png;base64,{img_str}"

def summarize_with_groq(api_key, text):
    """Generate summary using Groq API"""
    try:
        if not api_key or not api_key.strip():
            return "Please enter your Groq API key"
        
        if not text or not text.strip():
            return "No text to summarize"
        
        # Initialize Groq client
        client = Groq(api_key=api_key.strip())
        
        # System prompt for summarization
        system_prompt = """You are a highly capable language model specialized in document summarization. Your task is to read and understand the full content of a multi-page PDF document and generate a clear, accurate, and detailed summary of the entire document.
Focus on capturing all main ideas, key points, arguments, findings, and conclusions presented throughout the document. If the document is technical, legal, academic, or contains structured sections (e.g., introduction, methods, results, discussion), maintain the logical flow and structure while expressing the content in a comprehensive and accessible manner.
Avoid unnecessary simplification. Include important details, supporting evidence, and nuanced insights that reflect the depth of the original material. Do not copy the text verbatim.
Output only the summary. Do not explain your process. Use a neutral, professional, and informative tone. The summary should provide a full understanding of the document to someone who has not read it."""

        # Create completion with optimal hyperparameters
        completion = client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=[
                {
                    "role": "system",
                    "content": system_prompt
                },
                {
                    "role": "user", 
                    "content": f"Please summarize the following document:\n\n{text}"
                }
            ],
            temperature=0.3,        # Low randomness for factual, focused summaries
            max_completion_tokens=2048,  # Increased to allow longer summaries
            top_p=0.9,             # Allows some diversity while still grounded
            stream=False,
            stop=None,
        )
        
        summary = completion.choices[0].message.content
        return summary
    
    except Exception as e:
        error_msg = f"Error generating summary: {str(e)}"
        if "authentication" in str(e).lower() or "api" in str(e).lower():
            error_msg += "\n\nPlease check your Groq API key and ensure it's valid."
        return error_msg

def analyze_images_with_groq(api_key, images):
    """Analyze images using Groq API"""
    if not api_key:
        return "❌ Please enter your Groq API key."
    
    try:
        client = Groq(api_key=api_key)
        
        results = []
        
        for idx, img_data in enumerate(images):
            # Encode image to base64
            base64_image = encode_image_to_base64(img_data['image'])
            
            # Prepare messages for the API call
            messages = [
                {
                    "role": "system",
                    "content": """You are an advanced language model with strong capabilities in visual and textual understanding. Your task is to analyze all images, diagrams, and flowcharts within a PDF document. This includes:

1. Extracting and interpreting text from images and flowcharts.
2. Understanding the visual structure, logic, and relationships depicted in diagrams.
3. Summarizing the content and purpose of each visual element in a clear and informative manner.

After processing, be ready to answer user questions about any of the images or flowcharts, including their meaning, structure, data, process flows, or relationships shown.

Be accurate, concise, and visually aware. Clearly explain visual content in text form. Do not guess if visual data is unclear or ambiguous — instead, state what is observable.

Use a neutral, helpful tone. Do not include irrelevant information or commentary unrelated to the visual content. When summarizing or answering questions, assume the user may not have access to the original image or diagram."""
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": f"Please analyze this image from page {img_data.get('page', 'unknown')} of the PDF document. Provide a detailed analysis of all visual elements, text, diagrams, flowcharts, and their relationships."
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": base64_image
                            }
                        }
                    ]
                }
            ]
            
            # Make API call with optimal parameters
            completion = client.chat.completions.create(
                model="meta-llama/llama-4-scout-17b-16e-instruct",
                messages=messages,
                temperature=0.2,
                max_completion_tokens=2048,
                top_p=0.85,
                stream=False
            )
            
            analysis = completion.choices[0].message.content
            
            page_info = f"Page {img_data.get('page', 'N/A')}"
            if 'index' in img_data:
                page_info += f", Image {img_data['index']}"
            elif 'type' in img_data and img_data['type'] == 'full_page':
                page_info += " (Full Page)"
            
            results.append(f"## 📄 {page_info}\n\n{analysis}\n\n---\n")
        
        if not results:
            return "⚠️ No images found in the PDF document."
        
        return "\n".join(results)
        
    except Exception as e:
        return f"❌ Error analyzing images: {str(e)}"

def process_pdf_text_summary(api_key, pdf_file, progress=gr.Progress()):
    """Process PDF and generate text summary"""
    try:
        if not api_key or not api_key.strip():
            return "❌ Please enter your Groq API key", "", ""
        
        if pdf_file is None:
            return "❌ Please upload a PDF file", "", ""
        
        progress(0.1, desc="Reading PDF file...")
        
        # Extract text from PDF
        text, error = extract_text_from_pdf(pdf_file)
        if error:
            return f"❌ {error}", "", ""
        
        progress(0.4, desc="Text extracted successfully...")
        
        # Show preview of extracted text
        text_preview = text[:500] + "..." if len(text) > 500 else text
        
        progress(0.6, desc="Generating summary with Groq AI...")
        
        # Generate summary
        summary = summarize_with_groq(api_key, text)
        
        progress(1.0, desc="Summary generated successfully!")
        
        return "✅ Summary generated successfully!", text_preview, summary
    
    except Exception as e:
        error_traceback = traceback.format_exc()
        return f"❌ Unexpected error: {str(e)}\n\nTraceback:\n{error_traceback}", "", ""

def process_pdf_image_analysis(api_key, pdf_file, analysis_method, progress=gr.Progress()):
    """Process PDF and analyze images"""
    if pdf_file is None:
        return "⚠️ Please upload a PDF file."
    
    if not api_key or api_key.strip() == "":
        return "⚠️ Please enter your Groq API key."
    
    try:
        progress(0.1, desc="Processing PDF file...")
        
        # Create temporary file for PDF processing
        with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
            # Write uploaded file content to temporary file
            if hasattr(pdf_file, 'read'):
                tmp_file.write(pdf_file.read())
            else:
                with open(pdf_file.name, 'rb') as f:
                    tmp_file.write(f.read())
            tmp_file_path = tmp_file.name
        
        progress(0.3, desc="Extracting images...")
        
        images_to_analyze = []
        
        if analysis_method == "Extract embedded images only":
            # Extract only embedded images
            images_to_analyze = extract_images_from_pdf(tmp_file_path)
            if not images_to_analyze:
                return "⚠️ No embedded images found in the PDF. Try 'Full page analysis' to analyze the entire content."
        
        elif analysis_method == "Full page analysis":
            # Convert each page to image for comprehensive analysis
            images_to_analyze = convert_pdf_pages_to_images(tmp_file_path)
        
        else:  # Both methods
            # First try embedded images
            embedded_images = extract_images_from_pdf(tmp_file_path)
            # Then add full page analysis
            page_images = convert_pdf_pages_to_images(tmp_file_path)
            images_to_analyze = embedded_images + page_images
        
        # Clean up temporary file
        os.unlink(tmp_file_path)
        
        if not images_to_analyze:
            return "⚠️ No visual content found in the PDF document."
        
        progress(0.7, desc="Analyzing images with AI...")
        
        # Analyze images with Groq
        analysis_result = analyze_images_with_groq(api_key, images_to_analyze)
        
        progress(1.0, desc="Analysis complete!")
        
        return analysis_result
        
    except Exception as e:
        return f"❌ Error processing PDF: {str(e)}"

def clear_all_text():
    """Clear all text analysis fields"""
    return "", None, "", "", ""

def clear_all_image():
    """Clear all image analysis fields"""
    return "", None, "Full page analysis", ""

# Custom CSS for better styling
css = """
.gradio-container {
    max-width: 1400px !important;
    margin: auto !important;
}
.main-header {
    text-align: center;
    margin-bottom: 2rem;
}
.status-success {
    color: #28a745 !important;
}
.status-error {
    color: #dc3545 !important;
}
.info-box {
    background-color: #f8f9fa;
    padding: 1rem;
    border-radius: 0.5rem;
    border-left: 4px solid #007bff;
    margin: 1rem 0;
}
.feature-box {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    padding: 1.5rem;
    border-radius: 1rem;
    color: white;
    margin: 1rem 0;
}
"""

# Create Gradio interface
with gr.Blocks(css=css, title="Advanced PDF Analyzer with Groq AI") as demo:
    # Header
    gr.HTML("""
    <div class="main-header">
        <h1>🚀 Advanced PDF Analyzer with Groq AI</h1>
        <p>Comprehensive PDF analysis tool - Extract text summaries and analyze images/diagrams using state-of-the-art AI models</p>
    </div>
    """)
    
    # Feature overview
    gr.HTML("""
    <div class="feature-box">
        <h3>✨ What this tool can do:</h3>
        <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 1rem; margin-top: 1rem;">
            <div>
                <h4>📝 Text Analysis:</h4>
                <ul>
                    <li>Extract and summarize text content</li>
                    <li>Generate comprehensive document summaries</li>
                    <li>Maintain logical structure and key insights</li>
                </ul>
            </div>
            <div>
                <h4>🖼️ Visual Analysis:</h4>
                <ul>
                    <li>Analyze embedded images and diagrams</li>
                    <li>Process flowcharts and technical drawings</li>
                    <li>Extract text from images (OCR)</li>
                </ul>
            </div>
        </div>
    </div>
    """)
    
    # API Key section
    gr.HTML("""
    <div class="info-box">
        <strong>🔑 How to get your Groq API Key:</strong><br>
        1. Visit <a href="https://console.groq.com/" target="_blank">console.groq.com</a><br>
        2. Sign up or log in to your account<br>
        3. Navigate to API Keys section<br>
        4. Create a new API key and copy it<br>
        5. Paste it in the field below
    </div>
    """)
    
    # Global API key input
    api_key_input = gr.Textbox(
        label="🔑 Groq API Key",
        placeholder="Enter your Groq API key here (used for both text and image analysis)...",
        type="password",
        info="Your API key is not stored and only used for this session"
    )
    
    # Tabs for different functionalities
    with gr.Tabs():
        # Text Summary Tab
        with gr.TabItem("📝 Text Summary", elem_id="text-tab"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("## 📄 Text Analysis")
                    
                    pdf_file_text = gr.File(
                        label="Upload PDF Document",
                        file_types=[".pdf"]
                    )
                    
                    with gr.Row():
                        summarize_btn = gr.Button("📋 Generate Text Summary", variant="primary", size="lg")
                        clear_text_btn = gr.Button("🗑️ Clear All", variant="secondary")
                
                with gr.Column(scale=2):
                    gr.Markdown("## 📊 Text Analysis Results")
                    status_text_output = gr.Textbox(
                        label="Status",
                        interactive=False,
                        show_label=True
                    )
                    
                    with gr.Tabs():
                        with gr.TabItem("📝 Summary"):
                            summary_output = gr.Textbox(
                                label="AI Generated Summary",
                                lines=15,
                                interactive=False,
                                placeholder="Your PDF summary will appear here...",
                                show_copy_button=True
                            )
                        
                        with gr.TabItem("📄 Extracted Text Preview"):
                            text_preview_output = gr.Textbox(
                                label="Extracted Text (First 500 characters)",
                                lines=10,
                                interactive=False,
                                placeholder="Preview of extracted text will appear here...",
                                show_copy_button=True
                            )
        
        # Image Analysis Tab
        with gr.TabItem("🖼️ Image Analysis", elem_id="image-tab"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("## 🔍 Visual Analysis")
                    
                    pdf_file_image = gr.File(
                        label="Upload PDF Document",
                        file_types=[".pdf"]
                    )
                    
                    analysis_method = gr.Radio(
                        choices=[
                            "Extract embedded images only",
                            "Full page analysis", 
                            "Both (embedded + full pages)"
                        ],
                        value="Full page analysis",
                        label="Analysis Method",
                        info="Choose how to analyze the PDF content"
                    )
                    
                    with gr.Row():
                        analyze_images_btn = gr.Button("🔍 Analyze Images", variant="primary", size="lg")
                        clear_image_btn = gr.Button("🗑️ Clear All", variant="secondary")
                
                with gr.Column(scale=2):
                    gr.Markdown("## 📊 Image Analysis Results")
                    image_analysis_output = gr.Textbox(
                        label="Visual Analysis Results",
                        lines=20,
                        max_lines=50,
                        show_copy_button=True,
                        placeholder="Image analysis results will appear here..."
                    )
    
    # Usage instructions
    gr.HTML("""
    <div class="info-box">
        <strong>📋 Usage Instructions:</strong><br>
        <h4>For Text Summary:</h4>
        1. Enter your Groq API key above<br>
        2. Go to "Text Summary" tab<br>
        3. Upload a PDF document<br>
        4. Click "Generate Text Summary"<br>
        
        <h4>For Image Analysis:</h4>
        1. Enter your Groq API key above<br>
        2. Go to "Image Analysis" tab<br>
        3. Upload a PDF document<br>
        4. Choose analysis method<br>
        5. Click "Analyze Images"<br>
        
        <h4>Analysis Methods:</h4>
        • <strong>Extract embedded images:</strong> Analyzes only images embedded within the PDF<br>
        • <strong>Full page analysis:</strong> Converts each page to image for comprehensive analysis (recommended)<br>
        • <strong>Both:</strong> Combines both methods for maximum coverage
    </div>
    """)
    
    # Model information
    gr.HTML("""
    <div style="margin-top: 2rem; padding: 1rem; background-color: #e9ecef; border-radius: 0.5rem;">
        <strong>🤖 Model Information:</strong><br>
        • <strong>Text Analysis:</strong> Llama-3.3-70B-Versatile (optimized for summarization)<br>
        • <strong>Image Analysis:</strong> Llama-3.3-70B-Versatile (with vision capabilities)<br>
        • <strong>Temperature:</strong> 0.2-0.3 (focused, factual analysis)<br>
        • <strong>Max Tokens:</strong> 2048 (comprehensive outputs)
    </div>
    """)
    
    # Event handlers for text summary
    summarize_btn.click(
        fn=process_pdf_text_summary,
        inputs=[api_key_input, pdf_file_text],
        outputs=[status_text_output, text_preview_output, summary_output],
        show_progress=True
    )
    
    clear_text_btn.click(
        fn=clear_all_text,
        outputs=[api_key_input, pdf_file_text, status_text_output, text_preview_output, summary_output]
    )
    
    # Event handlers for image analysis
    analyze_images_btn.click(
        fn=process_pdf_image_analysis,
        inputs=[api_key_input, pdf_file_image, analysis_method],
        outputs=[image_analysis_output],
        show_progress=True
    )
    
    clear_image_btn.click(
        fn=clear_all_image,
        outputs=[api_key_input, pdf_file_image, analysis_method, image_analysis_output]
    )

# Launch the app
if __name__ == "__main__":
    print("🚀 Starting Advanced PDF Analyzer App...")
    print("📋 Make sure you have the required packages installed:")
    print("   pip install -r requirements.txt")
    print("\n🔑 Don't forget to get your Groq API key from: https://console.groq.com/")
    
    demo.launch(
        server_name="0.0.0.0",  # Allow external connections
        server_port=7860,       # Default port
        share=True,             # Create shareable link
        debug=True,             # Enable debug mode
        show_error=True         # Show detailed errors
    )