#!/usr/bin/env python3
"""
Document Extraction Tool - Simplified Version
Convert PDF and DOCX files to Markdown format
Suitable for Hugging Face Spaces deployment
"""

import gradio as gr
import os
import logging
import tempfile
from typing import Tuple, Optional
from pathlib import Path

# Configure logging with better formatting
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

# Pre-load docling at module level to avoid runtime errors
try:
    from docling.document_converter import DocumentConverter

    DOCLING_AVAILABLE = True
    logger.info("Docling loaded successfully")
except ImportError as e:
    DOCLING_AVAILABLE = False
    logger.warning(f"Docling not available: {e}")

# Global extractor instance - created at startup, ready to extract documents
_extractor = None


class SimpleDocumentExtractor:
    """Simplified document extractor - created at startup, ready to extract documents"""

    def __init__(self):
        """Initialize extractor"""
        self.converter = None

        # Try to initialize docling
        if DOCLING_AVAILABLE:
            try:
                self.converter = DocumentConverter()
                logger.info("✅ Docling extractor ready")
            except Exception as e:
                logger.warning(f"⚠️ Docling initialization failed: {e}")
                self.converter = None

        # Initialize fallback libraries
        try:
            import PyPDF2
            import docx

            self.pdf_lib = PyPDF2
            self.docx_lib = docx
            logger.info("✅ Fallback extractors ready")
        except ImportError as e:
            logger.error(f"❌ No document processing libraries available: {e}")
            raise RuntimeError("No document processing libraries available")

    def extract(self, file_path: str) -> str:
        """Extract document content"""
        file_path = Path(file_path)

        if not file_path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")

        file_ext = file_path.suffix.lower()
        logger.info(f"🔄 Extracting document: {file_path.name}")

        # Prefer docling
        if self.converter:
            try:
                result = self.converter.convert(str(file_path))
                content = result.document.export_to_markdown()
                if content and content.strip():
                    logger.info(
                        f"✅ Docling extraction successful: {len(content)} characters"
                    )
                    return content
            except Exception as e:
                logger.warning(f"⚠️ Docling extraction failed: {e}, using fallback")

        # Fallback methods
        if file_ext == ".pdf":
            return self._extract_pdf(str(file_path))
        elif file_ext == ".docx":
            return self._extract_docx(str(file_path))
        else:
            raise ValueError(f"Unsupported file type: {file_ext}")

    def _extract_pdf(self, file_path: str) -> str:
        """Extract PDF content with enhanced error handling"""
        try:
            with open(file_path, "rb") as file:
                # Try with strict=False to handle corrupted PDFs better
                pdf_reader = self.pdf_lib.PdfReader(file, strict=False)
                text_content = []

                # Check if PDF has pages
                if len(pdf_reader.pages) == 0:
                    logger.warning("PDF file appears to be empty or corrupted")
                    return (
                        "# PDF Document\n\n⚠️ PDF file appears to be empty or corrupted"
                    )

                for page_num, page in enumerate(pdf_reader.pages):
                    try:
                        page_text = page.extract_text()
                        if page_text and page_text.strip():
                            text_content.append(
                                f"## Page {page_num + 1}\n\n{page_text.strip()}\n"
                            )
                    except Exception as e:
                        logger.warning(f"Failed to extract page {page_num + 1}: {e}")
                        continue

                if text_content:
                    result = "# PDF Document\n\n" + "\n".join(text_content)
                    logger.info(
                        f"✅ PDF extraction successful: {len(result)} characters from {len(text_content)} pages"
                    )
                    return result
                else:
                    logger.warning("No text content could be extracted from PDF")
                    return "# PDF Document\n\n⚠️ Unable to extract text content. This PDF may be:\n- Image-based (scanned document)\n- Protected or encrypted\n- Corrupted or malformed\n- Empty"

        except Exception as e:
            logger.error(f"❌ PDF extraction failed: {e}")
            error_msg = str(e).lower()

            if "startxref" in error_msg or "xref" in error_msg:
                return "# PDF Extraction Error\n\n❌ **PDF Structure Error**\n\nThe PDF file has structural issues (corrupted xref table). This often happens with:\n- Damaged or incomplete downloads\n- Corrupted files\n- Non-standard PDF generation\n\nTry:\n1. Re-downloading the file\n2. Opening in a PDF viewer to verify it works\n3. Converting to a new PDF format"
            elif "decrypt" in error_msg or "password" in error_msg:
                return "# PDF Extraction Error\n\n❌ **Password Protected**\n\nThis PDF is password-protected or encrypted.\n\nPlease:\n1. Unlock the PDF first\n2. Save an unprotected version\n3. Try again with the unlocked file"
            else:
                return f"# PDF Extraction Error\n\n❌ **Extraction Failed**\n\nError: {str(e)}\n\nThe file may be corrupted, damaged, or in an unsupported format."

    def _extract_docx(self, file_path: str) -> str:
        """Extract DOCX content"""
        try:
            doc = self.docx_lib.Document(file_path)
            text_content = []

            # Extract paragraphs
            for paragraph in doc.paragraphs:
                if paragraph.text and paragraph.text.strip():
                    text_content.append(paragraph.text.strip())

            # Extract tables
            for table in doc.tables:
                table_text = []
                for row in table.rows:
                    row_text = [
                        cell.text.strip() for cell in row.cells if cell.text.strip()
                    ]
                    if row_text:
                        table_text.append(" | ".join(row_text))
                if table_text:
                    text_content.append(
                        "\n### Table\n\n" + "\n".join(table_text) + "\n"
                    )

            if text_content:
                result = "# DOCX Document\n\n" + "\n\n".join(text_content)
                logger.info(f"✅ DOCX extraction successful: {len(result)} characters")
                return result
            else:
                return "# DOCX Document\n\n⚠️ Document is empty"
        except Exception as e:
            logger.error(f"❌ DOCX extraction failed: {e}")
            return f"# DOCX Extraction Error\n\n❌ Extraction failed: {str(e)}"


