import gradio as gr
import openai
import fitz  # PyMuPDF for PDF processing
import os
import tempfile
import base64
from datetime import datetime

# Variable to store API key
api_key = ""

# Function to update API key
def set_api_key(key):
    global api_key
    api_key = key
    return "API Key Set Successfully!"

# Function to extract text from PDF
def extract_text_from_pdf(pdf_path):
    try:
        doc = fitz.open(pdf_path)
        text = "\n".join([page.get_text("text") for page in doc])
        return text
    except Exception as e:
        return f"Error extracting text from PDF: {str(e)}"

# Function to interact with OpenAI API for systematic review
def generate_systematic_review(pdf_files, review_question, include_tables=True):
    if not api_key:
        return "Please enter your OpenAI API key first."
    
    if not pdf_files:
        return "Please upload at least one PDF file."
    
    if not review_question:
        return "Please enter a review question."
    
    try:
        openai.api_key = api_key
        
        # Create the system message with systematic review guidelines
        system_prompt = """You are an expert academic assistant. Create a systematic review using academic research paper formatting. The Systematic Review must be in great details. Structure it using these steps:
        
        Step 1: Identify a Research Field
        The first step in writing a systematic review paper is to identify a research field. This involves selecting a specific area of study that you are interested in and want to explore further.
        
        Step 2: Generate a Research Question
        Once you have identified your research field, the next step is to generate a research question. This question should be specific, measurable, achievable, relevant, and time-bound (SMART). Consider using the PICO framework (Population, Intervention, Comparison, Outcome) to structure clinical questions.
        
        Step 3: Create a Protocol
        After generating your research question, create a detailed protocol. This is a comprehensive plan outlining your research methodology, including search strategies, databases to be used, and analysis techniques. The protocol should be registered in appropriate databases (e.g., PROSPERO) when applicable and follow PRISMA guidelines.
        
        Step 4: Define Inclusion and Exclusion Criteria
        Clearly articulate the criteria for including or excluding studies in your review. These criteria should be directly tied to your research question and may include: publication date range, study types, population characteristics, intervention specifications, outcome measures, and language restrictions.
        
        Step 5: Evaluate Relevant Literature
        Conduct a comprehensive literature search using multiple databases (e.g., PubMed, Scopus, Web of Science, CINAHL) with clearly defined search terms and Boolean operators. Document your search strategy in detail to ensure reproducibility. Consider both published and unpublished (gray) literature to minimize publication bias.
        
        Step 6: Quality Assessment of Studies
        Apply established quality assessment tools appropriate for your included study designs (e.g., Cochrane Risk of Bias Tool for randomized trials, ROBINS-I for non-randomized studies, CASP checklists). Document quality assessments for all included studies.
        
        Step 7: Data Extraction
        Create and use a standardized data extraction form to systematically collect relevant information from each study. This should include: study characteristics, participant demographics, intervention details, comparison groups, outcome measures, and results. Have multiple reviewers extract data independently when possible.
        
        Step 8: Data Synthesis and Analysis
        Synthesize the extracted data using appropriate methods. If statistical pooling is appropriate, conduct meta-analysis with suitable models (fixed or random effects). If heterogeneity precludes meta-analysis, provide a narrative synthesis with clear explanations.
        
        Step 9: Critical Analysis of Results
        Analyze your findings critically, examining patterns, inconsistencies, and relationships across studies. Address heterogeneity, publication bias (using funnel plots when applicable), and methodological limitations. Consider using the GRADE approach to evaluate certainty of evidence.
        
        Step 10: Interpreting Findings
        Interpret your findings in the context of the original research question. Discuss implications for practice, policy, and future research. Address limitations of your systematic review process and any potential biases.
        
        Step 11: Concluding Statements
        Provide clear, substantiated conclusions based on your review findings. Ensure conclusions are proportionate to the evidence presented and acknowledge uncertainty where appropriate. Offer specific recommendations for future research.
        
        Step 12: References and Documentation
        Include a comprehensive reference list following a specific citation style (APA, Vancouver, etc.). Provide links to source papers when available.
        
        Your response should be formatted in HTML (but avoid showing these tags ```html ```) but generate the content to look like a professional academic paper. Include proper section headers, abstracts, methodology sections, etc. Number all sections like an academic paper. Follow academic journal standards with double spacing, appropriate margins, and consistent formatting throughout.
        """
        
        # Extract text from each PDF
        pdf_texts = []
        pdf_names = []
        
        for pdf_file in pdf_files:
            if isinstance(pdf_file, str):  # If it's already a path
                pdf_path = pdf_file
            else:  # If it's a file object
                pdf_path = pdf_file.name
                
            pdf_name = os.path.basename(pdf_path)
            pdf_text = extract_text_from_pdf(pdf_path)
            
            pdf_texts.append(pdf_text)
            pdf_names.append(pdf_name)
        
        # Prepare the user prompt with the review question and instructions
        table_instruction = ""
        if include_tables:
            table_instruction = " Please include important new generated tables in your review."
            
        user_prompt = f"Please generate a systematic review of the following {len(pdf_files)} papers: {', '.join(pdf_names)}.{table_instruction}\n\nReview Question: {review_question}"
        
        # Create the messages for the API call
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt + "\n\n" + "\n\n".join([f"Paper {i+1} - {pdf_names[i]}:\n{pdf_texts[i]}" for i in range(len(pdf_texts))])}
        ]
        
        # Call the API with temperature=0.7 and top_p=1
        response = openai.ChatCompletion.create(
            model="gpt-4.1",
            messages=messages,
            temperature=0.7,
            top_p=1,
            max_tokens=16384
        )
        
        # Get the AI response
        review_content = response["choices"][0]["message"]["content"]
        
        # Apply professional academic paper styling
        styled_html = f"""
        <!DOCTYPE html>
        <html lang="en">
        <head>
            <meta charset="UTF-8">
            <meta name="viewport" content="width=device-width, initial-scale=1.0">
            <title>Systematic Review</title>
            <style>
                /* Academic Paper Styling */
                body {{
                    font-family: 'Times New Roman', Times, serif;
                    line-height: 1.6;
                    color: #333;
                    margin: 0;
                    padding: 0;
                    background-color: #f9f9f9;
                }}
                .paper-container {{
                    max-width: 800px;
                    margin: 0 auto;
                    padding: 40px;
                    background-color: white;
                    box-shadow: 0 0 20px rgba(0, 0, 0, 0.1);
                }}
                header {{
                    text-align: center;
                    margin-bottom: 30px;
                    border-bottom: 1px solid #ddd;
                    padding-bottom: 20px;
                }}
                h1 {{
                    font-size: 24px;
                    margin: 0 0 15px;
                    font-weight: bold;
                }}
                .author-info {{
                    font-size: 14px;
                    margin-bottom: 15px;
                }}
                .abstract {{
                    font-style: italic;
                    margin: 20px 0;
                    padding: 15px;
                    background-color: #f8f8f8;
                    border-left: 3px solid #ccc;
                }}
                h2 {{
                    font-size: 18px;
                    margin: 30px 0 15px;
                    border-bottom: 1px solid #eee;
                    padding-bottom: 5px;
                }}
                h3 {{
                    font-size: 16px;
                    margin: 25px 0 10px;
                }}
                p {{
                    margin: 0 0 15px;
                    text-align: justify;
                }}
                .section {{
                    margin-bottom: 30px;
                }}
                table {{
                    width: 100%;
                    border-collapse: collapse;
                    margin: 20px 0;
                    font-size: 14px;
                }}
                table, th, td {{
                    border: 1px solid #ddd;
                }}
                th, td {{
                    padding: 10px;
                    text-align: left;
                }}
                th {{
                    background-color: #f2f2f2;
                }}
                tr:nth-child(even) {{
                    background-color: #f9f9f9;
                }}
                .citation {{
                    font-size: 14px;
                    color: #555;
                }}
                .reference-list {{
                    margin-top: 40px;
                    border-top: 1px solid #ddd;
                    padding-top: 20px;
                }}
                .reference-list h2 {{
                    margin-top: 0;
                }}
                .reference-item {{
                    margin-bottom: 10px;
                    padding-left: 25px;
                    text-indent: -25px;
                }}
                ul, ol {{
                    margin: 15px 0;
                    padding-left: 25px;
                }}
                li {{
                    margin-bottom: 5px;
                }}
                .figure {{
                    margin: 25px 0;
                    text-align: center;
                }}
                .figure img {{
                    max-width: 100%;
                }}
                .figure-caption {{
                    font-size: 14px;
                    color: #666;
                    margin-top: 10px;
                }}
                .footnote {{
                    font-size: 12px;
                    color: #777;
                }}
                @media print {{
                    body {{
                        background-color: white;
                    }}
                    .paper-container {{
                        box-shadow: none;
                        padding: 0;
                    }}
                }}
            </style>
        </head>
        <body>
            <div class="paper-container">
                {review_content}
            </div>
        </body>
        </html>
        """
        
        return styled_html
    
    except Exception as e:
        return f"""
        <div style="color: red; padding: 20px; border: 1px solid red; border-radius: 5px; background-color: #ffecec;">
            <h3>Error Generating Systematic Review</h3>
            <p>{str(e)}</p>
        </div>
        """

