# pyright: ignore
#!/usr/bin/env python3
"""
Test fixed VLM pipeline configuration.
"""

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

# Test file
pdf_path = Path("output/debug_pages/TMUS_992_Q423_page_017.pdf")

print(f"Testing FIXED VLM pipeline on: {pdf_path}")
print("=" * 50)

# VLM pipeline configuration (matching working version)
vlm_options = VlmPipelineOptions(vlm_options=GRANITEDOCLING_TRANSFORMERS)
vlm_options.generate_page_images = True  # Enable efficient page batching
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 with FIXED VLM pipeline...")
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")

# Export to markdown
markdown = result.document.export_to_markdown()
print(f"
Markdown length: {len(markdown)} characters")

# Save output
output_path = Path("test_vlm_fixed_output.md")
with open(output_path, "w") as f:
    f.write(markdown)
print(f"Output saved to: {output_path}")

# Show preview
print("
First 2000 chars of markdown:")
print(markdown[:2000])

# Check if tables were extracted
if result.document.tables:
    print(f"
✓ SUCCESS: Found {len(result.document.tables)} tables!")
else:
    print("
✗ ERROR: No tables found")