import os
import shutil
import tempfile
import re
from pathlib import Path
from typing import Optional, Tuple, List, Dict, Any

import gradio as gr

from doctra.parsers.structured_pdf_parser import StructuredPDFParser
from doctra.parsers.table_chart_extractor import ChartTablePDFParser


def _gather_outputs(out_dir: Path, allowed_kinds: Optional[List[str]] = None, zip_filename: Optional[str] = None, is_structured_parsing: bool = False) -> Tuple[List[tuple[str, str]], List[str], str]:
    gallery_items: List[tuple[str, str]] = []
    file_paths: List[str] = []

    if out_dir.exists():
        if is_structured_parsing:
            # For structured parsing, show ALL files in the directory
            for file_path in sorted(out_dir.rglob("*")):
                if file_path.is_file():
                    file_paths.append(str(file_path))
        else:
            # For full parsing, use the original logic
            # Always add main output files (HTML, Markdown, etc.) regardless of allowed_kinds
            main_files = [
                "result.html",
                "result.md", 
                "tables.html",
                "tables.xlsx"
            ]
            
            for main_file in main_files:
                file_path = out_dir / main_file
                if file_path.exists():
                    file_paths.append(str(file_path))
            
            # Add image files based on allowed_kinds or all images if not specified
            if allowed_kinds:
                for kind in allowed_kinds:
                    # ChartTablePDFParser saves directly to charts/ and tables/ directories
                    p = out_dir / kind
                    if p.exists():
                        for img in sorted(p.glob("*.png")):  # ChartTablePDFParser saves as .png
                            file_paths.append(str(img))
                    
                    # Also check images/ subdirectories (for StructuredPDFParser)
                    images_dir = out_dir / "images" / kind
                    if images_dir.exists():
                        for img in sorted(images_dir.glob("*.jpg")):  # StructuredPDFParser saves as .jpg
                            file_paths.append(str(img))
            else:
                # Fallback: look in both direct directories and images/ subdirectories
                for p in (out_dir / "charts").glob("*.png"):
                    file_paths.append(str(p))
                for p in (out_dir / "tables").glob("*.png"):
                    file_paths.append(str(p))
                for p in (out_dir / "images").rglob("*.jpg"):
                    file_paths.append(str(p))

            # Add Excel files based on extraction target (for structured parsing)
            if allowed_kinds:
                if "charts" in allowed_kinds and "tables" in allowed_kinds:
                    excel_files = ["parsed_tables_charts.xlsx"]
                elif "charts" in allowed_kinds:
                    excel_files = ["parsed_charts.xlsx"]
                elif "tables" in allowed_kinds:
                    excel_files = ["parsed_tables.xlsx"]
                else:
                    excel_files = []
                
                for excel_file in excel_files:
                    excel_path = out_dir / excel_file
                    if excel_path.exists():
                        file_paths.append(str(excel_path))

    kinds = allowed_kinds if allowed_kinds else ["tables", "charts", "figures"]
    for sub in kinds:
        # Look in both direct directories and images/ subdirectories
        # First try direct directories (for ChartTablePDFParser)
        p = out_dir / sub
        if p.exists():
            for img in sorted(p.glob("*.png")):  # ChartTablePDFParser saves as .png
                gallery_items.append((str(img), f"{sub}: {img.name}"))
        
        # Also try images/ subdirectories (for StructuredPDFParser)
        images_dir = out_dir / "images" / sub
        if images_dir.exists():
            for img in sorted(images_dir.glob("*.jpg")):  # StructuredPDFParser saves as .jpg
                gallery_items.append((str(img), f"{sub}: {img.name}"))

    tmp_zip_dir = Path(tempfile.mkdtemp(prefix="doctra_zip_"))
    
    # Use custom filename if provided, otherwise use default
    if zip_filename:
        # Clean the filename to be safe for file systems
        safe_filename = re.sub(r'[<>:"/\\|?*]', '_', zip_filename)
        zip_base = tmp_zip_dir / safe_filename
    else:
        zip_base = tmp_zip_dir / "doctra_outputs"
    
    # Create a filtered copy of the output directory excluding temp files
    filtered_dir = tmp_zip_dir / "filtered_outputs"
    shutil.copytree(out_dir, filtered_dir, ignore=shutil.ignore_patterns('~$*', '*.tmp', '*.temp'))
    
    zip_path = shutil.make_archive(str(zip_base), 'zip', root_dir=str(filtered_dir))

    return gallery_items, file_paths, zip_path