# Function to save uploaded files
def save_uploaded_files(files):
    if not files:
        return []
    
    saved_paths = []
    for file in files:
        if file is not None:
            # Create a temporary file
            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
                tmp_file.write(file)
                saved_paths.append(tmp_file.name)
    
    return saved_paths

# Function to create a downloadable HTML file
def create_html_download_link(html_content):
    if not html_content or "<div style=\"color: red; padding: 20px;" in html_content or "Please upload at least one PDF file" in html_content:
        return None
    
    # Create timestamp for the filename
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"systematic_review_{timestamp}.html"
    
    # Encode the HTML content for download
    b64_html = base64.b64encode(html_content.encode()).decode()
    download_link = f'<a href="data:text/html;base64,{b64_html}" download="{filename}" class="download-button">Download HTML</a>'
    
    return download_link

# Add CSS styling for the Gradio interface
custom_css = """
<style>
    /* Main UI */
    .gradio-container {
        font-family: 'Arial', sans-serif;
        background-color: #f9f9f9;
    }
    
    /* Header */
    h1 {
        font-size: 28px;
        color: #333;
        margin-bottom: 20px;
        text-align: center;
        padding-bottom: 10px;
        border-bottom: 2px solid #4a00e0;
    }
    
    /* Primary Button */
    #generate_button {
        background: linear-gradient(135deg, #4a00e0 0%, #8e2de2 100%);
        color: white;
        font-weight: bold;
        padding: 10px 20px;
        border-radius: 5px;
        transition: all 0.3s ease;
    }
    #generate_button:hover {
        background: linear-gradient(135deg, #5b10f1 0%, #9f3ef3 100%);
        transform: translateY(-2px);
        box-shadow: 0 4px 8px rgba(0,0,0,0.1);
    }
    
    /* API Key Button */
    #api_key_button {
        background: linear-gradient(135deg, #68d391 0%, #48bb78 100%);
        color: white;
        font-weight: bold;
        margin-top: 27px;
        padding: 10px 20px;
        border-radius: 5px;
        transition: all 0.3s ease;
    }
    #api_key_button:hover {
        background: linear-gradient(135deg, #38a169 0%, #68d391 100%);
        transform: translateY(-2px);
        box-shadow: 0 4px 8px rgba(0,0,0,0.1);
    }
    
    /* Form Elements */
    .input-container {
        background-color: white;
        padding: 20px;
        border-radius: 8px;
        box-shadow: 0 2px 10px rgba(0,0,0,0.05);
        margin-bottom: 20px;
    }
    
    /* Labels */
    label {
        font-weight: 600;
        color: #555;
        margin-bottom: 8px;
    }
    
    /* Instructions Accordion */
    .accordion {
        background-color: white;
        border: 1px solid #e0e0e0;
        border-radius: 8px;
        margin-bottom: 20px;
    }
    
    /* Output Container */
    .output-container {
        background-color: white;
        padding: 15px;
        border-radius: 8px;
        box-shadow: 0 2px 10px rgba(0,0,0,0.05);
    }
    
    /* File Upload Area */
    .file-upload {
        border: 2px dashed #ccc;
        border-radius: 5px;
        padding: 20px;
        text-align: center;
        margin-bottom: 20px;
    }
    
    /* Download Button */
    .download-button {
        display: inline-block;
        background: linear-gradient(135deg, #4a00e0 0%, #8e2de2 100%);
        color: white;
        font-weight: bold;
        padding: 8px 16px;
        border-radius: 4px;
        text-decoration: none;
        margin-bottom: 10px;
        transition: all 0.3s ease;
    }
    .download-button:hover {
        background: linear-gradient(135deg, #5b10f1 0%, #9f3ef3 100%);
        transform: translateY(-2px);
        box-shadow: 0 4px 8px rgba(0,0,0,0.1);
    }
    
    /* Download Container */
    #download-container {
        display: flex;
        justify-content: center;
        margin: 20px 0;
        padding: 15px;
        background-color: #f5f5f5;
        border-radius: 8px;
    }
    
    /* Responsive adjustments */
    @media screen and (max-width: 768px) {
        .gradio-container {
            padding: 10px;
        }
    }
</style>
"""

