"""Test : cropper la zone checkbox et demander à GLM-OCR."""
from pathlib import Path
from PIL import Image
from pipeline.ocr_glm import GLMOCR
from pipeline.ingest import pdf_to_images

# Zone checkbox sur la fiche recueil : bas à droite, ~60-85% largeur, 82-88% hauteur
ZONE = (0.55, 0.82, 0.92, 0.90)

CASES = [
    ("2018 CARC/OGC 7.pdf",  1, "accord"),
    ("2018 CARC/OGC 55.pdf", 1, "désaccord"),
    ("2018 CARC/OGC 27.pdf", 1, "désaccord"),
    ("2018 CARC/OGC 86.pdf", 1, "désaccord"),  # ajout pour confirmer
]

PROMPTS = {
    "P1_naturel": """Sur cette image, deux cases à cocher : "Accord" et "Désaccord". Quelle case contient une croix ou est remplie ? Réponds par UN SEUL mot : accord OU désaccord.""",
    "P2_json":    """Retourne UNIQUEMENT ce JSON :\n{"case_cochee": "accord" ou "désaccord"}\nRegarde les deux cases à cocher sur l'image et identifie celle qui est cochée (X, V ou noire).""",
}


def crop_rel(img: Image.Image, z):
    w, h = img.size
    return img.crop((int(z[0]*w), int(z[1]*h), int(z[2]*w), int(z[3]*h)))


def main():
    ocr = GLMOCR()
    print(f"VRAM = {ocr.vram_gb:.2f} Go\n")

    Path("/tmp/ogc_crops").mkdir(exist_ok=True)
    for pdf, page, expected in CASES:
        images = pdf_to_images(pdf)
        img = Image.open(images[page - 1])
        crop = crop_rel(img, ZONE)
        crop_path = f"/tmp/ogc_crops/{Path(pdf).stem.replace(' ', '_')}_cb.png"
        crop.save(crop_path)
        print(f"=== {Path(pdf).stem} (attendu: {expected})   crop={crop.size} ===")
        for name, prompt in PROMPTS.items():
            res = ocr.run(crop_path, prompt, max_new_tokens=64)
            print(f"  [{name}] → {res['text'].strip()[:120]}")
        print()


if __name__ == "__main__":
    main()
