try:
    import spaces
except ImportError:
    class spaces:
        @staticmethod
        def GPU(fn=None, duration=None):
            if fn is None:
                def decorator(f):
                    return f
                return decorator
            return fn

import os
import shutil
import tempfile
import zipfile
import time
import uuid
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
import gradio as gr

def has_selectable_text(pdf_path):
    """Checks if the PDF has selectable text to skip OCR for digital documents."""
    try:
        import pypdf
        reader = pypdf.PdfReader(pdf_path)
        text_sample = ""
        num_pages = len(reader.pages)
        for i in range(min(10, num_pages)):
            page_text = reader.pages[i].extract_text()
            if page_text:
                text_sample += page_text
                if len(text_sample.strip()) > 50:
                    return True
        return len(text_sample.strip()) > 50
    except Exception as e:
        print(f"Error checking text selectability: {e}")
        return True

def convert_single_doc(file_path, filename, digital_converter, scanned_converter, output_dir):
    """Processes a single document. Executed inside the GPU node's local thread pool."""
    start_time = time.time()
    try:
        print(f"[{filename}] STARTING conversion on GPU...", flush=True)
        is_digital = has_selectable_text(file_path)
        converter = digital_converter if is_digital else scanned_converter
        
        result = converter.convert(file_path)
        markdown_content = result.document.export_to_markdown()
        
        # Save output markdown file
        base_name = os.path.splitext(filename)[0]
        output_path = os.path.join(output_dir, f"{base_name}.md")
        with open(output_path, "w", encoding="utf-8") as f:
            f.write(markdown_content)
        
        elapsed = time.time() - start_time
        try:
            import pypdf
            reader = pypdf.PdfReader(file_path)
            pages = len(reader.pages)
        except Exception:
            pages = "N/A"
            
        print(f"[{filename}] COMPLETED successfully in {elapsed:.1f}s ({pages} pages)", flush=True)
        return {
            "name": filename,
            "pages": pages,
            "status": "Success",
            "time": f"{elapsed:.1f}s",
            "content": markdown_content
        }
    except Exception as e:
        elapsed = time.time() - start_time
        print(f"[{filename}] FAILED after {elapsed:.1f}s: {e}", flush=True)
        return {
            "name": filename,
            "pages": "N/A",
            "status": f"Error: {str(e)}",
            "time": f"{elapsed:.1f}s",
            "content": ""
        }

# Define a batch conversion function decorated with ZeroGPU.
# This processes a list of files on a SINGLE GPU instance concurrently in VRAM
# using a local thread pool, maximizing GPU core occupancy and minimizing GPU-second duration.
@spaces.GPU(duration=480)
def convert_file_batch_on_single_gpu(batch_info: list, output_dir: str) -> list:
    # Local imports to prevent global PyTorch initialization on CPU container
    from docling.datamodel.base_models import InputFormat
    from docling.datamodel.pipeline_options import PdfPipelineOptions, EasyOcrOptions
    from docling.document_converter import DocumentConverter, PdfFormatOption
    from docling.backend.pypdfium2_backend import PyPdfiumDocumentBackend
    from docling.datamodel.accelerator_options import AcceleratorOptions, AcceleratorDevice

    # 1. Configure Docling pipeline options
    def get_pipeline_options(do_ocr):
        opts = PdfPipelineOptions()
        opts.do_ocr = do_ocr
        if do_ocr:
            opts.ocr_options = EasyOcrOptions(
                lang=["en"],
                use_gpu=True
            )
        opts.do_picture_classification = False
        opts.do_picture_description = False
        opts.do_chart_extraction = False
        opts.do_code_enrichment = False
        opts.do_formula_enrichment = False
        opts.generate_page_images = False
        opts.generate_picture_images = False
        opts.accelerator_options = AcceleratorOptions(
            device=AcceleratorDevice.CUDA,
            num_threads=2
        )
        return opts

    # Pre-instantiate converters to share VRAM weights across the concurrent threads
    # One for digital (no OCR) and one for scanned (OCR enabled)
    digital_converter = DocumentConverter(
        format_options={
            InputFormat.PDF: PdfFormatOption(
                pipeline_options=get_pipeline_options(do_ocr=False),
                backend=PyPdfiumDocumentBackend
            )
        }
    )
    scanned_converter = DocumentConverter(
        format_options={
            InputFormat.PDF: PdfFormatOption(
                pipeline_options=get_pipeline_options(do_ocr=True),
                backend=PyPdfiumDocumentBackend
            )
        }
    )

    results = []
    
    # Process files concurrently inside the A10G/RTX6000 VRAM using a local thread pool
    # With 48GB VRAM (Blackwell), we can easily handle 8 concurrent document streams sharing weights
    local_concurrency = min(8, len(batch_info))
    print(f"Executing local thread pool inside GPU container with {local_concurrency} workers...", flush=True)
    
    with ThreadPoolExecutor(max_workers=local_concurrency) as local_executor:
        futures = {
            local_executor.submit(
                convert_single_doc,
                file_path, filename,
                digital_converter, scanned_converter,
                output_dir
            ): filename
            for file_path, filename in batch_info
        }
        for future in as_completed(futures):
            results.append(future.result())
            
    return results

