import os
import re
import tempfile
import requests
import gradio as gr
from PyPDF2 import PdfReader
import logging
import webbrowser
from gradio_client import Client

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Initialize Hugging Face models
HUGGINGFACE_MODELS = {
    "Phi-3 Mini 128k": "eswardivi/Phi-3-mini-128k-instruct",
}

# Common context window sizes
CONTEXT_SIZES = {
    "4K": 4000,
    "8K": 8000,
    "32K": 32000,
    "128K": 128000,
    "200K": 200000
}

def copy_to_clipboard(text):
    return text

def open_chatgpt():
    webbrowser.open('https://chat.openai.com/')
    return "Opening ChatGPT in browser..."

# Utility Functions
def extract_text_from_pdf(pdf_path):
    """Extract text content from PDF file."""
    try:
        reader = PdfReader(pdf_path)
        text = ""
        for page_num, page in enumerate(reader.pages, start=1):
            page_text = page.extract_text()
            if page_text:
                text += page_text + "\n"
            else:
                logging.warning(f"No text found on page {page_num}.")
        if not text.strip():
            return "Error: No extractable text found in the PDF."
        return text
    except Exception as e:
        logging.error(f"Error reading PDF file: {e}")
        return f"Error reading PDF file: {e}"

def format_content(text, format_type):
    """Format extracted text according to specified format."""
    if format_type == 'txt':
        return text
    elif format_type == 'md':
        paragraphs = text.split('\n\n')
        return '\n\n'.join(paragraphs)
    elif format_type == 'html':
        paragraphs = text.split('\n\n')
        return ''.join([f'<p>{para.strip()}</p>' for para in paragraphs if para.strip()])
    else:
        logging.error(f"Unsupported format: {format_type}")
        return f"Unsupported format: {format_type}"

def split_into_snippets(text, context_size):
    """Split text into manageable snippets based on context size."""
    sentences = re.split(r'(?<=[.!?]) +', text)
    snippets = []
    current_snippet = ""

    for sentence in sentences:
        if len(current_snippet) + len(sentence) + 1 > context_size:
            if current_snippet:
                snippets.append(current_snippet.strip())
                current_snippet = sentence + " "
            else:
                snippets.append(sentence.strip())
                current_snippet = ""
        else:
            current_snippet += sentence + " "

    if current_snippet.strip():
        snippets.append(current_snippet.strip())

    return snippets

def build_prompts(snippets, prompt_instruction, custom_prompt, snippet_num=None):
    """Build formatted prompts from text snippets."""
    if snippet_num is not None:
        if 1 <= snippet_num <= len(snippets):
            selected_snippets = [snippets[snippet_num - 1]]
        else:
            return f"Error: Invalid snippet number. Please choose between 1 and {len(snippets)}."
    else:
        selected_snippets = snippets

    prompts = []
    base_prompt = custom_prompt if custom_prompt else prompt_instruction
    
    for idx, snippet in enumerate(selected_snippets, start=1):
        if len(selected_snippets) > 1:
            prompt_header = f"{base_prompt} Part {idx} of {len(selected_snippets)}: ---\n"
        else:
            prompt_header = f"{base_prompt} ---\n"
        
        framed_prompt = f"{prompt_header}{snippet}\n---"
        prompts.append(framed_prompt)
    
    return "\n\n".join(prompts)

def send_to_huggingface(prompt, model_name):
    """Send prompt to Hugging Face model using gradio_client."""
    try:
        client = Client(model_name)
        response = client.predict(
            prompt,  # message
            0.9,    # temperature
            True,   # sampling
            512,    # max_new_tokens
            api_name="/chat"
        )
        return response
    except Exception as e:
        logging.error(f"Error interacting with Hugging Face model: {e}")
        return f"Error interacting with Hugging Face model: {e}"

