# pyright: ignore
#!/usr/bin/env python3
"""
Test VLM pipeline with full PDF (not extracted pages).
"""

from pathlib import Path
from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import VlmPipelineOptions
from docling.datamodel.vlm_model_specs import GRANITEDOCLING_TRANSFORMERS
from docling.datamodel.accelerator_options import AcceleratorDevice
from docling.pipeline.vlm_pipeline import VlmPipeline

# Use FULL PDF, not extracted page
pdf_path = Path("documents/TMUS_992_Q423.pdf")

print(f"Testing VLM on FULL PDF: {pdf_path}")
print("=" * 50)

# VLM pipeline configuration (matching working version)
vlm_options = VlmPipelineOptions(vlm_options=GRANITEDOCLING_TRANSFORMERS)
vlm_options.generate_page_images = True
vlm_options.accelerator_options.device = AcceleratorDevice.CUDA

vlm_format_options = PdfFormatOption(
    pipeline_cls=VlmPipeline,
    pipeline_options=vlm_options
)

converter = DocumentConverter(
    format_options={InputFormat.PDF: vlm_format_options}
)

print("Converting FULL PDF with VLM (this may take a while)...")
print("Processing only page 17 for comparison...")

# Convert full PDF
result = converter.convert(str(pdf_path))

print("
Conversion complete!")
print(f"Document has {len(result.document.pages)} pages")
print(f"Document has {len(result.document.tables)} tables")

# Check page 17 specifically (index 16)
if len(result.document.pages) >= 17:
    print("
--- Page 17 Analysis ---")

    # Check if page 17 has tables
    page_17_tables = [
        table for table in result.document.tables
        if table.prov and len(table.prov) > 0 and table.prov[0].page_no == 17
    ]

    print(f"Tables on page 17: {len(page_17_tables)}")

    # Export just a portion to see structure
    markdown = result.document.export_to_markdown()

    # Find the Financial & Operational Tables section
    if "Financial" in markdown and "Operational" in markdown:
        start_idx = markdown.find("Financial")
        if start_idx > 0:
            excerpt = markdown[start_idx:start_idx+2000]
            print("
--- Page 17 Markdown Excerpt ---")
            print(excerpt)
            print("
✓ SUCCESS: VLM correctly processed page 17 in full document!")
    else:
        print("
✗ ERROR: Could not find Financial & Operational Tables section")
else:
    print(f"
✗ ERROR: Document only has {len(result.document.pages)} pages")