def get_extractor() -> SimpleDocumentExtractor:
    """Get document extractor instance (singleton pattern)"""
    global _extractor
    if _extractor is None:
        _extractor = SimpleDocumentExtractor()
    return _extractor


def extract_document(file) -> Tuple[str, str]:
    """Extract document content"""
    if file is None:
        return "", "❌ Please upload a file"

    try:
        # Get file path
        file_path = _extract_file_path(file)
        if not file_path:
            return "", f"❌ Invalid file format: {type(file)}"

        # Simple validation
        file_path = Path(file_path)
        if not file_path.exists():
            return "", f"❌ File not found: {file_path.name}"

        file_ext = file_path.suffix.lower()
        if file_ext not in [".pdf", ".docx"]:
            return "", f"❌ Unsupported file type: {file_ext}"

        # Extract content
        extractor = get_extractor()
        content = extractor.extract(str(file_path))

        if not content or not content.strip():
            return "", f"❌ Unable to extract content from file: {file_path.name}"

        # Statistics
        char_count = len(content)
        word_count = len(content.split())

        return (
            content,
            f"✅ Successfully extracted **{file_path.name}** ({char_count:,} characters, ~{word_count:,} words)",
        )

    except Exception as e:
        logger.error(f"❌ Document extraction error: {str(e)}")
        return "", f"❌ Extraction failed: {str(e)}"


def _extract_file_path(file) -> Optional[str]:
    """
    Extract file path from various file object types.

    Args:
        file: File object in various formats

    Returns:
        File path string or None if extraction fails
    """
    try:
        # Gradio file object with .name attribute
        if hasattr(file, "name") and file.name:
            return str(file.name)

        # String path
        if isinstance(file, str) and file.strip():
            return file.strip()

        # Dict-like object (some Gradio versions)
        if isinstance(file, dict) and "name" in file:
            return str(file["name"])

        # Path-like object
        if hasattr(file, "__fspath__"):
            return str(file)

        # Temporary file object
        if hasattr(file, "name") and hasattr(file, "read"):
            return str(file.name)

        logger.error(f"❌ Unknown file format: {file}, type: {type(file)}")
        return None

    except Exception as e:
        logger.error(f"❌ Error extracting file path: {e}")
        return None


