import gradio as gr
import pandas as pd
import os
import json
import logging
import time
import uuid
from pathlib import Path
from citation_validator import CitationValidator, CitationValidatorError, APIKeyError

# Configure logging first
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# File upload validation constants
MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB
ALLOWED_EXTENSIONS = {'.pdf'}
MAX_TEXT_SIZE = 100 * 1024  # 100KB for text input


try:
    from gradio_pdf import PDF
    PDF_AVAILABLE = True
    logger.info("PDF component loaded successfully")
except ImportError:
    PDF_AVAILABLE = False
    logger.warning("gradio_pdf not available, PDF preview will be disabled")

logger.info("Reference Checker app starting...")

def update_pdf_preview(file_obj):
    """Show PDF preview immediately after upload"""
    if file_obj is not None:
        return file_obj
    return None

def check_references(file_obj, text_input, input_type, ai_verification, basic_formatting, api_key_input):
    """
    Check references from either PDF or text input.

    Args:
        file_obj: Uploaded PDF file object
        text_input: Text string input
        input_type: Either "PDF" or "Text"
        ai_verification: Whether to use AI verification for failed citations
        basic_formatting: Whether to use basic rule-based formatting instead of AI
        api_key_input: Optional user-provided Google API key

    Returns:
        Tuple of (DataFrame, api_results, grounding_metadata)
    """
    try:
        logger.info(f"Starting reference check - Input type: {input_type}, AI verification: {ai_verification}")

        # API key from input takes precedence over environment variable
        try:
            api_key = api_key_input.strip() if api_key_input and api_key_input.strip() else None
            checker = CitationValidator(api_key=api_key)
        except APIKeyError as e:
            logger.error(f"API key error: {e}")
            error_df = pd.DataFrame({
                "Error": ["API Key Error"],
                "Details": [str(e)],
                "Solution": ["Please either enter your Google API key above or set GOOGLE_API_KEY in .env file"]
            })
            return error_df, None, None

        # Process based on input type
        if input_type == "PDF":
            if file_obj is None:
                logger.warning("PDF selected but no file uploaded")
                error_df = pd.DataFrame({"Error": ["Please upload a PDF file"]})
                return error_df, None, None

            if not hasattr(file_obj, 'name') or not file_obj.name:
                logger.error("Invalid file object")
                error_df = pd.DataFrame({"Error": ["Invalid file upload. Please try again."]})
                return error_df, None, None

            # Validate file size
            if hasattr(file_obj, 'size') and file_obj.size > MAX_FILE_SIZE:
                logger.warning(f"File too large: {file_obj.size / 1024 / 1024:.1f}MB")
                error_df = pd.DataFrame({
                    "Error": [f"File too large: {file_obj.size / 1024 / 1024:.1f}MB (max 10MB)"]
                })
                return error_df, None, None

            # Validate file extension
            file_path = Path(file_obj.name)
            if file_path.suffix.lower() not in ALLOWED_EXTENSIONS:
                logger.warning(f"Invalid file type: {file_path.suffix}")
                error_df = pd.DataFrame({
                    "Error": [f"Only PDF files allowed (received: {file_path.suffix})"]
                })
                return error_df, None, None

            # Validate against path traversal
            try:
                file_path = file_path.resolve()
                if not file_path.exists():
                    logger.error(f"File does not exist: {file_path}")
                    error_df = pd.DataFrame({"Error": ["File not found. Please upload again."]})
                    return error_df, None, None
            except (OSError, RuntimeError) as e:
                logger.error(f"Invalid file path: {e}")
                error_df = pd.DataFrame({"Error": ["Invalid file path. Please upload again."]})
                return error_df, None, None

            logger.info(f"Processing PDF: {file_obj.name}")
            results = checker.check_references(str(file_path), ai_verify_failed=ai_verification, use_basic_formatting=basic_formatting)

        elif input_type == "Text":
            if not text_input or not text_input.strip():
                logger.warning("Text selected but no content provided")
                error_df = pd.DataFrame({"Error": ["Please enter some text with references"]})
                return error_df, None, None

            # Validate text size
            text_size = len(text_input.encode('utf-8'))
            if text_size > MAX_TEXT_SIZE:
                logger.warning(f"Text too large: {text_size / 1024:.1f}KB")
                error_df = pd.DataFrame({
                    "Error": [f"Text too large: {text_size / 1024:.1f}KB (max {MAX_TEXT_SIZE / 1024:.0f}KB)"]
                })
                return error_df, None, None

            logger.info(f"Processing text input ({len(text_input)} characters)")
            results = checker.check_references_from_text(text_input, ai_verify_failed=ai_verification, use_basic_formatting=basic_formatting)

        else:
            logger.error(f"Invalid input type: {input_type}")
            raise ValueError(f"Invalid input type: {input_type}")

        # Convert main results to DataFrame
        if isinstance(results['results'], list):
            df = pd.DataFrame(results['results'])
        else:
            df = pd.DataFrame([results['results']])

        # Format URLs as clickable links
        if 'url' in df.columns:
            df['url'] = df['url'].apply(
                lambda x: f'<a href={x} target="_blank">{x}</a>' if pd.notna(x) and x else ""
            )
        if 'link' in df.columns:
            df['link'] = df['link'].apply(
                lambda x: f'<a href={x} target="_blank">{x}</a>' if pd.notna(x) and x else ""
            )

        logger.info(f"Successfully processed {len(df)} references")
        return df, results.get('api_results', None), results.get('grounding_metadata', None)

    except FileNotFoundError as e:
        logger.error(f"File not found: {e}")
        error_df = pd.DataFrame({"Error": [f"File not found: {str(e)}"]})
        return error_df, None, None

    except ValueError as e:
        logger.error(f"Validation error: {e}")
        error_df = pd.DataFrame({"Error": [str(e)]})
        return error_df, None, None

    except CitationValidatorError as e:
        logger.error(f"Reference checker error: {e}")
        error_df = pd.DataFrame({"Error": [f"Reference checking failed: {str(e)}"]})
        return error_df, None, None

    except Exception as e:
        logger.exception(f"Unexpected error processing {input_type}: {e}")
        error_df = pd.DataFrame({
            "Error": [f"Unexpected error processing {input_type.lower()}"],
            "Details": [str(e)],
            "Action": ["Please try again or contact support if the problem persists"]
        })
        return error_df, None, None

