"""Diagnostic script for VLM OCR with OpenAI.

Tests whether page.apply_ocr(engine="vlm", model="gpt-4o-mini", client=client)
returns useful content, and falls back to raw OpenAI API calls for comparison.
"""

import base64
import io
import logging
import os
import sys

# Load environment variables from .env
from dotenv import load_dotenv

load_dotenv(os.path.join(os.path.dirname(__file__), "..", ".env"))

api_key = os.environ.get("OPENAI_API_KEY")
if not api_key:
    print("ERROR: OPENAI_API_KEY not found in environment. Check .env file.")
    sys.exit(1)
print(f"API key loaded: {api_key[:8]}...{api_key[-4:]}")

from openai import OpenAI

client = OpenAI()

# Enable debug logging for natural_pdf
import natural_pdf

natural_pdf.configure_logging(level=logging.DEBUG)
logger = logging.getLogger(__name__)

# Also make sure we see vlm_ocr and vlm_client logs
logging.getLogger("natural_pdf.ocr.vlm_ocr").setLevel(logging.DEBUG)
logging.getLogger("natural_pdf.core.vlm_client").setLevel(logging.DEBUG)
logging.getLogger("natural_pdf.core.vlm_prompts").setLevel(logging.DEBUG)

print("\n" + "=" * 70)
print("STEP 1: Load PDF")
print("=" * 70)

pdf = natural_pdf.PDF(
    "https://github.com/jsoma/natural-pdf/raw/refs/heads/main/pdfs/needs-ocr.pdf"
)
page = pdf.pages[0]
print(f"Page loaded: {page.width} x {page.height}")
print(f"Existing text elements: {len(page.find_all('text'))}")
print(f"Existing text content: {repr(page.extract_text()[:200])}")

print("\n" + "=" * 70)
print("STEP 2: Apply VLM OCR with gpt-4o-mini")
print("=" * 70)

try:
    result = page.apply_ocr(engine="vlm", model="gpt-4o-mini", client=client)
    text = page.extract_text()
    elements = page.find_all("text")
    print(f"\nResult type: {type(result)}")
    print(f"Text elements after OCR: {len(elements)}")
    print(f"Extracted text length: {len(text)}")
    print(f"Extracted text (first 500 chars):\n{repr(text[:500])}")

    if elements:
        print("\nFirst 5 elements:")
        for i, el in enumerate(elements[:5]):
            print(
                f"  [{i}] text={repr(el.text[:80] if el.text else None)}, "
                f"bbox=({el.x0:.1f}, {el.top:.1f}, {el.x1:.1f}, {el.bottom:.1f}), "
                f"source={getattr(el, 'source', '?')}"
            )
except Exception as e:
    print(f"ERROR with gpt-4o-mini: {type(e).__name__}: {e}")
    import traceback

    traceback.print_exc()

print("\n" + "=" * 70)
print("STEP 3: Apply VLM OCR with gpt-4o")
print("=" * 70)

try:
    # Remove previous OCR elements first
    page.remove_ocr_elements()

    result = page.apply_ocr(engine="vlm", model="gpt-4o", client=client)
    text = page.extract_text()
    elements = page.find_all("text")
    print(f"\nResult type: {type(result)}")
    print(f"Text elements after OCR: {len(elements)}")
    print(f"Extracted text length: {len(text)}")
    print(f"Extracted text (first 500 chars):\n{repr(text[:500])}")

    if elements:
        print("\nFirst 5 elements:")
        for i, el in enumerate(elements[:5]):
            print(
                f"  [{i}] text={repr(el.text[:80] if el.text else None)}, "
                f"bbox=({el.x0:.1f}, {el.top:.1f}, {el.x1:.1f}, {el.bottom:.1f}), "
                f"source={getattr(el, 'source', '?')}"
            )
except Exception as e:
    print(f"ERROR with gpt-4o: {type(e).__name__}: {e}")
    import traceback

    traceback.print_exc()

