"""A/B test prompts page recueil — champs qui échouent en V1.

Cible : dp_libelle, praticien_conseil, codage_reco.dp, ghs_reco.
Ground truth = JSON legacy (sauf quand on l'a corrigé manuellement).
"""
import json
import time
from pathlib import Path
from pipeline.ocr_glm import GLMOCR
from pipeline.ingest import pdf_to_images
from pipeline.extract import parse_json_output


# ======== Prompts à comparer ========

PROMPT_V1 = """Lis la fiche médicale OGC et renvoie STRICTEMENT le JSON suivant, sans commentaire ni markdown.
Si un champ est illisible, laisse une chaîne vide.

{
  "codage_etab":  {"dp": "", "dp_libelle": "", "dr": "", "das": [{"code":"","position":"","libelle":""}]},
  "codage_reco":  {"dp": "", "dr": "", "das": [{"code":"","position":""}]},
  "ghm_etab": "", "ghs_etab": "",
  "ghm_reco": "", "ghs_reco": "",
  "praticien_conseil": ""
}"""

PROMPT_V2 = """Lis cette fiche médicale OGC (contrôle T2A). Renvoie STRICTEMENT le JSON suivant, sans commentaire ni markdown.

RÈGLES IMPORTANTES :
- Le tableau "Codage de l'Établissement / Recodage" a DEUX colonnes distinctes : les codes "Recodage" sont dans la colonne la plus à DROITE, séparés des codes "Établissement" (à gauche). Ne recopie JAMAIS les codes Établissement dans Recodage.
- "dp_libelle" = texte descriptif à droite du code DP (ex: "HEMORR. ET HEMATOME COMPLIQ. UN ACTE, NCA").
- "ghm_etab", "ghs_etab", "ghm_reco", "ghs_reco" sont sur UNE SEULE ligne en bas, dans cet ordre (4 valeurs).
- "praticien_conseil" = nom manuscrit (DR + nom) en bas de page sous "Nom du praticien conseil responsable du codage".
- Si un champ est illisible, laisse une chaîne vide. Ne devine pas.

{
  "codage_etab":  {"dp": "", "dp_libelle": "", "dr": "", "das": [{"code":"","position":"","libelle":""}]},
  "codage_reco":  {"dp": "", "dr": "", "das": [{"code":"","position":""}]},
  "ghm_etab": "", "ghs_etab": "",
  "ghm_reco": "", "ghs_reco": "",
  "praticien_conseil": ""
}"""


CASES = [
    ("2018 CARC/OGC 7.pdf",  1),
    ("2018 CARC/OGC 27.pdf", 1),
    ("2018 CARC/OGC 55.pdf", 1),
    ("2018 CARC/OGC 86.pdf", 1),
]


def get(d, path, default=""):
    for k in path.split("."):
        d = d.get(k, {}) if isinstance(d, dict) else default
    return d if d else default


def compare_fields(label, extracted, legacy):
    fields = {
        "codage_etab.dp":         ("codage_etab.dp",         "codage_etab.dp"),
        "codage_etab.dp_libelle": ("codage_etab.dp_libelle", "codage_etab.dp_libelle"),
        "codage_reco.dp":         ("codage_reco.dp",         "codage_reco.dp"),
        "ghm_etab":               ("ghm_etab",               "ghm_etab"),
        "ghs_etab":               ("ghs_etab",               "ghs_etab"),
        "ghm_reco":               ("ghm_reco",               "ghm_reco"),
        "ghs_reco":               ("ghs_reco",               "ghs_reco"),
        "praticien_conseil":      ("praticien_conseil",      "praticien_conseil"),
    }
    print(f"    --- {label} vs legacy ---")
    for f, (pe, pl) in fields.items():
        v_ext = str(get(extracted or {}, pe)).strip()
        v_leg = str(get(legacy, pl)).strip()
        # Comparaison tolérante sur dp_libelle (tronqué dans legacy)
        if f == "codage_etab.dp_libelle":
            match = v_leg in v_ext or v_ext in v_leg if (v_ext and v_leg) else (v_ext == v_leg)
        else:
            match = v_ext == v_leg
        mark = "✓" if match else ("∅" if not v_ext and not v_leg else "✗")
        print(f"      {mark} {f:26s} ext={v_ext!r:45s} leg={v_leg!r}")


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

    for pdf, page in CASES:
        name = Path(pdf).stem
        images = pdf_to_images(pdf)
        img = images[page - 1]
        with open(f"output/{name}.json") as f:
            legacy = json.load(f)["recueil"]["parsed"]

        print(f"=========== {name} ===========")
        for label, prompt in [("V1 (actuel)", PROMPT_V1), ("V2 (précisé)", PROMPT_V2)]:
            t0 = time.time()
            res = ocr.run(img, prompt, max_new_tokens=2048)
            parsed = parse_json_output(res["text"])
            print(f"  [{label}] ({time.time()-t0:.1f}s)")
            compare_fields(label, parsed, legacy)
        print()


if __name__ == "__main__":
    main()