def show_metadata(evt: gr.SelectData, results_df, crossref_data, metadata):
    """Display metadata for the selected result"""
    try:
        if metadata is None or crossref_data is None:
            return {"message": "No metadata available for this result"}
            
        # Get the selected row index
        row_idx = evt.index[0]
        
        # Create a combined metadata object
        result_metadata = {
            "api_data": crossref_data[row_idx] if isinstance(crossref_data, list) else crossref_data,
            "grounding_metadata": metadata
        }
        
        return result_metadata
    except Exception as e:
        return {"error": str(e)}

def download_results(df, format_type):
    """
    Download results as CSV or Excel file.

    Args:
        df: DataFrame containing results
        format_type: Either "CSV" or "Excel"

    Returns:
        Path to the generated file or None if failed
    """
    try:
        if df is None or df.empty:
            logger.warning("Attempted to download empty results")
            return None

        logger.info(f"Generating {format_type} download file")

        # Generate unique filename
        unique_id = f"{int(time.time())}_{uuid.uuid4().hex[:8]}"

        if format_type.lower() == "csv":
            output_path = f"citation_validation_results_{unique_id}.csv"
            df.to_csv(output_path, index=False)
            logger.info(f"CSV file created: {output_path}")
        else:
            # Excel format
            output_path = f"citation_validation_results_{unique_id}.xlsx"
            try:

                df.to_excel(output_path, index=False, engine='openpyxl')
                logger.info(f"Excel file created with openpyxl: {output_path}")
            except ImportError:
                logger.warning("openpyxl not available, trying xlsxwriter")
                try:
                    # Try with xlsxwriter as fallback
                    df.to_excel(output_path, index=False, engine='xlsxwriter')
                    logger.info(f"Excel file created with xlsxwriter: {output_path}")
                except ImportError:
                    # If all Excel engines fail, fall back to CSV
                    logger.warning("No Excel engines available, falling back to CSV")
                    output_path = f"citation_validation_results_{unique_id}.csv"
                    df.to_csv(output_path, index=False)
                    logger.info(f"Fallback CSV file created: {output_path}")
            except Exception as e:
                logger.error(f"Excel export failed: {e}, falling back to CSV")
                output_path = f"citation_validation_results_{unique_id}.csv"
                df.to_csv(output_path, index=False)

        # Make sure the file exists before returning
        if os.path.exists(output_path):
            logger.info(f"Download file ready: {output_path}")
            return output_path

        logger.error("Download file was not created")
        return None

    except Exception as e:
        logger.exception(f"Download error: {e}")
        return None