def create_interface():
    """Create enhanced Gradio interface with better UX"""

    # Custom CSS for better styling
    custom_css = """
    .main-header { 
        text-align: center; 
        color: #2563eb; 
        margin-bottom: 2rem;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        padding: 2rem;
        border-radius: 1rem;
        color: white;
    }
    .upload-section { 
        background: linear-gradient(135deg, #f8fafc 0%, #e2e8f0 100%);
        padding: 2rem; 
        border-radius: 1rem; 
        margin-bottom: 2rem;
        border: 1px solid #cbd5e1;
    }
    .output-section { 
        background: #ffffff; 
        border: 2px solid #e2e8f0; 
        border-radius: 1rem;
        min-height: 400px;
    }
    .feature-list {
        background: #f1f5f9;
        padding: 1rem;
        border-radius: 0.5rem;
        margin: 1rem 0;
    }
    .stats-box {
        background: #ecfdf5;
        border: 1px solid #10b981;
        border-radius: 0.5rem;
        padding: 1rem;
        margin: 0.5rem 0;
    }
    """

    with gr.Blocks(
        title="📄 Document Extraction Tool",
        theme=gr.themes.Soft(),
        css=custom_css,
    ) as app:
        # Header
        gr.HTML("""
        <div class="main-header">
            <h1>📄 Document Extraction Tool</h1>
            <p style="font-size: 1.2em; margin-top: 1rem;">
                Convert PDF and DOCX files to clean Markdown format
            </p>
            <p style="opacity: 0.9; margin-top: 0.5rem;">
                Powered by Docling AI + PyPDF2 fallback • Fast • Reliable • Free
            </p>
        </div>
        """)

        with gr.Row():
            # Left column - Upload section
            with gr.Column(scale=1):
                with gr.Group(elem_classes="upload-section"):
                    gr.HTML("<h3>📂 Upload Document</h3>")

                    file_input = gr.File(
                        label="Select PDF or DOCX file",
                        file_types=[".pdf", ".docx"],
                        type="filepath",
                        height=120,
                    )

                    with gr.Row():
                        extract_btn = gr.Button(
                            "🔄 Extract Content", variant="primary", size="lg", scale=2
                        )
                        clear_btn = gr.Button("🗑️ Clear", variant="secondary", scale=1)

                    # Feature highlights
                    gr.HTML("""
                    <div class="feature-list">
                        <h4>✨ Features:</h4>
                        <ul style="margin: 0.5rem 0;">
                            <li>🚀 AI-powered extraction (Docling)</li>
                            <li>📄 PDF & DOCX support</li>
                            <li>🔄 Multiple fallback methods</li>
                            <li>📝 Clean Markdown output</li>
                            <li>💾 Copy-paste ready</li>
                        </ul>
                    </div>
                    """)

            # Right column - Results section
            with gr.Column(scale=2):
                with gr.Group(elem_classes="output-section"):
                    gr.HTML("<h3>📝 Extraction Results</h3>")

                    # Status with stats
                    with gr.Row():
                        status_output = gr.Textbox(
                            label="📊 Status", interactive=False, max_lines=2, scale=3
                        )

                    # Main content output
                    content_output = gr.Textbox(
                        label="📄 Markdown Content",
                        lines=25,
                        max_lines=50,
                        interactive=False,
                        show_copy_button=True,
                        placeholder="Extracted content will appear here...\n\nSupported formats:\n• PDF documents\n• DOCX documents\n\nUpload a file to get started!",
                    )

                    # Download button (hidden initially)
                    download_btn = gr.DownloadButton(
                        "💾 Download Markdown", visible=False, variant="secondary"
                    )

        # Event handlers
        def process_and_update(file):
            """Process file and update interface"""
            content, status = extract_document(file)

            # Show download button if content exists
            show_download = bool(content and content.strip())

            # Prepare download file if content exists
            download_file = None
            if show_download and file:
                try:
                    file_name = Path(_extract_file_path(file) or "document").stem
                    temp_file = tempfile.NamedTemporaryFile(
                        mode="w",
                        suffix=".md",
                        prefix=f"{file_name}_",
                        delete=False,
                        encoding="utf-8",
                    )
                    temp_file.write(content)
                    temp_file.close()
                    download_file = temp_file.name
                except Exception as e:
                    logger.warning(f"Failed to create download file: {e}")

            return (
                content,
                status,
                gr.update(visible=show_download, value=download_file),
            )

        def clear_all():
            """Clear all inputs and outputs"""
            return (
                None,  # file_input
                "",  # content_output
                "Ready to process new document",  # status_output
                gr.update(visible=False),  # download_btn
            )

        # Bind events
        extract_btn.click(
            fn=process_and_update,
            inputs=[file_input],
            outputs=[content_output, status_output, download_btn],
        )

        # Auto-extract when file is uploaded
        file_input.change(
            fn=process_and_update,
            inputs=[file_input],
            outputs=[content_output, status_output, download_btn],
        )

        # Clear button
        clear_btn.click(
            fn=clear_all,
            outputs=[file_input, content_output, status_output, download_btn],
        )

        # Footer
        gr.HTML("""
        <div style="text-align: center; margin-top: 2rem; padding: 2rem; background: #f8fafc; border-radius: 1rem;">
            <p style="color: #64748b; margin-bottom: 1rem;">
                <strong>Supported formats:</strong> PDF, DOCX • <strong>Max file size:</strong> 50MB
            </p>
            <p style="color: #64748b; font-size: 0.9em;">
                Also available via Claude Desktop's MCP integration for seamless workflow
            </p>
        </div>
        """)

    return app


def main():
    """Launch application"""
    try:
        logger.info("🚀 Starting document extraction tool...")

        # Pre-initialize extractor
        logger.info("🔧 Initializing extractor...")
        get_extractor()  # Initialize global extractor instance
        logger.info("✅ Extractor initialization complete, ready to extract")

        # Create interface
        logger.info("🎨 Creating user interface...")
        app = create_interface()

        # Launch configuration
        server_port = int(os.getenv("GRADIO_SERVER_PORT", "7860"))

        logger.info(f"🌐 Starting server on port: {server_port}")
        logger.info(f"📱 Access URL: http://127.0.0.1:{server_port}")

        # Launch application
        app.launch(
            server_name="0.0.0.0",
            server_port=server_port,
            mcp_server=True,
            share=False,
            show_error=True,
        )

    except KeyboardInterrupt:
        logger.info("🛑 Application stopped by user")
    except Exception as e:
        logger.error(f"❌ Application startup failed: {str(e)}")
        raise


if __name__ == "__main__":
    main()