# Gradio UI Layout with improved styling
with gr.Blocks(css=custom_css) as demo:
    gr.Markdown("# Systematic Review Generator for Research Papers")
    
    with gr.Accordion("How to Use This App", open=False):
        gr.Markdown(""" 
        ### Getting Started:
        1. Enter your OpenAI API key in the field below and click "Set API Key"
        2. Upload multiple PDF research papers (2 or more recommended)
        3. Enter your review question or topic
        4. Check the "Include Tables" option if you want the review to include comparison tables
        5. Click "Generate Systematic Review" to start the process
        6. After generation, you can download the review as HTML
        
        ### Tips for Best Results:
        - Upload papers that are related to the same research topic or field
        - Be specific in your review question to get more focused results
        - The generated review will follow a systematic structure including research field identification, data extraction, analysis, and conclusions
        - The more papers you upload, the more comprehensive the review will be
        - The review will be formatted as a professional academic paper with proper sections and citations
        """)
    
    # API Key Input in a styled container
    with gr.Row(elem_classes="input-container"):
        with gr.Column(scale=3):
            api_key_input = gr.Textbox(label="Enter OpenAI API Key", type="password", placeholder="sk-...")
        with gr.Column(scale=1):
            api_key_button = gr.Button("Set API Key", elem_id="api_key_button")
        api_key_output = gr.Textbox(label="API Key Status", interactive=False)
    
    # PDF Upload and Review Settings
    with gr.Row(elem_classes="input-container"):
        with gr.Column():
            gr.Markdown("### Upload Research Papers")
            pdf_files = gr.File(label="Upload PDF Research Papers", file_count="multiple", type="binary", elem_classes="file-upload")
            review_question = gr.Textbox(
                label="Review Question or Topic", 
                value="Please generate a systematic review of the following papers.",
                placeholder="e.g., What are the effects of mindfulness meditation on stress reduction?"
            )
            include_tables = gr.Checkbox(label="Include Comparison Tables", value=True)
            generate_button = gr.Button("Generate Systematic Review", elem_id="generate_button", size="large")
    
    # Download link container
    download_html_output = gr.HTML(label="Download Options")
    
    # Output with improved styling
    with gr.Row(elem_classes="output-container"):
        review_output = gr.HTML(label="Systematic Review")
    
    # Button actions
    api_key_button.click(set_api_key, inputs=[api_key_input], outputs=[api_key_output])
    
    # Generate systematic review
    def process_files_and_generate_review(files, question, include_tables):
        if not files:
            return ("""
            <div style="padding: 20px; border: 1px solid #e0e0e0; border-radius: 5px; background-color: #f9f9f9;">
                <h3 style="color: #666;">Please upload at least one PDF file.</h3>
                <p>To generate a systematic review, upload one or more research papers in PDF format.</p>
            </div>
            """, "")
        
        # Save uploaded files
        saved_paths = save_uploaded_files(files)
        
        # Show loading message
        loading_message = """
        <div style="padding: 20px; text-align: center;">
            <h3>Generating Systematic Review...</h3>
            <p>This may take a few minutes depending on the number and size of papers.</p>
            <div style="width: 100%; height: 4px; background-color: #f0f0f0; margin: 20px 0; border-radius: 2px; overflow: hidden;">
                <div style="width: 30%; height: 100%; background: linear-gradient(90deg, #4a00e0, #8e2de2); animation: progress 2s infinite linear;"></div>
            </div>
            <style>
                @keyframes progress {
                    0% { margin-left: -30%; }
                    100% { margin-left: 100%; }
                }
            </style>
        </div>
        """
        
        yield loading_message, ""
        
        # Generate review
        review = generate_systematic_review(saved_paths, question, include_tables)
        
        # Create HTML download link
        html_link = create_html_download_link(review)
        
        # Create download link HTML
        download_link = f"""
        <div id="download-container">
            <div>
                <h3>Download Option:</h3>
                {html_link or ""}
            </div>
        </div>
        """
        
        # Clean up temporary files
        for path in saved_paths:
            try:
                os.remove(path)
            except:
                pass
        
        yield review, download_link
    
    generate_button.click(
        process_files_and_generate_review, 
        inputs=[pdf_files, review_question, include_tables], 
        outputs=[review_output, download_html_output]
    )

# Launch the app
if __name__ == "__main__":
    demo.launch(share=True)