css = """
:root {
    --primary-color: #6366F1;
    --secondary-color: #4F46E5;
    --background-color: #F9FAFB;
    --text-color: #1F2937;
    --error-color: #EF4444;
    --success-color: #10B981;
    --card-bg: #FFFFFF;
    --border-radius: 12px;
    --shadow: 0 10px 15px -3px rgba(0,0,0,0.1), 0 4px 6px -2px rgba(0,0,0,0.05);
}

.gradio-container {
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
}

.app-header {
    text-align: center;
    margin-bottom: 2rem;
    padding: 1.5rem 0;
}

.app-title {
    color: var(--primary-color);
    font-size: 2.8rem;
    font-weight: 800;
    margin-bottom: 0.75rem;
    background: linear-gradient(90deg, var(--primary-color), var(--secondary-color));
    -webkit-background-clip: text;
}

.app-description {
    font-size: 1.2rem;
    opacity: 0.85;
    max-width: 800px;
    margin: 0 auto;
    line-height: 1.6;
}

.input-panel, .pdf-panel, .results-panel {
    border-radius: var(--border-radius);
    padding: 1.5rem;
    margin-bottom: 1.5rem;
}

.panel-title {
    font-size: 1.5rem;
    font-weight: 600;
    margin-bottom: 1.25rem;
    color: var(--secondary-color);
}

/* Button styling */
.custom-button {
    background: linear-gradient(90deg, var(--primary-color), var(--secondary-color)) !important;
    color: white !important;
    font-weight: 600 !important;
    border-radius: 8px !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 6px -1px rgba(99, 102, 241, 0.4) !important;
}

.custom-button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 8px -1px rgba(99, 102, 241, 0.6) !important;
}

/* Form elements */
input, select, textarea {
    border-radius: 8px !important;
    border: 2px solid rgba(99, 102, 241, 0.2) !important;
    padding: 0.75rem 1rem !important;
    transition: all 0.2s !important;
}

input:focus, select:focus, textarea:focus {
    border-color: var(--primary-color) !important;
    box-shadow: 0 0 0 3px rgba(99, 102, 241, 0.2) !important;
}

/* Ensure input fields inherit background from theme */
.gradio-container input[type="text"],
.gradio-container input[type="password"],
.gradio-container textarea {
    background: var(--input-background-fill) !important;
}

/* Table styling */
table.dataframe {
    border-collapse: separate !important;
    border-spacing: 0 !important;
    width: 100% !important;
    border-radius: 8px !important;
    overflow: hidden !important;
}

table.dataframe th {
    padding: 12px !important;
    font-weight: 600 !important;
    text-align: left !important;
}

table.dataframe td {
    padding: 12px !important;
    border-bottom: 1px solid rgba(99, 102, 241, 0.1) !important;
}

.metadata-panel {
    margin-top: 1.5rem;
    border: 1px solid rgba(99, 102, 241, 0.2);
    border-radius: var(--border-radius);
}

.metadata-help {
    color: var(--text-color);
    opacity: 0.7;
    font-style: italic;
    margin-bottom: 1rem;
}

.json-viewer {
    max-height: 500px;
    overflow-y: auto;
    font-family: 'Fira Code', monospace;
    font-size: 14px;
    line-height: 1.6;
}

/* PDF display */
.pdf-container {
    height: 800px;
    border-radius: var(--border-radius);
    overflow: hidden;
    box-shadow: var(--shadow);
}

/* Download section */
.download-section {
    margin-top: 1.5rem;
    display: flex;
    justify-content: flex-end;
    gap: 0.5rem;
}

/* Links that open in new tab */
a {
    color: var(--primary-color) !important;
    text-decoration: none !important;
    font-weight: 500 !important;
}

a:hover {
    text-decoration: underline !important;
    color: var(--secondary-color) !important;
}

/* JSON formatting */
pre.json {
    background-color: #f8fafc;
    border-radius: 8px;
    padding: 1rem;
    overflow: auto;
    max-height: 500px;
    font-family: 'Fira Code', monospace;
    font-size: 14px;
    line-height: 1.6;
    border: 1px solid rgba(99, 102, 241, 0.2);
}

.input-description {
    margin-bottom: 1.5rem;
    padding: 1rem;
    background-color: rgba(99, 102, 241, 0.05);
    border-left: 4px solid var(--primary-color);
    border-radius: 6px;
    font-size: 0.95rem;
    line-height: 1.6;
}

/* Custom radio button styling */
.custom-radio > div {
    display: flex;
    gap: 1rem;
}

.custom-radio label {
    padding: 0.75rem 1.25rem;
    border: 2px solid rgba(99, 102, 241, 0.3);
    border-radius: 8px;
    transition: all 0.2s;
    cursor: pointer;
    font-weight: 500;
}

.custom-radio input[type="radio"]:checked + label {
    background-color: var(--primary-color);
    color: white;
    border-color: var(--primary-color);
}

.custom-radio input[type="radio"] {
    display: none;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Row(elem_classes=["app-header"]):
        gr.HTML("""<h1 class="app-title">Citation Validator</h1>
                 <p class="app-description">Validate academic citations against multiple scholarly databases and AI verification.</p>""")

    # State variables to store results
    crossref_results_state = gr.State(None)
    metadata_state = gr.State(None)
    results_df_state = gr.State(None)

    # Top section: Input (1/3 width) and PDF preview (2/3 width)
    with gr.Row():
        # Input section (1/3 width)
        with gr.Column(scale=1, elem_classes=["input-panel"]):
            gr.Markdown("### Input Options", elem_classes=["panel-title"])

            # Add description and disclaimer here
            gr.Markdown("""
            **Citation Validator** verifies citations in academic documents against
            CrossRef, arXiv, OpenAlex, and AI-powered validation.

            **How it works:** Upload PDF or paste text → AI extracts citations → Validates against
            multiple databases → Provides verification results with source links.

            ⚠️ **Rate Limiting Notice**: During high traffic, you may encounter rate limiting. **For best performance and no rate limits, run this tool locally**.

            **Disclaimer:** Automated validation may not catch all errors. Always manually verify
            results for critical submissions. Documents are not stored after processing.
            """, elem_classes=["input-description"])

            # API Key input (optional - for users who want to use their own key)
            api_key_input = gr.Textbox(
                label="Google API Key (Required)",
                placeholder="Enter your Google API key here (local users can set .env variable)",
                type="password",
                info="Your API key is never stored. Get a free key at https://aistudio.google.com/apikey (free tier has limits for 20 documents/day)"
            )

            # Then modify the radio button with custom styling
            input_type = gr.Radio(
                ["PDF", "Text"],
                label="Choose Input Type",
                value="PDF",
                elem_classes=["custom-radio"]
            )

            # PDF Upload Section
            with gr.Group(visible=True) as pdf_group:
                file_upload = gr.File(
                    label="Upload PDF Document",
                    file_types=[".pdf"],
                    elem_classes=["file-upload"]
                )

            # Text Input Section
            with gr.Group(visible=False) as text_group:
                text_input = gr.Textbox(
                    label="Paste Text with References",
                    placeholder="Enter your text here...",
                    lines=10
                )

            basic_formatting = gr.Checkbox(
                label="Use Basic Formatting (No AI, faster, saves API calls)",
                value=True,
                info="Use rule-based formatting instead of AI for analyzing API results"
            )
                
            ai_verification = gr.Checkbox(
                label="AI Verification for Failed Citations (Uses a lot more API calls)",
                value=False,
                info="Verify failed citations one-by-one with AI + Google Search grounding"
            )



            check_button = gr.Button(
                "Validate References",
                variant="primary",
                elem_classes=["custom-button"]
            )

        # PDF Preview (2/3 width)
        with gr.Column(scale=2, elem_classes=["pdf-panel"]):
            gr.Markdown("### PDF Preview", elem_classes=["panel-title"])
            if PDF_AVAILABLE:
                pdf_preview = PDF(
                    label="",
                    visible=True,
                    height=800,
                    starting_page=1,
                    elem_classes=["pdf-container"]
                )
            else:
                pdf_preview = gr.Markdown("PDF preview not available (gradio-pdf not installed)")

    # Results section
    with gr.Row(elem_classes=["results-panel"]):
        with gr.Column():
            gr.Markdown("### Results", elem_classes=["panel-title"])

            results_df = gr.Dataframe(
                headers=["Title", "Exists", "Link","URL",  "Explanation", "First Author", "Year", "Journal", "Issues"],
                elem_classes=["results-table"],
                datatype=["str", "str", "html", "html", "str", "str", "str", "str", "str"],
            )

            with gr.Accordion("Result Details", open=True, elem_classes=["metadata-panel"]):
                gr.Markdown("Click on a result to view detailed metadata", elem_classes=["metadata-help"])
                metadata_json = gr.JSON(label="", elem_classes=["json-viewer"])

            # Download section
            with gr.Row(elem_classes=["download-section"]):
                download_format = gr.Radio(
                    ["CSV", "Excel"],
                    label="Download Format",
                    value="CSV"
                )
                download_button = gr.Button(
                    "Download Results",
                    variant="secondary",
                    elem_classes=["custom-button"]
                )
                download_output = gr.File(
                    label="Download"
                )



    # Toggle between PDF and Text input
    input_type.change(
        lambda x: [gr.update(visible=x=="PDF"), gr.update(visible=x=="Text")],
        inputs=[input_type],
        outputs=[pdf_group, text_group]
    )

    # Update PDF preview when file is uploaded (only if PDF component is available)
    if PDF_AVAILABLE:
        file_upload.change(
            fn=update_pdf_preview,
            inputs=[file_upload],
            outputs=[pdf_preview]
        ).then(
            lambda x: gr.update(visible=x is not None),
            inputs=[pdf_preview],
            outputs=[pdf_preview]
        )

    # Check references and store results
    check_button.click(
        fn=check_references,
        inputs=[file_upload, text_input, input_type, ai_verification, basic_formatting, api_key_input],
        outputs=[results_df, crossref_results_state, metadata_state]
    ).then(
        lambda df: df,
        inputs=[results_df],
        outputs=[results_df_state]
    )

    results_df.select(
        fn=show_metadata,
        inputs=[results_df_state, crossref_results_state, metadata_state],
        outputs=[metadata_json]
    )

    # Download results
    download_button.click(
        fn=download_results,
        inputs=[results_df_state, download_format],
        outputs=[download_output]
    )

demo.launch(
    show_error=True,  # Show errors in UI for debugging
    inbrowser=True,  # Auto-open browser
    max_threads=5,  # Rate limiting: max 5 concurrent requests
    max_file_size="10mb"  # Enforce file size limit at Gradio level
)