"""Parsing JSON tolérant pour les sorties des VLM.

Les VLM (Qwen, GLM-OCR, GOT-OCR…) produisent du JSON avec des anomalies
fréquentes :

- Encadrement par des fences markdown ```json ... ```
- Virgules manquantes entre objets ou éléments de tableau
- Boucles pathologiques d'objets vides `{"code":"","position":""}` répétés
  jusqu'à `max_new_tokens`, ce qui tronque le JSON sans le fermer proprement

Ce module expose `parse_json_output()` qui applique plusieurs stratégies de
récupération avant d'abandonner. En dernier recours, il renvoie un dict avec
`_raw` + `_parse_error` pour audit, jamais `None` (ce qui casserait le pipeline).
"""
from __future__ import annotations

import json
import re


# Pattern qui matche un objet vide générique {"code":"","position":"",...}
_EMPTY_OBJ_PATTERN = re.compile(
    r'\{\s*"code"\s*:\s*""\s*,\s*"position"\s*:\s*""\s*(?:,\s*"libelle"\s*:\s*""\s*)?\}',
    re.DOTALL,
)

_FENCE_OPEN_RE = re.compile(r"^```(?:json)?\s*")
_FENCE_CLOSE_RE = re.compile(r"\s*```$")
_MISSING_COMMA_OBJ_RE = re.compile(r"\}\s*\n(\s*\{)")
_MISSING_COMMA_ARR_RE = re.compile(r"\]\s*\n(\s*\[)")


def strip_fences(text: str) -> str:
    """Retire un éventuel encadrement ```json ... ```."""
    text = _FENCE_OPEN_RE.sub("", text.strip())
    text = _FENCE_CLOSE_RE.sub("", text)
    return text


def patch_missing_commas(text: str) -> str:
    """Ajoute les virgules manquantes entre `}\\n{` et `]\\n[`.

    Les VLM omettent fréquemment ces virgules dans leurs sorties JSON.
    """
    text = _MISSING_COMMA_OBJ_RE.sub(r"},\n\1", text)
    text = _MISSING_COMMA_ARR_RE.sub(r"],\n\1", text)
    return text


def truncate_empty_loop(text: str, max_consecutive: int = 2) -> str:
    """Tronque les boucles d'objets vides (`{"code":"","position":""}` répété).

    Cas d'usage : quand un tableau DAS ou Actes est vide dans l'image, le
    VLM a parfois tendance à générer le même objet vide en boucle jusqu'à
    saturer `max_new_tokens`. La sortie est alors tronquée sans fermer le
    JSON → parse error. Garder au maximum `max_consecutive` occurrences.
    """
    matches = list(_EMPTY_OBJ_PATTERN.finditer(text))
    if len(matches) <= max_consecutive:
        return text
    cut_at = matches[max_consecutive - 1].end()
    return text[:cut_at]


def close_open_json(text: str) -> str:
    """Ajoute les brackets/braces manquants pour fermer un JSON tronqué.

    Compte les accolades et crochets non-balancés en ignorant ceux à
    l'intérieur d'une chaîne, puis ferme dans le bon ordre (tableaux
    ouverts d'abord, puis objets).
    """
    depth_brace = 0
    depth_bracket = 0
    in_string = False
    escape = False
    for c in text:
        if escape:
            escape = False
            continue
        if c == "\\":
            escape = True
            continue
        if c == '"':
            in_string = not in_string
            continue
        if in_string:
            continue
        if c == "{":
            depth_brace += 1
        elif c == "}":
            depth_brace -= 1
        elif c == "[":
            depth_bracket += 1
        elif c == "]":
            depth_bracket -= 1
    closed = text.rstrip().rstrip(",")
    closed += "]" * max(0, depth_bracket)
    closed += "}" * max(0, depth_brace)
    return closed


def parse_json_output(raw: str) -> dict | None:
    """Parse une sortie VLM en dict. Applique plusieurs stratégies :

    1. strip des fences markdown ```json
    2. parse direct
    3. patch des virgules manquantes
    4. troncature des boucles d'objets vides + fermeture du JSON

    En cas d'échec de toutes les stratégies, retourne
    `{"_raw": raw, "_parse_error": str}` pour permettre l'audit manuel
    plutôt que de casser le pipeline.
    """
    if not raw:
        return None
    text = strip_fences(raw)
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    patched = patch_missing_commas(text)
    try:
        return json.loads(patched)
    except json.JSONDecodeError:
        pass

    trimmed = truncate_empty_loop(patched)
    closed = close_open_json(trimmed)
    try:
        result = json.loads(closed)
        if isinstance(result, dict):
            result["_truncated_loop"] = True
        return result
    except json.JSONDecodeError as e:
        return {"_raw": raw, "_parse_error": str(e)}