print("Extract Text from PDF Documents with IONOS AI Model Hub")
print("========================================================\n")


print("Prerequisite: Access API Token from environment variable")
print("--------------------------------------------------------")

from dotenv import load_dotenv
import os

load_dotenv()
IONOS_API_TOKEN = os.getenv("IONOS_API_TOKEN")

print("IONOS_API_TOKEN: ", IONOS_API_TOKEN[:10], "\n")


print("Setup: Initialize client and PDF rendering helper")
print("-------------------------------------------------")

import base64
import io
import time

import pypdfium2 as pdfium
from openai import OpenAI

client = OpenAI(
    api_key=IONOS_API_TOKEN,
    base_url="https://openai.inference.de-txl.ionos.com/v1",
)

LIGHTON_MODEL = "lightonai/LightOnOCR-2-1B"
MISTRAL_MODEL = "mistralai/Mistral-Small-24B-Instruct"

# Replace with the path to your local PDF file
pdf_path = "document.pdf"


def pdf_page_to_base64(pdf_path: str, page_index: int, scale: float = 2.0) -> str:
    doc = pdfium.PdfDocument(pdf_path)
    bitmap = doc[page_index].render(scale=scale)
    buf = io.BytesIO()
    bitmap.to_pil().save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode()


from pathlib import Path

if not Path(pdf_path).exists():
    print(f"PDF file '{pdf_path}' not found.")
    print("Place a PDF at this path and re-run the script.\n")
    raise SystemExit(1)


print("Step 1: Extract text from page 1 using LightOnOCR-2-1B")
print("-------------------------------------------------------")

image_b64 = pdf_page_to_base64(pdf_path, page_index=0)

start = time.time()
response = client.chat.completions.create(
    model=LIGHTON_MODEL,
    messages=[
        {
            "role": "user",
            "content": [
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/png;base64,{image_b64}"},
                }
            ],
        }
    ],
    max_tokens=4096,
    temperature=0.0,
)
print(f"Response time: {time.time() - start:.2f}s")
print(f"Extracted text:\n{response.choices[0].message.content}\n")


print("Step 2: Extract text from page 1 using Mistral Small 24B")
print("---------------------------------------------------------")

start = time.time()
response = client.chat.completions.create(
    model=MISTRAL_MODEL,
    messages=[
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": "Extract all text from this document page and preserve the structure.",
                },
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/png;base64,{image_b64}"},
                },
            ],
        }
    ],
    max_tokens=4096,
    temperature=0.0,
)
print(f"Response time: {time.time() - start:.2f}s")
print(f"Extracted text:\n{response.choices[0].message.content}\n")


print("Step 3: Process all pages")
print("-------------------------")

doc = pdfium.PdfDocument(pdf_path)
total_pages = len(doc)
print(f"Document has {total_pages} page(s). Processing with LightOnOCR-2-1B...\n")

parts = []
for i in range(total_pages):
    img_b64 = pdf_page_to_base64(pdf_path, i)
    resp = client.chat.completions.create(
        model=LIGHTON_MODEL,
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/png;base64,{img_b64}"},
                    }
                ],
            }
        ],
        max_tokens=4096,
        temperature=0.0,
    )
    parts.append(f"--- Page {i + 1} ---\n\n{resp.choices[0].message.content}")
    print(f"Page {i + 1}/{total_pages} done.")

print("\nFull document output:")
print("\n\n".join(parts))