"""
prepare_doctr_dataset.py
-------------------------
Convertit dataset_lines/ vers le format doctr :
  dataset_doctr/
    train/images/ + labels.json
    val/images/   + labels.json

Copie de référence pour la reproductibilité (documentation technique) —
attend l'arborescence annotator/ du dépôt d'origine, non incluse ici.

Doit être lancé depuis le dossier train/ du dépôt d'origine :
  cd train/
  python prepare_doctr_dataset.py
"""

import os
import json
import shutil
from pathlib import Path

# Chemins depuis train/ → annotator/
SOURCE_DIR = os.path.join("..", "annotator", "dataset_lines")
TARGET_DIR = os.path.join("..", "annotator", "dataset_doctr")


def convert_split(split: str) -> int:
    src_images = os.path.join(SOURCE_DIR, split, "images")
    src_labels = os.path.join(SOURCE_DIR, split, "labels")
    dst_images = os.path.join(TARGET_DIR, split, "images")
    os.makedirs(dst_images, exist_ok=True)

    if not os.path.isdir(src_images):
        print(f"  ⚠️  Introuvable : {src_images}")
        return 0

    labels_dict = {}
    n_ok = n_skip = 0

    for img_file in sorted(os.listdir(src_images)):
        if not img_file.lower().endswith((".jpg", ".jpeg", ".png")):
            continue
        stem  = Path(img_file).stem
        lbl_f = os.path.join(src_labels, f"{stem}.txt")
        if not os.path.exists(lbl_f):
            n_skip += 1
            continue
        text = open(lbl_f, encoding="utf-8").read().strip()
        if not text:
            n_skip += 1
            continue
        shutil.copy2(os.path.join(src_images, img_file),
                     os.path.join(dst_images, img_file))
        labels_dict[img_file] = text
        n_ok += 1

    out = os.path.join(TARGET_DIR, split, "labels.json")
    with open(out, "w", encoding="utf-8") as f:
        json.dump(labels_dict, f, indent=2, ensure_ascii=False)

    print(f"  {split:6s} : {n_ok} images, {n_skip} ignorées → {out}")
    return n_ok


def show_vocab():
    chars = set()
    for split in ("train", "val"):
        p = os.path.join(TARGET_DIR, split, "labels.json")
        if os.path.exists(p):
            for txt in json.load(open(p)).values():
                chars.update(txt)
    print(f"\n  Vocabulaire ({len(chars)} chars) : {''.join(sorted(chars))}")
    print(f"  → Le vocab doctr 'french' couvre tous ces caractères ✅")


if __name__ == "__main__":
    if not os.path.isdir(SOURCE_DIR):
        print(f"❌ Introuvable : {SOURCE_DIR}")
        print(f"   Lance d'abord depuis annotator/ : python split_lines.py")
        raise SystemExit(1)

    print(f"🔄 Conversion {SOURCE_DIR} → {TARGET_DIR}\n")
    n_train = convert_split("train")
    n_val   = convert_split("val")
    show_vocab()
    print(f"\n✅ {n_train} train + {n_val} val prêts pour le fine-tuning")
    print(f"\n👉 Prochaine étape :")
    print(f"   pip install python-doctr[torch]")
    print(f"   python finetune_doctr.py")