# Main Interface
with gr.Blocks(theme=gr.themes.Default()) as demo:
    # Header
    gr.Markdown("# 📄 Smart PDF Summarizer")
    gr.Markdown("Upload a PDF document and get AI-powered summaries using OpenAI or Hugging Face models.")
    
    # Main Content
    with gr.Row():
        # Left Column - Input Options
        with gr.Column(scale=1):
            pdf_input = gr.File(
                label="📁 Upload PDF",
                file_types=[".pdf"]
            )
            
            with gr.Row():
                format_type = gr.Radio(
                    choices=["txt", "md", "html"],
                    value="txt",
                    label="📝 Output Format"
                )
            
            gr.Markdown("### Context Window Size")
            with gr.Row():
                for size_name, size_value in CONTEXT_SIZES.items():
                    if gr.Button(size_name).click:
                        context_size.value = size_value
                        
            context_size = gr.Slider(
                minimum=1000,
                maximum=200000,
                step=1000,
                value=32000,
                label="📏 Custom Context Size"
            )
            
            snippet_number = gr.Number(
                label="🔢 Snippet Number",
                value=1,
                precision=0
            )
            
            custom_prompt = gr.Textbox(
                label="✍️ Custom Prompt",
                placeholder="Enter your custom prompt here...",
                lines=2
            )
            
            model_choice = gr.Radio(
                choices=["OpenAI ChatGPT", "Hugging Face Model"],
                value="OpenAI ChatGPT",
                label="🤖 Model Selection"
            )
            
            hf_model = gr.Dropdown(
                choices=list(HUGGINGFACE_MODELS.keys()),
                label="🔧 Hugging Face Model",
                visible=False
            )
            
            # Authentication moved down
            with gr.Row(visible=False) as auth_row:
                openai_api_key = gr.Textbox(
                    label="🔑 OpenAI API Key",
                    type="password",
                    placeholder="Enter your OpenAI API key (optional)"
                )

        # Right Column - Output
        with gr.Column(scale=1):
            with gr.Row():
                process_button = gr.Button("🚀 Process PDF", variant="primary")
                
            progress_status = gr.Textbox(
                label="📊 Progress",
                interactive=False
            )
            
            generated_prompt = gr.Textbox(
                label="📋 Generated Prompt",
                lines=10
            )
            
            with gr.Row():
                copy_prompt_button = gr.Button("📋 Copy Prompt")
                open_chatgpt_button = gr.Button("🌐 Open ChatGPT")
            
            summary_output = gr.Textbox(
                label="📝 Summary",
                lines=15
            )
            
            with gr.Row():
                copy_summary_button = gr.Button("📋 Copy Summary")
                download_files = gr.Files(
                    label="📥 Download Files"
                )

    # Event Handlers
    def toggle_hf_model(choice):
        return gr.update(visible=choice == "Hugging Face Model"), gr.update(visible=choice == "OpenAI ChatGPT")

    def process_pdf(pdf, fmt, ctx_size, snippet_num, prompt, model_selection, hf_model_choice):
        try:
            if not pdf:
                return "Please upload a PDF file.", "", "", None
            
            # Extract text
            text = extract_text_from_pdf(pdf.name)
            if text.startswith("Error"):
                return text, "", "", None
            
            # Format content
            formatted_text = format_content(text, fmt)
            
            # Split into snippets
            snippets = split_into_snippets(formatted_text, ctx_size)
            
            # Build prompts
            default_prompt = "Summarize the following text:"
            full_prompt = build_prompts(snippets, default_prompt, prompt, snippet_num)
            
            if isinstance(full_prompt, str) and full_prompt.startswith("Error"):
                return full_prompt, "", "", None
            
            # Generate summary based on model choice
            if model_selection == "Hugging Face Model":
                summary = send_to_huggingface(full_prompt, HUGGINGFACE_MODELS[hf_model_choice])
            else:
                summary = "Please use the Copy Prompt button and paste into ChatGPT."
            
            # Save files for download
            files_to_download = []
            
            with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.txt') as prompt_file:
                prompt_file.write(full_prompt)
                files_to_download.append(prompt_file.name)
                
            if summary != "Please use the Copy Prompt button and paste into ChatGPT.":
                with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.txt') as summary_file:
                    summary_file.write(summary)
                    files_to_download.append(summary_file.name)
            
            return "Processing complete!", full_prompt, summary, files_to_download
            
        except Exception as e:
            logging.error(f"Error processing PDF: {e}")
            return f"Error processing PDF: {str(e)}", "", "", None

    # Connect event handlers
    model_choice.change(
        toggle_hf_model,
        inputs=[model_choice],
        outputs=[hf_model, auth_row]
    )
    
    process_button.click(
        process_pdf,
        inputs=[
            pdf_input,
            format_type,
            context_size,
            snippet_number,
            custom_prompt,
            model_choice,
            hf_model
        ],
        outputs=[
            progress_status,
            generated_prompt,
            summary_output,
            download_files
        ]
    )
    
    copy_prompt_button.click(
        copy_to_clipboard,
        inputs=[generated_prompt],
        outputs=[progress_status]
    )
    
    copy_summary_button.click(
        copy_to_clipboard,
        inputs=[summary_output],
        outputs=[progress_status]
    )
    
    open_chatgpt_button.click(
        open_chatgpt,
        outputs=[progress_status]
    )

    # Instructions
    gr.Markdown("""
    ### 📌 Instructions:
    1. Upload a PDF document
    2. Choose output format and context window size
    3. Select snippet number (default: 1) or enter custom prompt
    4. Select between OpenAI ChatGPT or Hugging Face model
    5. Click 'Process PDF' to generate summary
    6. Use 'Copy Prompt' and 'Open ChatGPT' for manual processing
    7. Download generated files as needed

    ### ⚙️ Features:
    - Support for multiple PDF formats
    - Flexible text formatting options
    - Predefined context window sizes (4K to 200K)
    - Copy to clipboard functionality
    - Direct ChatGPT integration
    - Downloadable outputs
    """)

# Launch the interface
if __name__ == "__main__":
    demo.launch(share=False, debug=True)