print("\n" + "=" * 70)
print("STEP 4: Direct OpenAI API call (bypass natural_pdf)")
print("=" * 70)

try:
    from PIL import Image

    # Render the page to an image
    img = page.render(resolution=144)
    print(f"Rendered image size: {img.size}")

    # Encode as base64 JPEG
    buf = io.BytesIO()
    if img.mode in ("RGBA", "LA", "P"):
        img = img.convert("RGB")
    img.save(buf, format="JPEG", quality=90)
    b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
    data_uri = f"data:image/jpeg;base64,{b64}"
    print(f"Base64 image size: {len(b64)} chars")

    # Build the same prompt that natural_pdf uses for OpenAI
    from natural_pdf.core.vlm_prompts import OPENAI_OCR_PROMPT

    print(f"\nPrompt being used:\n{OPENAI_OCR_PROMPT}")

    # Call OpenAI directly
    print("\nCalling gpt-4o-mini directly...")
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": data_uri}},
                    {"type": "text", "text": OPENAI_OCR_PROMPT},
                ],
            }
        ],
        max_completion_tokens=4096,
        response_format={"type": "json_object"},
    )
    print(f"\nFull response object:")
    print(f"  finish_reason: {response.choices[0].finish_reason}")
    print(f"  usage: {response.usage}")
    refusal = getattr(response.choices[0].message, "refusal", None)
    print(f"  refusal: {refusal}")
    content = response.choices[0].message.content
    print(f"  content length: {len(content) if content else 0}")
    print(f"  content (first 1000 chars):\n{content[:1000] if content else '(None)'}")

    # Also try WITHOUT json_object mode
    print("\n\nCalling gpt-4o-mini WITHOUT json_object mode...")
    response2 = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": data_uri}},
                    {"type": "text", "text": OPENAI_OCR_PROMPT},
                ],
            }
        ],
        max_completion_tokens=4096,
    )
    content2 = response2.choices[0].message.content
    print(f"  finish_reason: {response2.choices[0].finish_reason}")
    print(f"  content length: {len(content2) if content2 else 0}")
    print(f"  content (first 1000 chars):\n{content2[:1000] if content2 else '(None)'}")

    # Try with gpt-4o as well
    print("\n\nCalling gpt-4o directly...")
    response3 = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": data_uri}},
                    {"type": "text", "text": OPENAI_OCR_PROMPT},
                ],
            }
        ],
        max_completion_tokens=4096,
        response_format={"type": "json_object"},
    )
    content3 = response3.choices[0].message.content
    print(f"  finish_reason: {response3.choices[0].finish_reason}")
    print(f"  content length: {len(content3) if content3 else 0}")
    print(f"  content (first 1000 chars):\n{content3[:1000] if content3 else '(None)'}")

    # Parse the responses to see if natural_pdf's parser would work
    print("\n\n--- Testing natural_pdf parser on raw responses ---")
    from natural_pdf.ocr.vlm_ocr import parse_grounding_response

    for label, raw in [
        ("gpt-4o-mini (json_object)", content),
        ("gpt-4o-mini (no format)", content2),
        ("gpt-4o (json_object)", content3),
    ]:
        if raw:
            parsed = parse_grounding_response(raw, family="openai")
            print(f"\n  {label}: {len(parsed)} results parsed")
            if parsed:
                for j, p in enumerate(parsed[:3]):
                    print(f"    [{j}] bbox={p['bbox']}, text={repr(p['text'][:60])}")
            else:
                print(f"    Raw (first 300): {raw[:300]}")
        else:
            print(f"\n  {label}: content was None/empty")

except Exception as e:
    print(f"ERROR in direct API call: {type(e).__name__}: {e}")
    import traceback

    traceback.print_exc()

print("\n" + "=" * 70)
print("DONE")
print("=" * 70)

pdf.close()