def run_batch_conversion(uploaded_files, max_workers):
    if not uploaded_files:
        yield "No files uploaded.", gr.update(visible=False), gr.update(visible=False), [], gr.update(visible=False), {}
        return

    session_id = uuid.uuid4().hex
    temp_workspace = os.path.join(tempfile.gettempdir(), f"docling_workspace_{session_id}")
    input_dir = os.path.join(temp_workspace, "inputs")
    output_dir = os.path.join(temp_workspace, "outputs")
    os.makedirs(input_dir, exist_ok=True)
    os.makedirs(output_dir, exist_ok=True)

    files_to_process = []

    for file_obj in uploaded_files:
        path = file_obj.name
        filename = os.path.basename(path)
        
        if filename.lower().endswith(".zip"):
            print(f"Extracting ZIP archive: {filename}")
            zip_extract_dir = os.path.join(input_dir, f"zip_{uuid.uuid4().hex}")
            os.makedirs(zip_extract_dir, exist_ok=True)
            try:
                with zipfile.ZipFile(path, 'r') as zip_ref:
                    zip_ref.extractall(zip_extract_dir)
                for root, _, files in os.walk(zip_extract_dir):
                    for f in files:
                        ext = os.path.splitext(f)[1].lower()
                        if ext in [".pdf", ".docx", ".pptx", ".html", ".png", ".jpg", ".jpeg"]:
                            full_path = os.path.join(root, f)
                            files_to_process.append((full_path, f))
            except Exception as e:
                print(f"Error extracting ZIP: {e}")
        else:
            dest_path = os.path.join(input_dir, filename)
            shutil.copy(path, dest_path)
            files_to_process.append((dest_path, filename))

    if not files_to_process:
        yield "No supported document files found to convert.", gr.update(visible=False), gr.update(visible=False), [], gr.update(visible=False), {}
        return

    # Process files in GPU-efficient chunks.
    # Grouping files into chunks of 8 enables concurrent processing in VRAM on a single GPU node.
    chunk_size = 8
    chunks = [files_to_process[i:i + chunk_size] for i in range(0, len(files_to_process), chunk_size)]
    
    results = []
    print(f"Processing {len(files_to_process)} files in {len(chunks)} GPU batches (max concurrent batches={max_workers})...", flush=True)
    
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {
            executor.submit(convert_file_batch_on_single_gpu, chunk, output_dir): idx
            for idx, chunk in enumerate(chunks)
        }
        for future in as_completed(futures):
            chunk_results = future.result()
            results.extend(chunk_results)
            
            # Package current ZIP
            zip_output_path = os.path.join(temp_workspace, "converted_markdown_files.zip")
            with zipfile.ZipFile(zip_output_path, 'w', zipfile.ZIP_DEFLATED) as zip_out:
                for root, _, files in os.walk(output_dir):
                    for f in files:
                        full_path = os.path.join(root, f)
                        arcname = os.path.relpath(full_path, output_dir)
                        zip_out.write(full_path, arcname)

            table_data = [
                [r["name"], r["pages"], r["status"], r["time"]]
                for r in results
            ]

            success_count = len([r for r in results if r['status'] == 'Success'])
            summary_text = f"Successfully processed {success_count} out of {len(files_to_process)} files."
            
            preview_dict = {r["name"]: r["content"] for r in results if r["status"] == "Success"}
            dropdown_choices = list(preview_dict.keys())
            dropdown_update = gr.update(choices=dropdown_choices, value=dropdown_choices[0] if dropdown_choices else None, visible=bool(dropdown_choices))
            preview_box_update = gr.update(visible=bool(dropdown_choices))

            yield (
                summary_text,
                gr.update(value=zip_output_path, visible=True),
                table_data,
                dropdown_update,
                preview_box_update,
                preview_dict
            )