def _parse_markdown_by_pages(md_content: str) -> List[Dict[str, Any]]:
    """
    Parse markdown content and organize it by pages.
    Returns a list of page dictionaries with content, tables, charts, and figures.
    """
    import re
    
    pages = []
    current_page = None
    
    lines = md_content.split('\n')
    i = 0
    
    while i < len(lines):
        line = lines[i].strip()
        
        # Check for page header
        if line.startswith('## Page '):
            # Save previous page if exists
            if current_page:
                pages.append(current_page)
            
            # Start new page
            page_num = line.replace('## Page ', '').strip()
            current_page = {
                'page_num': page_num,
                'content': [],
                'tables': [],
                'charts': [],
                'figures': [],
                'images': [],
                'full_content': []  # Store full content with inline images
            }
            i += 1
            continue
        
        # Check for images (tables, charts, figures)
        if line.startswith('![') and '](images/' in line:
            # Extract image info
            match = re.match(r'!\[([^\]]+)\]\(([^)]+)\)', line)
            if match:
                caption = match.group(1)
                img_path = match.group(2)
                
                # Categorize by type
                if 'Table' in caption:
                    current_page['tables'].append({'caption': caption, 'path': img_path})
                elif 'Chart' in caption:
                    current_page['charts'].append({'caption': caption, 'path': img_path})
                elif 'Figure' in caption:
                    current_page['figures'].append({'caption': caption, 'path': img_path})
                
                current_page['images'].append({'caption': caption, 'path': img_path})
                
                # Add to full content with proper markdown formatting
                current_page['full_content'].append(f"![{caption}]({img_path})")
        
        # Regular content
        elif current_page:
            if line:  # Only add non-empty lines
                current_page['content'].append(line)
            current_page['full_content'].append(line)
        
        i += 1
    
    # Add the last page
    if current_page:
        pages.append(current_page)
    
    return pages


def run_full_parse(
    pdf_file: str,
    use_vlm: bool,
    vlm_provider: str,
    vlm_api_key: str,
    layout_model_name: str,
    dpi: int,
    min_score: float,
    ocr_lang: str,
    ocr_psm: int,
    ocr_oem: int,
    ocr_extra_config: str,
    box_separator: str,
) -> Tuple[str, Optional[str], List[tuple[str, str]], List[str], str]:
    if not pdf_file:
        return ("No file provided.", None, [], [], "")

    # Extract filename from the uploaded file path
    # Gradio provides the original filename in the file path
    original_filename = Path(pdf_file).stem
    
    tmp_dir = Path(tempfile.mkdtemp(prefix="doctra_"))
    # Use original filename for temp file so parser creates correct output directory
    input_pdf = tmp_dir / f"{original_filename}.pdf"
    shutil.copy2(pdf_file, input_pdf)

    # Validate VLM configuration
    if use_vlm and not vlm_api_key:
        return ("❌ Error: VLM API key is required when using VLM", None, [], [], "")
    
    if use_vlm and vlm_api_key:
        # Basic API key validation
        if len(vlm_api_key.strip()) < 10:
            return ("❌ Error: VLM API key appears to be too short or invalid", None, [], [], "")
        if vlm_api_key.strip().startswith('sk-') and len(vlm_api_key.strip()) < 20:
            return ("❌ Error: OpenAI API key appears to be invalid (too short)", None, [], [], "")

    parser = StructuredPDFParser(
        use_vlm=use_vlm,
        vlm_provider=vlm_provider,
        vlm_api_key=vlm_api_key or None,
        layout_model_name=layout_model_name,
        dpi=int(dpi),
        min_score=float(min_score),
        ocr_lang=ocr_lang,
        ocr_psm=int(ocr_psm),
        ocr_oem=int(ocr_oem),
        ocr_extra_config=ocr_extra_config or "",
        box_separator=box_separator or "\n",
    )

    try:
        parser.parse(str(input_pdf))
    except Exception as e:
        import traceback
        traceback.print_exc()
        # Safely encode error message for return value
        try:
            error_msg = str(e).encode('utf-8', errors='replace').decode('utf-8')
            return (f"❌ VLM processing failed: {error_msg}", None, [], [], "")
        except Exception:
            return (f"❌ VLM processing failed: <Unicode encoding error>", None, [], [], "")

    outputs_root = Path("outputs")
    out_dir = outputs_root / original_filename / "full_parse"
    if not out_dir.exists():
        # fallback: search latest created dir under outputs
        candidates = sorted(outputs_root.glob("*/"), key=lambda p: p.stat().st_mtime, reverse=True)
        if candidates:
            out_dir = candidates[0] / "full_parse"
        else:
            out_dir = outputs_root

    md_file = next(out_dir.glob("*.md"), None)
    md_preview = None
    if md_file and md_file.exists():
        try:
            with md_file.open("r", encoding="utf-8", errors="ignore") as f:
                md_preview = f.read()  # Return the full markdown content
        except Exception:
            md_preview = None

    gallery_items, file_paths, zip_path = _gather_outputs(out_dir, zip_filename=original_filename, is_structured_parsing=False)
    return (f"✅ Parsing completed successfully!\n📁 Output directory: {out_dir}", md_preview, gallery_items, file_paths, zip_path)


