import json
import os
import tempfile
from typing import List, Dict, Any

import fitz  # PyMuPDF
import gradio as gr
from openai import OpenAI

# Load API key from environment variable
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")


def load_pitfalls() -> List[Dict[str, Any]]:
    """Load pitfalls from the JSON file."""
    try:
        with open("pitfalls.json", "r") as f:
            data = json.load(f)
            return data.get("pitfalls", [])
    except FileNotFoundError:
        gr.Warning("pitfalls.json file not found!")
        return []
    except json.JSONDecodeError:
        gr.Warning("Invalid JSON in pitfalls.json file!")
        return []


def extract_text_from_pdf(pdf_file) -> str:
    """Extract text content from a PDF file."""
    try:
        pdf_path = pdf_file.name
        doc = fitz.open(pdf_path)
        text_content = ""

        for page_num in range(len(doc)):
            page = doc[page_num]
            text_content += f"\n--- Page {page_num + 1} ---\n"
            text_content += page.get_text()

        doc.close()
        return text_content
    except Exception as e:
        raise gr.Error(f"Error extracting text from {pdf_file.name}: {str(e)}")


def format_paper_text(paper_text: str) -> Dict[str, Any]:
    """First stage: Format the paper text to make it more readable and suitable for analysis."""

    # Check if API key is available
    if not OPENROUTER_API_KEY:
        return {
            "formatted_text": None,
            "success": False,
            "error": "OpenRouter API key not found. Please set the OPENROUTER_API_KEY environment variable.",
        }

    # Initialize OpenAI client with OpenRouter
    client = OpenAI(
        base_url="https://openrouter.ai/api/v1",
        api_key=OPENROUTER_API_KEY,
    )

    format_prompt = f"""You are an expert academic text processor. Your task is to clean and format the following research paper text to make it more readable and suitable for detailed analysis.

Please:
1. Remove excessive whitespace and formatting artifacts
2. Organize the text into clear sections (Abstract, Introduction, Methods, Results, Discussion, Conclusion, References, Appendix)
3. Preserve all important content including figures, tables, and equations
4. Make the text flow better while maintaining academic integrity
5. Ensure all evaluation-related content is clearly identifiable
6. Keep the text under 8000 characters while preserving key information

Original paper text:
{paper_text}

Please provide the cleaned and formatted text:"""

    try:
        completion = client.chat.completions.create(
            extra_headers={
                "HTTP-Referer": "https://github.com/paper-eval-checker",
                "X-Title": "Paper Evaluation Pitfall Checker",
            },
            model="x-ai/grok-4-fast:free",
            messages=[{"role": "user", "content": format_prompt}],
            temperature=0.1,  # Very low temperature for consistent formatting
            max_tokens=3000,
        )

        return {
            "formatted_text": completion.choices[0].message.content,
            "success": True,
            "error": None,
        }
    except Exception as e:
        return {"formatted_text": None, "success": False, "error": str(e)}


def analyze_paper_for_pitfalls(
    formatted_text: str, pitfalls: List[Dict[str, Any]]
) -> Dict[str, Any]:
    """Second stage: Use OpenRouter API with Grok model to analyze the formatted paper for potential pitfalls."""

    # Check if API key is available
    if not OPENROUTER_API_KEY:
        return {
            "analysis": None,
            "success": False,
            "error": "OpenRouter API key not found. Please set the OPENROUTER_API_KEY environment variable.",
        }

    # Initialize OpenAI client with OpenRouter
    client = OpenAI(
        base_url="https://openrouter.ai/api/v1",
        api_key=OPENROUTER_API_KEY,
    )

    # Create the prompt for pitfall analysis
    pitfalls_description = "\n\n".join(
        [
            f"**{pitfall['name']}** {pitfall['emoji']}\n"
            f"Category: {pitfall['category']}\n"
            f"Description: {pitfall['description']}\n"
            f"Subjective/Objective: {pitfall['subjective_objective']}\n"
            f"Actors Most Affected: {', '.join(pitfall['actors_most_affected'])}\n"
            f"Evaluation Use: {pitfall['evaluation_use']}\n"
            f"Modalities: {', '.join(pitfall['modalities'])}"
            for pitfall in pitfalls
        ]
    )

    analysis_prompt = f"""You are an expert research paper reviewer specializing in identifying evaluation pitfalls in academic papers. 

Your task is to analyze the provided formatted research paper text and identify any potential pitfalls from the following list:

{pitfalls_description}

Please analyze the paper carefully and provide:
1. A list of potential pitfalls found (if any)
2. For each pitfall found, provide:
   - The pitfall name (and emoji)
   - Specific evidence from the paper that suggests this pitfall
   - The section/page where this evidence appears
   - Suggestions for improvement
3. Be concise, and use markdown formatting.
4. If you find evidence of a pitfall, make sure to look at ALL of the paper to see if it is mitigated elsewhere -- make sure to check the appendix of the paper as well.

The output format:

# Overall
![](https://img.shields.io/severity/high-red) (for low use green, for medium use yellow, for high use red)
![](https://img.shields.io/evaluation/objective-blue) (either write 'subjective', 'objective', or include two images in case both are present in the paper)
[One sentence summary of evaluation use]

# Pitfall 

## Evidence
"specific evidence from the paper"

If no pitfalls are found, please state that clearly.

Formatted paper text to analyze:
{formatted_text}

Please provide your analysis in a structured format."""

    try:
        completion = client.chat.completions.create(
            extra_headers={
                "HTTP-Referer": "https://github.com/paper-eval-checker",
                "X-Title": "Paper Evaluation Pitfall Checker",
            },
            model="x-ai/grok-4-fast:free",
            messages=[{"role": "user", "content": analysis_prompt}],
            temperature=0.3,  # Lower temperature for more consistent analysis
            max_tokens=2000,
        )

        return {
            "analysis": completion.choices[0].message.content,
            "success": True,
            "error": None,
        }
    except Exception as e:
        return {"analysis": None, "success": False, "error": str(e)}


