"""Download official model weights and verify them on a synthetic page."""

from __future__ import annotations

import argparse
import os
from pathlib import Path
import tempfile

from PIL import Image, ImageDraw


ROOT = Path(os.environ.get("OCR_MODEL_HOME", Path(__file__).resolve().parents[1] / ".cache")).expanduser().resolve()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--classic", action="store_true")
    parser.add_argument("--vision", action="store_true")
    parser.add_argument("--device", default="cpu", help="Vision device, for example cpu or gpu:0")
    args = parser.parse_args()
    if not args.classic and not args.vision:
        parser.error("choose --classic and/or --vision")

    ROOT.mkdir(parents=True, exist_ok=True)
    os.environ["PADDLE_PDX_CACHE_HOME"] = str(ROOT)
    os.environ["HF_HOME"] = str(ROOT / "hf")
    os.environ.pop("HF_HUB_OFFLINE", None)
    os.environ.pop("TRANSFORMERS_OFFLINE", None)
    with tempfile.TemporaryDirectory(prefix="ocr-synthetic-setup-") as temporary:
        image = Path(temporary) / "synthetic.png"
        page = Image.new("RGB", (1000, 600), "white")
        draw = ImageDraw.Draw(page)
        draw.text((50, 50), "SYNTHETIC TEST / NO PERSONAL DATA", fill="black")
        draw.text((50, 130), "D.3 SAMPLE 123", fill="black")
        page.save(image)
        if args.classic:
            from paddleocr import PaddleOCR
            engine = PaddleOCR(
                text_detection_model_name="PP-OCRv5_mobile_det",
                text_recognition_model_name="latin_PP-OCRv5_mobile_rec",
                use_doc_orientation_classify=False,
                use_doc_unwarping=False,
                use_textline_orientation=False,
                device="cpu",
            )
            list(engine.predict(str(image)))
        if args.vision:
            from paddleocr import PaddleOCRVL
            engine = PaddleOCRVL(
                pipeline_version="v1.6", use_layout_detection=True,
                use_doc_orientation_classify=False, use_doc_unwarping=False,
                device=args.device,
            )
            list(engine.predict(str(image)))

    models = ROOT / "official_models"
    checks = []
    if args.classic:
        checks.extend((models / name / "inference.pdiparams" for name in
                       ("PP-OCRv5_mobile_det", "latin_PP-OCRv5_mobile_rec")))
    if args.vision:
        checks.extend((models / "PP-DocLayoutV3" / "inference.pdiparams",
                       models / "PaddleOCR-VL-1.6" / "model.safetensors"))
    if any(not path.is_file() for path in checks):
        raise SystemExit("Official model download did not create all expected files")
    print("Official models ready; only synthetic pixels were used")


if __name__ == "__main__":
    main()