def run_extract(
    pdf_file: str,
    target: str,
    use_vlm: bool,
    vlm_provider: str,
    vlm_api_key: str,
    layout_model_name: str,
    dpi: int,
    min_score: float,
) -> Tuple[str, str, List[tuple[str, str]], List[str], str]:
    if not pdf_file:
        return ("No file provided.", "", [], [], "")
    
    # Validate VLM configuration
    if use_vlm and not vlm_api_key:
        return ("❌ Error: VLM API key is required when using VLM", "", [], [], "")
    
    if use_vlm and vlm_api_key:
        # Basic API key validation
        if len(vlm_api_key.strip()) < 10:
            return ("❌ Error: VLM API key appears to be too short or invalid", "", [], [], "")
        if vlm_api_key.strip().startswith('sk-') and len(vlm_api_key.strip()) < 20:
            return ("❌ Error: OpenAI API key appears to be invalid (too short)", "", [], [], "")

    # Extract filename from the uploaded file path
    original_filename = Path(pdf_file).stem
    
    tmp_dir = Path(tempfile.mkdtemp(prefix="doctra_"))
    # Use original filename for temp file so parser creates correct output directory
    input_pdf = tmp_dir / f"{original_filename}.pdf"
    shutil.copy2(pdf_file, input_pdf)

    parser = ChartTablePDFParser(
        extract_charts=(target in ("charts", "both")),
        extract_tables=(target in ("tables", "both")),
        use_vlm=use_vlm,
        vlm_provider=vlm_provider,
        vlm_api_key=vlm_api_key or None,
        layout_model_name=layout_model_name,
        dpi=int(dpi),
        min_score=float(min_score),
    )

    output_base = Path("outputs")
    parser.parse(str(input_pdf), str(output_base))

    outputs_root = output_base
    out_dir = outputs_root / original_filename / "structured_parsing"
    if not out_dir.exists():
        if outputs_root.exists():
            candidates = sorted(outputs_root.glob("*/"), key=lambda p: p.stat().st_mtime, reverse=True)
            if candidates:
                out_dir = candidates[0] / "structured_parsing"
            else:
                out_dir = outputs_root
        else:
            outputs_root.mkdir(parents=True, exist_ok=True)
            out_dir = outputs_root

    # Determine which kinds to include in outputs based on target selection
    allowed_kinds: Optional[List[str]] = None
    if target in ("tables", "charts"):
        allowed_kinds = [target]
    elif target == "both":
        allowed_kinds = ["tables", "charts"]

    gallery_items, file_paths, zip_path = _gather_outputs(out_dir, allowed_kinds, zip_filename=original_filename, is_structured_parsing=True)

    # Build tables HTML preview from Excel data (when VLM enabled)
    tables_html = ""
    try:
        if use_vlm:
            # Find Excel file based on target
            excel_filename = None
            if target in ("tables", "charts"):
                if target == "tables":
                    excel_filename = "parsed_tables.xlsx"
                else:  # charts
                    excel_filename = "parsed_charts.xlsx"
            elif target == "both":
                excel_filename = "parsed_tables_charts.xlsx"
            
            if excel_filename:
                excel_path = out_dir / excel_filename
                if excel_path.exists():
                    import pandas as pd
                    import html as _html
                    
                    # Read Excel file and create HTML tables
                    xl_file = pd.ExcelFile(excel_path)
                    html_blocks = []
                    
                    for sheet_name in xl_file.sheet_names:
                        df = pd.read_excel(excel_path, sheet_name=sheet_name)
                        if not df.empty:
                            # Create table with title
                            title = f"<h3>{_html.escape(sheet_name)}</h3>"
                            
                            # Convert DataFrame to HTML table
                            table_html = df.to_html(
                                classes="doc-table",
                                table_id=None,
                                escape=True,
                                index=False,
                                na_rep=""
                            )
                            
                            html_blocks.append(title + table_html)
                    
                    tables_html = "\n".join(html_blocks)
    except Exception as e:
        # Safely encode error message to handle Unicode characters
        try:
            error_msg = str(e).encode('utf-8', errors='replace').decode('utf-8')
            print(f"Error building tables HTML: {error_msg}")
        except Exception:
            print(f"Error building tables HTML: <Unicode encoding error>")
        tables_html = ""

    return (f"✅ Parsing completed successfully!\n📁 Output directory: {out_dir}", tables_html, gallery_items, file_paths, zip_path)