def process_paper(pdf_file, progress=gr.Progress()):
    """Main function to process a research paper for pitfall detection using two-stage approach."""

    if not pdf_file:
        return gr.Markdown(
            "⚠️ No PDF file uploaded. Please upload a research paper PDF."
        )

    if not OPENROUTER_API_KEY:
        return gr.Markdown(
            "⚠️ OpenRouter API key not found. Please set the OPENROUTER_API_KEY environment variable."
        )

    try:
        # Step 1: Load pitfalls
        progress(0.1, desc="Loading pitfalls definitions...")
        pitfalls = load_pitfalls()

        if not pitfalls:
            return gr.Markdown(
                "⚠️ No pitfalls definitions found. Please check pitfalls.json file."
            )

        # Step 2: Extract text from PDF
        progress(0.2, desc="Extracting text from PDF...")
        paper_text = extract_text_from_pdf(pdf_file)

        if not paper_text.strip():
            return gr.Markdown(
                "⚠️ No text content found in the PDF. Please check if the PDF contains readable text."
            )

        # Step 3: Format paper text (First AI call)
        progress(0.3, desc="Formatting paper text for analysis...")
        format_result = format_paper_text(paper_text)

        if not format_result["success"]:
            return gr.Markdown(
                f"❌ Error during text formatting: {format_result['error']}"
            )

        # Step 4: Analyze for pitfalls (Second AI call)
        progress(0.7, desc="Analyzing paper for potential pitfalls...")
        analysis_result = analyze_paper_for_pitfalls(
            format_result["formatted_text"], pitfalls
        )

        if not analysis_result["success"]:
            return gr.Markdown(f"❌ Error during analysis: {analysis_result['error']}")

        # Step 5: Format final results
        progress(0.9, desc="Preparing final report...")
        analysis_text = analysis_result["analysis"]

        # Create a formatted markdown report
        report = f"""# Research Paper Pitfall Analysis Report

## Analysis Results

{analysis_text}

---
*Analysis completed using OpenRouter API with Grok model (two-stage processing)*
"""

        progress(1.0, desc="Analysis complete!")
        return gr.Markdown(report)

    except Exception as e:
        return gr.Markdown(f"❌ An error occurred: {str(e)}")


# Define the Gradio interface
with gr.Blocks(title="Research Paper Pitfall Checker") as demo:
    gr.HTML(
        """<h1 style='text-align: center;'>🔍 Research Paper Pitfall Checker</h1>
        <center><i>Identify potential evaluation pitfalls in academic research papers</i></center>"""
    )

    gr.HTML(
        """
        <div style="max-width: 800px; margin: 0 auto; padding: 20px;">
            <h3>How it works:</h3>
            <ol>
                <li><strong>Upload a PDF</strong> of your research paper</li>
                <li><strong>Click "Analyze Paper"</strong> to scan for potential pitfalls</li>
                <li><strong>Review the analysis</strong> to identify areas for improvement</li>
                </ol>
            
            <h3>Supported Pitfalls:</h3>
            <ul>
                <li>🔒 The Lock-In Effect</li>
                <li>🍎🍊 Apples-to-Oranges Comparisons</li>
                <li>💧 Contamination Leak</li>
                <li>🤖❓ Unvalidated Automation</li>
                <li>🧐 Vague Scales</li>
            </ul>
        </div>
        """
    )

    with gr.Row():
        with gr.Column(scale=3):
            pdf_file = gr.File(label="Upload Research Paper (PDF)", file_types=[".pdf"])

        with gr.Column(scale=1):
            analyze_button = gr.Button(
                "🔍 Analyze Paper for Pitfalls",
                variant="primary",
                size="lg",
                elem_id="analyze-btn",
            )

    with gr.Row():
        results = gr.Markdown(
            value="Upload a PDF to get started with pitfall analysis.",
            elem_id="results",
        )

    # Add loading animation CSS
    demo.css = """
    #analyze-btn {
        background: linear-gradient(45deg, #ff6b6b, #4ecdc4, #45b7d1, #96ceb4, #feca57);
        background-size: 400% 400%;
        animation: gradient 3s ease infinite;
        border: none;
        color: white;
        font-weight: bold;
    }
    
    @keyframes gradient {
        0% { background-position: 0% 50%; }
        50% { background-position: 100% 50%; }
        100% { background-position: 0% 50%; }
    }
    
    #results {
        min-height: 200px;
        padding: 20px;
        border: 1px solid #e0e0e0;
        border-radius: 8px;
        background-color: #ffffff;
        color: #333333;
    }
    
    #results .markdown {
        color: #333333 !important;
    }
    
    #results h1, #results h2, #results h3, #results h4, #results h5, #results h6 {
        color: #2c3e50 !important;
    }
    
    #results p, #results li, #results div {
        color: #333333 !important;
    }
    
    #results code {
        background-color: #f8f9fa;
        color: #e83e8c;
        padding: 2px 4px;
        border-radius: 3px;
    }
    
    #results pre {
        background-color: #f8f9fa;
        color: #333333;
        padding: 10px;
        border-radius: 5px;
        border: 1px solid #e9ecef;
    }
    """

    # Connect the button to the processing function
    analyze_button.click(
        fn=process_paper,
        inputs=[pdf_file],
        outputs=[results],
    )

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