theme = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="blue",
    neutral_hue="slate",
).set(
    body_background_fill="*neutral_950",
    block_background_fill="*neutral_900",
    block_border_color="*neutral_800",
    button_primary_background_fill="linear-gradient(90deg, *primary_600, *secondary_600)",
    button_primary_text_color="white",
    block_title_text_color="*primary_400",
)

css = """
body {
    background-color: #0b0f19;
    font-family: 'Outfit', 'Inter', sans-serif;
}
.gradio-container {
    max-width: 1200px !important;
    margin: 0 auto !important;
}
h1 {
    background: linear-gradient(90deg, #818cf8, #3b82f6);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    font-weight: 800;
}
.glass-panel {
    background: rgba(30, 41, 59, 0.4) !important;
    backdrop-filter: blur(12px);
    border: 1px solid rgba(255, 255, 255, 0.05);
    border-radius: 12px;
}
"""

with gr.Blocks(theme=theme, css=css, title="Docling GPU Batch Converter") as demo:
    preview_state = gr.State({})

    with gr.Column():
        gr.Markdown(
            """
            # 📄 Docling GPU Batch Converter
            ### IBM Docling parallel document-to-markdown parsing powered by Hugging Face ZeroGPU.
            """
        )
        
        with gr.Row():
            with gr.Column(scale=1, elem_classes=["glass-panel"]):
                gr.Markdown("### 📥 Upload Documents")
                file_input = gr.File(
                    file_count="multiple",
                    label="Upload PDFs, DOCX, PPTX or a ZIP archive containing them",
                    file_types=[".pdf", ".docx", ".pptx", ".html", ".png", ".jpg", ".zip"]
                )
                
                with gr.Row():
                    workers_slider = gr.Slider(
                        minimum=1,
                        maximum=4,
                        value=1,
                        step=1,
                        label="Parallel GPU Nodes"
                    )
                
                submit_btn = gr.Button("⚡ Convert Batch", variant="primary")
                download_output = gr.File(label="📦 Download Converted MD (ZIP)", visible=False)
                
            with gr.Column(scale=2, elem_classes=["glass-panel"]):
                gr.Markdown("### 📊 Conversion Progress & Status")
                status_summary = gr.Markdown("Ready to process.")
                progress_table = gr.Dataframe(
                    headers=["File Name", "Pages", "Status", "Time Taken"],
                    datatype=["str", "str", "str", "str"],
                    value=[]
                )
                
        with gr.Column(visible=False) as preview_container:
            gr.Markdown("### 🔍 Document Markdown Preview")
            preview_selector = gr.Dropdown(label="Select document to preview", choices=[], interactive=True)
            preview_markdown = gr.Markdown(label="Converted Output")

    def update_preview_text(selected_file, data_dict):
        if selected_file in data_dict:
            return data_dict[selected_file]
        return "No content available."

    submit_btn.click(
        fn=run_batch_conversion,
        inputs=[file_input, workers_slider],
        outputs=[status_summary, download_output, progress_table, preview_selector, preview_container, preview_state]
    )

    preview_selector.change(
        fn=update_preview_text,
        inputs=[preview_selector, preview_state],
        outputs=preview_markdown
    )

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