THEME = gr.themes.Soft(primary_hue="indigo", neutral_hue="slate")

CUSTOM_CSS = """
.gradio-container {max-width: 100% !important; padding-left: 24px; padding-right: 24px}
.container {max-width: 100% !important}
.app {max-width: 100% !important}
.header {margin-bottom: 8px}
.subtitle {color: var(--body-text-color-subdued)}
.card {border:1px solid var(--border-color); border-radius:12px; padding:8px}
.status-ok {color: var(--color-success)}

/* Page content styling */
.page-content img {
    max-width: 100% !important;
    height: auto !important;
    display: block !important;
    margin: 10px auto !important;
    border: 1px solid #ddd !important;
    border-radius: 8px !important;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1) !important;
}

.page-content {
    max-height: none !important;
    overflow: visible !important;
}

/* Table styling */
.page-content table.doc-table { 
    width: 100% !important; 
    border-collapse: collapse !important; 
    margin: 12px 0 !important; 
}
.page-content table.doc-table th,
.page-content table.doc-table td { 
    border: 1px solid #e5e7eb !important; 
    padding: 8px 10px !important; 
    text-align: left !important; 
}
.page-content table.doc-table thead th { 
    background: #f9fafb !important; 
    font-weight: 600 !important; 
}
.page-content table.doc-table tbody tr:nth-child(even) td { 
    background: #fafafa !important; 
}

/* Clickable image buttons */
.image-button {
    background: #0066cc !important;
    color: white !important;
    border: none !important;
    padding: 5px 10px !important;
    border-radius: 4px !important;
    cursor: pointer !important;
    margin: 2px !important;
    font-size: 14px !important;
}

.image-button:hover {
    background: #0052a3 !important;
}
"""


def build_demo() -> gr.Blocks:
    with gr.Blocks(title="Doctra - Document Parser", theme=THEME, css=CUSTOM_CSS) as demo:
        gr.Markdown(
            """
<div class="header">
  <h2 style="margin:0">Doctra — Document Parser</h2>
  <div class="subtitle">Parse PDFs, extract tables/charts, preview markdown, and download outputs.</div>
</div>
            """
        )
        

        with gr.Tab("Full Parse"):
            with gr.Row():
                pdf = gr.File(file_types=[".pdf"], label="PDF")
                use_vlm = gr.Checkbox(label="Use VLM (optional)", value=False)
                vlm_provider = gr.Dropdown(["gemini", "openai", "anthropic", "openrouter"], value="gemini", label="VLM Provider")
                vlm_api_key = gr.Textbox(type="password", label="VLM API Key", placeholder="Optional if VLM disabled")

            with gr.Accordion("Advanced", open=False):
                with gr.Row():
                    layout_model = gr.Textbox(value="PP-DocLayout_plus-L", label="Layout model")
                    dpi = gr.Slider(100, 400, value=200, step=10, label="DPI")
                    min_score = gr.Slider(0, 1, value=0.0, step=0.05, label="Min layout score")
                with gr.Row():
                    ocr_lang = gr.Textbox(value="eng", label="OCR Language")
                    ocr_psm = gr.Slider(0, 13, value=4, step=1, label="Tesseract PSM")
                    ocr_oem = gr.Slider(0, 3, value=3, step=1, label="Tesseract OEM")
                with gr.Row():
                    ocr_config = gr.Textbox(value="", label="Extra OCR config")
                    box_sep = gr.Textbox(value="\n", label="Box separator")

            run_btn = gr.Button("▶ Run Full Parse", variant="primary")
            status = gr.Textbox(label="Status", elem_classes=["status-ok"])
            
            # Page selector for extracted content
            page_selector = gr.Dropdown(label="Select Page to Display", interactive=True, visible=False)
            
            # Full Parse components
            with gr.Row():
                with gr.Column():
                    md_preview = gr.HTML(label="Extracted Content", visible=True, elem_classes=["page-content"])
                with gr.Column():
                    page_image = gr.Image(label="Page image", interactive=False)
            gallery = gr.Gallery(label="Extracted images (tables/charts/figures)", columns=4, height=420, preview=True)
            files_out = gr.Files(label="Download individual output files")
            zip_out = gr.File(label="Download all outputs (ZIP)")
            
            # Hidden state to store pages data and all images
            pages_state = gr.State([])
            all_images_state = gr.State([])
            pdf_path_state = gr.State("")
            page_images_state = gr.State([])  # list of file paths per page index (1-based)
            
            # Hidden components for image filtering
            filter_trigger = gr.Button(visible=False)
            current_image_path = gr.State("")
            current_image_caption = gr.State("")
            image_filter_input = gr.Textbox(visible=False, elem_id="image_filter_input")

            def parse_markdown_by_pages(md_content: str):
                """Parse markdown content and organize it by pages."""
                import re
                
                pages = []
                current_page = None
                
                lines = md_content.split('\n')
                i = 0
                
                
                # First, let's find all page headers
                page_headers = []
                for i, line in enumerate(lines):
                    if line.strip().startswith('## Page '):
                        page_num = line.strip().replace('## Page ', '').strip()
                        page_headers.append((i, page_num, line))
                
                
                # Now parse content for each page
                for i, (line_idx, page_num, header_line) in enumerate(page_headers):
                    # Find the end of this page (start of next page or end of document)
                    start_line = line_idx
                    if i + 1 < len(page_headers):
                        end_line = page_headers[i + 1][0]
                    else:
                        end_line = len(lines)
                    
                    # Extract content for this page
                    page_content = lines[start_line:end_line]
                    
                    page = {
                        'page_num': page_num,
                        'content': page_content
                    }
                    pages.append(page)
                
                return pages

            def update_page_selector(pages_data):
                """Update the page selector dropdown with available pages."""
                if not pages_data:
                    return gr.Dropdown(choices=[], value=None, visible=False)
                
                page_choices = [f"Page {page['page_num']}" for page in pages_data]
                return gr.Dropdown(choices=page_choices, value=page_choices[0], visible=True)

            def display_selected_page(selected_page, pages_data, pdf_path, page_images):
                """Display the content of the selected page and the rendered page image."""
                if not selected_page or not pages_data:
                    return "", None
                
                
                # Find the selected page
                page_num = selected_page.replace("Page ", "")
                page = next((p for p in pages_data if p['page_num'] == page_num), None)
                
                if not page:
                    return "Page not found", None
                
                # Build HTML with inline base64 images, render markdown tables, and preserve paragraphs/line breaks
                import html as _html, base64, re as _re
                base_dir = None
                try:
                    stem = Path(pdf_path).stem if pdf_path else ""
                    if stem:
                        base_dir = Path("outputs") / stem / "full_parse"
                except Exception:
                    base_dir = None
                processed_content = []
                paragraph_buffer = []
                def flush_paragraph():
                    nonlocal paragraph_buffer
                    if paragraph_buffer:
                        joined = '<br/>'.join(_html.escape(l) for l in paragraph_buffer)
                        processed_content.append(f'<p>{joined}</p>')
                        paragraph_buffer = []

                # Simple markdown table detection and rendering
                def is_md_table_header(s: str) -> bool:
                    return '|' in s and ('---' in s or '—' in s)

                def render_md_table(lines: List[str]) -> str:
                    rows = [l.strip().strip('|').split('|') for l in lines]
                    rows = [[_html.escape(c.strip()) for c in r] for r in rows]
                    if len(rows) < 2:
                        return ""
                    header = rows[0]
                    body = rows[2:] if len(rows) > 2 else []
                    thead = '<thead><tr>' + ''.join(f'<th>{c}</th>' for c in header) + '</tr></thead>'
                    tbody = '<tbody>' + ''.join('<tr>' + ''.join(f'<td>{c}</td>' for c in r) + '</tr>' for r in body) + '</tbody>'
                    return f'<table class="doc-table">{thead}{tbody}</table>'

                i = 0
                lines = page['content']
                n = len(lines)
                while i < n:
                    raw_line = lines[i]
                    line = raw_line.rstrip('\r\n')
                    stripped = line.strip()
                    if stripped.startswith('![') and ('](images/' in stripped or '](images\\' in stripped):
                        flush_paragraph()
                        match = _re.match(r'!\[([^\]]+)\]\(([^)]+)\)', stripped)
                        if match and base_dir is not None:
                            caption = match.group(1)
                            rel_path = match.group(2).replace('\\\\', '/').replace('\\', '/').lstrip('/')
                            abs_path = (base_dir / rel_path).resolve()
                            try:
                                with open(abs_path, 'rb') as f:
                                    b64 = base64.b64encode(f.read()).decode('ascii')
                                processed_content.append(f'<figure><img src="data:image/jpeg;base64,{b64}" alt="{_html.escape(caption)}"/><figcaption>{_html.escape(caption)}</figcaption></figure>')
                            except Exception:
                                processed_content.append(f'<div>{_html.escape(caption)} (image not found)</div>')
                        else:
                            paragraph_buffer.append(raw_line)
                        i += 1
                        continue

                    # Detect markdown table blocks - only if line starts with | or has multiple | characters
                    if (stripped.startswith('|') or stripped.count('|') >= 2) and i + 1 < n and is_md_table_header(lines[i + 1]):
                        flush_paragraph()
                        table_block = [stripped]
                        i += 1
                        table_block.append(lines[i].strip())
                        i += 1
                        while i < n:
                            nxt = lines[i].rstrip('\r\n')
                            if nxt.strip() == '' or (not nxt.strip().startswith('|') and nxt.count('|') < 2):
                                break
                            table_block.append(nxt.strip())
                            i += 1
                        html_table = render_md_table(table_block)
                        if html_table:
                            processed_content.append(html_table)
                        else:
                            for tl in table_block:
                                paragraph_buffer.append(tl)
                        continue

                    if stripped.startswith('## '):
                        flush_paragraph()
                        processed_content.append(f'<h3>{_html.escape(stripped[3:])}</h3>')
                    elif stripped.startswith('# '):
                        flush_paragraph()
                        processed_content.append(f'<h2>{_html.escape(stripped[2:])}</h2>')
                    elif stripped == '':
                        flush_paragraph()
                        processed_content.append('<br/>')
                    else:
                        paragraph_buffer.append(raw_line)
                    i += 1
                flush_paragraph()
                
                # Join the processed content lines
                content = "\n".join(processed_content)

                # Ensure page images are prepared
                try:
                    if pdf_path and not page_images:
                        from doctra.utils.pdf_io import render_pdf_to_images
                        tmp_img_dir = Path(tempfile.mkdtemp(prefix="doctra_pages_"))
                        pil_pages = render_pdf_to_images(pdf_path)
                        saved_paths: List[str] = []
                        for idx, (im, _, _) in enumerate(pil_pages, start=1):
                            out_path = tmp_img_dir / f"page_{idx:03d}.jpg"
                            im.save(out_path, format="JPEG", quality=90)
                            saved_paths.append(str(out_path))
                        page_images = saved_paths
                        page_images_state.value = saved_paths  # cache
                except Exception as e:
                    pass

                # Select image for the current page number (1-based)
                page_img = None
                try:
                    page_index = int(page_num)
                    if page_images and 1 <= page_index <= len(page_images):
                        page_img = page_images[page_index - 1]
                except Exception:
                    page_img = None

                return content, page_img

            def filter_gallery_by_image(img_path, caption, all_images):
                """Filter gallery to show only the selected image."""
                if not img_path or not all_images:
                    return all_images
                
                # Find the selected image
                filtered_images = []
                for stored_img_path, stored_caption in all_images:
                    if stored_caption == caption:
                        filtered_images.append((stored_img_path, stored_caption))
                        break
                
                return filtered_images

            def trigger_image_filter(filter_input):
                """Trigger image filtering when input changes."""
                if not filter_input:
                    return "", ""
                
                # Parse the input (format: "img_path|caption")
                parts = filter_input.split("|", 1)
                if len(parts) == 2:
                    img_path, caption = parts
                    return img_path, caption
                return "", ""

            def filter_gallery_by_trigger(img_path, caption, all_images):
                """Filter gallery based on trigger values."""
                if not img_path or not caption or not all_images:
                    return all_images
                
                # Find the selected image
                filtered_images = []
                for stored_img_path, stored_caption in all_images:
                    if stored_caption == caption:
                        filtered_images.append((stored_img_path, stored_caption))
                        break
                
                return filtered_images

            def run_full_parse_with_pages(*args):
                """Run full parse and parse the markdown into pages."""
                result = run_full_parse(*args)
                status_msg, md_content, gallery_items, file_paths, zip_path = result
                
                # Parse markdown into pages
                pages_data = []
                first_page_content = ""
                all_images = []
                if md_content:
                    pages_data = parse_markdown_by_pages(md_content)
                    
                    # Collect all images from all pages
                    for page in pages_data:
                        for line in page['content']:
                            if line.strip().startswith('![') and ('](images/' in line or '](images\\' in line):
                                import re
                                match = re.match(r'!\[([^\]]+)\]\(([^)]+)\)', line.strip())
                                if match:
                                    caption = match.group(1)
                                    img_path = match.group(2)
                                    all_images.append((img_path, caption))
                    
                    
                    # Show only Page 1 content initially
                    if pages_data:
                        first_page = pages_data[0]
                        first_page_content = "\n".join(first_page['content'])
                
                # Prepare first page image immediately and cache page images
                input_pdf_path = args[0]
                first_page_image = None
                saved_paths: List[str] = []
                try:
                    if input_pdf_path:
                        from doctra.utils.pdf_io import render_pdf_to_images
                        tmp_img_dir = Path(tempfile.mkdtemp(prefix="doctra_pages_"))
                        pil_pages = render_pdf_to_images(input_pdf_path)
                        for idx, (im, _, _) in enumerate(pil_pages, start=1):
                            out_path = tmp_img_dir / f"page_{idx:03d}.jpg"
                            im.save(out_path, format="JPEG", quality=90)
                            saved_paths.append(str(out_path))
                        if saved_paths:
                            first_page_image = saved_paths[0]
                except Exception as e:
                    pass

                # Build initial HTML with inline images and proper blocks for first page
                if pages_data:
                    import html as _html, base64, re as _re
                    base_dir = None
                    try:
                        stem = Path(input_pdf_path).stem if input_pdf_path else ""
                        if stem:
                            base_dir = Path("outputs") / stem / "full_parse"
                    except Exception:
                        base_dir = None
                    html_lines: List[str] = []
                    for raw_line in pages_data[0]['content']:
                        line = raw_line.strip()
                        if line.startswith('![') and ('](images/' in line or '](images\\' in line):
                            match = _re.match(r'!\[([^\]]+)\]\(([^)]+)\)', line)
                            if match and base_dir is not None:
                                caption = match.group(1)
                                rel_path = match.group(2).replace('\\\\', '/').replace('\\', '/').lstrip('/')
                                abs_path = (base_dir / rel_path).resolve()
                                try:
                                    with open(abs_path, 'rb') as f:
                                        b64 = base64.b64encode(f.read()).decode('ascii')
                                    html_lines.append(f'<figure><img src="data:image/jpeg;base64,{b64}" alt="{_html.escape(caption)}"/><figcaption>{_html.escape(caption)}</figcaption></figure>')
                                except Exception:
                                    html_lines.append(f'<div>{_html.escape(caption)} (image not found)</div>')
                            else:
                                html_lines.append(f'<p>{_html.escape(raw_line)}</p>')
                        else:
                            if line.startswith('## '):
                                html_lines.append(f'<h3>{_html.escape(line[3:])}</h3>')
                            elif line.startswith('# '):
                                html_lines.append(f'<h2>{_html.escape(line[2:])}</h2>')
                            elif line == '':
                                html_lines.append('<br/>')
                            else:
                                html_lines.append(f'<p>{_html.escape(raw_line)}</p>')
                    first_page_content = "\n".join(html_lines)

                return status_msg, first_page_content, first_page_image, gallery_items, file_paths, zip_path, pages_data, all_images, input_pdf_path, saved_paths

            run_btn.click(
                fn=run_full_parse_with_pages,
                inputs=[pdf, use_vlm, vlm_provider, vlm_api_key, layout_model, dpi, min_score, ocr_lang, ocr_psm, ocr_oem, ocr_config, box_sep],
                outputs=[status, md_preview, page_image, gallery, files_out, zip_out, pages_state, all_images_state, pdf_path_state, page_images_state],
            ).then(
                fn=update_page_selector,
                inputs=[pages_state],
                outputs=[page_selector],
            )

            page_selector.change(
                fn=display_selected_page,
                inputs=[page_selector, pages_state, pdf_path_state, page_images_state],
                outputs=[md_preview, page_image],
            )

            image_filter_input.change(
                fn=trigger_image_filter,
                inputs=[image_filter_input],
                outputs=[current_image_path, current_image_caption],
            ).then(
                fn=filter_gallery_by_trigger,
                inputs=[current_image_path, current_image_caption, all_images_state],
                outputs=[gallery],
            )

        with gr.Tab("Extract Tables/Charts"):
            with gr.Row():
                pdf_e = gr.File(file_types=[".pdf"], label="PDF")
                target = gr.Dropdown(["tables", "charts", "both"], value="both", label="Target")
                use_vlm_e = gr.Checkbox(label="Use VLM (optional)", value=False)
                vlm_provider_e = gr.Dropdown(["gemini", "openai", "anthropic", "openrouter"], value="gemini", label="VLM Provider")
                vlm_api_key_e = gr.Textbox(type="password", label="VLM API Key", placeholder="Optional if VLM disabled")
            with gr.Accordion("Advanced", open=False):
                with gr.Row():
                    layout_model_e = gr.Textbox(value="PP-DocLayout_plus-L", label="Layout model")
                    dpi_e = gr.Slider(100, 400, value=200, step=10, label="DPI")
                    min_score_e = gr.Slider(0, 1, value=0.0, step=0.05, label="Min layout score")

            run_btn_e = gr.Button("▶ Run Extraction", variant="primary")
            status_e = gr.Textbox(label="Status")
            # Dropdown to select specific item
            item_selector_e = gr.Dropdown(label="Select Item", visible=False, interactive=True)
            
            # Display extracted data and images
            with gr.Row():
                tables_preview_e = gr.HTML(label="Extracted Data", elem_classes=["page-content"])
                image_e = gr.Image(label="Selected Image", interactive=False)
            
            # Keep gallery for reference but make it smaller
            gallery_e = gr.Gallery(label="All Extracted Images", columns=4, height=200, preview=True)
            files_out_e = gr.Files(label="Download individual output files")
            zip_out_e = gr.File(label="Download all outputs (ZIP)")

            # State to store output directory
            out_dir_state = gr.State("")
            
            def capture_out_dir(status_text):
                if not status_text:
                    return ""
                try:
                    if "Output directory:" in status_text:
                        return status_text.split("Output directory:", 1)[1].strip()
                except Exception:
                    pass
                return ""
            
            def build_item_selector(out_dir_path, target, use_vlm):
                if not out_dir_path or not use_vlm:
                    return gr.Dropdown(choices=[], value=None, visible=False)
                
                try:
                    out_dir = Path(out_dir_path)
                    mapping = out_dir / "vlm_items.json"
                    if not mapping.exists():
                        return gr.Dropdown(choices=[], value=None, visible=False)
                    
                    import json
                    data = json.loads(mapping.read_text(encoding="utf-8"))
                    choices = []
                    
                    for entry in data:
                        kind = entry.get("kind")
                        # Filter based on target
                        if target == "both" or (target == "tables" and kind == "table") or (target == "charts" and kind == "chart"):
                            title = entry.get("title") or f"{kind.title()}"
                            page = entry.get("page")
                            rel_path = entry.get("image_rel_path")
                            label = f"{title} — Page {page}"
                            choices.append((label, rel_path))
                    
                    return gr.Dropdown(choices=choices, value=choices[0][1] if choices else None, visible=bool(choices))
                except Exception:
                    return gr.Dropdown(choices=[], value=None, visible=False)
            
            def show_selected_item(rel_path, out_dir_path):
                if not rel_path or not out_dir_path:
                    return "", None
                
                try:
                    out_dir = Path(out_dir_path)
                    mapping = out_dir / "vlm_items.json"
                    if not mapping.exists():
                        return "", None
                    
                    import json, html as _html
                    data = json.loads(mapping.read_text(encoding="utf-8"))
                    
                    for entry in data:
                        if entry.get("image_rel_path") == rel_path:
                            headers = entry.get("headers") or []
                            rows = entry.get("rows") or []
                            title = entry.get("title") or "Data"
                            kind = entry.get("kind", "table")
                            
                            # Create HTML table
                            if headers and rows:
                                thead = '<thead><tr>' + ''.join(f'<th>{_html.escape(str(h))}</th>' for h in headers) + '</tr></thead>'
                                tbody = '<tbody>' + ''.join('<tr>' + ''.join(f'<td>{_html.escape(str(c))}</td>' for c in r) + '</tr>' for r in rows) + '</tbody>'
                                html_table = f'<h3>{_html.escape(title)} ({kind.title()})</h3><table class="doc-table">{thead}{tbody}</table>'
                            else:
                                html_table = f'<h3>{_html.escape(title)} ({kind.title()})</h3><p>No structured data available</p>'
                            
                            # Get image path
                            img_abs = str((out_dir / rel_path).resolve())
                            return html_table, img_abs
                    
                    return "", None
                except Exception:
                    return "", None

            run_btn_e.click(
                fn=lambda f, t, a, b, c, d, e, g: run_extract(
                    f.name if f else "",
                    t,
                    a,
                    b,
                    c,
                    d,
                    e,
                    g,
                ),
                inputs=[pdf_e, target, use_vlm_e, vlm_provider_e, vlm_api_key_e, layout_model_e, dpi_e, min_score_e],
                outputs=[status_e, tables_preview_e, gallery_e, files_out_e, zip_out_e],
            ).then(
                fn=capture_out_dir,
                inputs=[status_e],
                outputs=[out_dir_state]
            ).then(
                fn=build_item_selector,
                inputs=[out_dir_state, target, use_vlm_e],
                outputs=[item_selector_e]
            ).then(
                fn=show_selected_item,
                inputs=[item_selector_e, out_dir_state],
                outputs=[tables_preview_e, image_e]
            )
            
            # Handle dropdown selection changes
            item_selector_e.change(
                fn=show_selected_item,
                inputs=[item_selector_e, out_dir_state],
                outputs=[tables_preview_e, image_e]
            )


        gr.Markdown(
            """
<div class="card">
  <b>Tips</b>
  <ul>
    <li>On Spaces, set a secret <code>VLM_API_KEY</code> to enable VLM features.</li>
    <li>Outputs are saved under <code>outputs/&lt;pdf_stem&gt;/</code>.</li>
  </ul>
</div>
            """
        )

    return demo


def launch_ui(server_name: str = "0.0.0.0", server_port: int = 7860, share: bool = False):
    demo = build_demo()
    demo.launch(server_name=server_name, server_port=server_port, share=share)


