json_utils.py

"""
json_utils.py — Extração cirúrgica e recuperação de JSON malformado
════════════════════════════════════════════════════════════════════

Arquitetura em cascata:
  1. Parse direto                     (happy path — O(1))
  2. Sanitização textual              (patterns conhecidos de malformação)
  3. Extração key-by-key              (balanceamento de profundidade)
  4. Inject defaults                  (campos required ausentes → defaults tipados)

Caso de uso primário: recuperar `failed_generation` de erros 400
`json_validate_failed` do provider (OpenAI-compat), onde o LLM produziu
JSON válido mas faltando um campo obrigatório no schema.

Integração em external_worker.py → TrindadePipeline.api_call():

    elif result.code == 400:
        recovered = JsonUtils.repair_failed_json(
            result.raw_error_body,          # body bruto do HTTP 400
            schema   = schema_step,
            wrapper  = pipeline.wrapper_key,
        )
        if recovered:
            logger.warning(f"♻️  [{label}] JSON recuperado via repair_failed_json")
            return recovered, None, False   # trata como END_SUCCESS
        logger.error(f"Erro 400 irrecuperável: {result.content}")
        return None, None, False

Nota: para que `result.raw_error_body` esteja disponível, providers.py
precisa salvá-lo em APIResult.raw_error_body ao receber status 400.
"""

from __future__ import annotations

import json
import logging
import re
from typing import Any, Dict, List, Optional, Tuple

logger = logging.getLogger("json_utils")


# ═══════════════════════════════════════════════════════════════════════════════
#  CONSTANTES
# ═══════════════════════════════════════════════════════════════════════════════

# Defaults por tipo JSON Schema
_TYPE_DEFAULTS: Dict[str, Any] = {
    "array":   [],
    "object":  {},
    "string":  "",
    "integer": 0,
    "number":  0.0,
    "boolean": False,
    "null":    None,
}

# Patterns de sanitização em ordem de aplicação.
# Cada entrada: (pattern_regex, replacement, flags)
_SANITIZE_PATTERNS: List[Tuple[str, str, int]] = [
    # "campo",":",valor  →  "campo": valor   (vírgula antes dos dois-pontos)
    (r'"(\w+)"\s*,\s*":"\s*,\s*', r'"\1": ', 0),
    # "campo",": valor   →  "campo": valor   (vírgula antes do separador sem espaço)
    (r'"(\w+)"\s*,\s*":\s*', r'"\1": ', 0),
    # vírgulas duplicadas
    (r',\s*,+', ',', 0),
    # trailing comma antes de ] ou }
    (r',\s*(?=[\]\}])', '', 0),
    # colchete duplo no fechamento de array  ]]  →  ]
    # (só quando não é início de novo array)
    (r'\]\s*\](?=\s*[,\}\]])', ']', 0),
    # colchete duplo seguido de }  ]]}  →  ]}
    (r'\]\s*\]\s*\}', ']}', 0),
]


# ═══════════════════════════════════════════════════════════════════════════════
#  JsonUtils
# ═══════════════════════════════════════════════════════════════════════════════

class JsonUtils:
    """
    Utilitários de extração, sanitização e recuperação de JSON malformado.

    Todos os métodos são @classmethod — sem estado de instância.
    Thread-safe (apenas operações puras em strings/dicts).
    """

    # ──────────────────────────────────────────────────────────────────────────
    #  1. ENTRY POINT: pipeline cascata completo
    # ──────────────────────────────────────────────────────────────────────────

    @classmethod
    def extract_json(
        cls,
        text: str,
        required_keys: Optional[List[str]] = None,
    ) -> Optional[Dict]:
        """
        Tenta extrair um dict JSON de `text` em cascata:
          1. json.loads direto
          2. Isola o bloco { ... } externo e parseia
          3. Sanitiza e repete 1-2
          4. Extração key-by-key (required_keys)

        Retorna None se todas as camadas falharem.
        """
        # ── camada 1: parse direto ─────────────────────────────────────────
        result = cls._try_parse(text)
        if result is not None:
            return result

        # ── camada 2: isola bloco externo { } ─────────────────────────────
        block = cls._isolate_outer_block(text)
        if block and block != text:
            result = cls._try_parse(block)
            if result is not None:
                return result

        # ── camada 3: sanitiza → tenta de novo ────────────────────────────
        sanitized = cls.sanitize_json(block or text)
        result = cls._try_parse(sanitized)
        if result is not None:
            return result

        # isola de novo após sanitização
        block2 = cls._isolate_outer_block(sanitized)
        if block2 and block2 != sanitized:
            result = cls._try_parse(block2)
            if result is not None:
                return result

        # ── camada 4: key-by-key (requer required_keys) ───────────────────
        #if required_keys:
        #    assembled = cls._extract_by_keys(sanitized or text, required_keys)
        #    if assembled:
        #        return assembled

        logger.debug("extract_json: todas as camadas falharam")
        return None

    # ──────────────────────────────────────────────────────────────────────────
    #  2. SANITIZAÇÃO TEXTUAL
    # ──────────────────────────────────────────────────────────────────────────

    @classmethod
    def sanitize_json(cls, text: str) -> str:
        """
        Aplica patterns conhecidos de malformação em sequência.
        Não lança exceções — pior caso retorna o texto original.

        Patterns tratados:
          • "campo",":",valor   →  "campo": valor
          • ,,                  →  ,
          • ,]  ,}              →  ] }
          • ]]                  →  ]  (array duplo no fechamento)
        """
        result = text
        for pattern, replacement, flags in _SANITIZE_PATTERNS:
            try:
                new = re.sub(pattern, replacement, result, flags=flags)
                if new != result:
                    logger.debug(f"sanitize_json: aplicou pattern {pattern!r}")
                result = new
            except re.error as e:
                logger.warning(f"sanitize_json: regex error em {pattern!r}: {e}")
        return result

    # ──────────────────────────────────────────────────────────────────────────
    #  3. EXTRAÇÃO DE KEY INDIVIDUAL (balanceamento de profundidade)
    # ──────────────────────────────────────────────────────────────────────────

    @classmethod
    def extract_key(cls, text: str, key: str) -> Optional[str]:
        """
        Localiza `"key":` no texto e extrai o valor completo usando
        balanceamento de profundidade para {} e [].

        Suporta valores: objeto {}, array [], string "", número, bool, null.
        Retorna a substring do valor (sem key), ou None se não encontrar.
        """
        pattern = re.compile(r'"' + re.escape(key) + r'"\s*:\s*')
        m = pattern.search(text)
        if not m:
            return None

        start = m.end()
        return cls._extract_value_at(text, start)

    # ──────────────────────────────────────────────────────────────────────────
    #  4. INJEÇÃO DE DEFAULTS (campos required ausentes)
    # ──────────────────────────────────────────────────────────────────────────

    @classmethod
    def inject_defaults(
        cls,
        data:        Dict,
        schema:      Optional[Dict] = None,
        wrapper_key: str            = "manifestacao_juridica",
    ) -> Dict:
        """
        Percorre o schema recursivamente e injeta valores padrão para campos
        `required` ausentes em `data`.

        Se `schema` for None, retorna `data` sem modificação.

        Estratégia:
          • type=array   → []
          • type=object  → {}  (recursivo com subschema)
          • type=string  → ""
          • type=integer → 0
          • type=number  → 0.0
          • type=boolean → false
          • type=null    → null
          • items com required (array de objetos) → injeta em cada elemento
        """
        if schema is None:
            return data

        result = json.loads(json.dumps(data))  # deep copy via JSON
        cls._inject_recursive(result, schema)
        return result

    # ──────────────────────────────────────────────────────────────────────────
    #  5. RECUPERAÇÃO DE `failed_generation` (erro 400 json_validate_failed)
    # ──────────────────────────────────────────────────────────────────────────

    @classmethod
    def repair_failed_json(
        cls,
        raw_error_body: Optional[str],
        schema:         Optional[Dict] = None,
        wrapper_key:    str            = "manifestacao_juridica",
    ) -> Optional[str]:
        """
        Dado o body bruto de um HTTP 400 com code=json_validate_failed,
        extrai o campo `failed_generation`, tenta parseá-lo, injeta
        defaults para campos required ausentes e retorna JSON string
        pronto para `json.loads()`.

        Fluxo:
          1. Parseia o envelope de erro
          2. Confirma code == "json_validate_failed"
          3. Extrai failed_generation
          4. Tenta parse direto (happy path — quase sempre válido)
          5. Fallback: extract_json com sanitização
          6. inject_defaults com schema
          7. Retorna json.dumps(resultado)

        Retorna None se não conseguir recuperar.
        """
        if not raw_error_body:
            return None

        # ── parseia envelope de erro ───────────────────────────────────────
        try:
            envelope = json.loads(raw_error_body)
        except json.JSONDecodeError:
            # tenta extrair JSON de dentro do texto
            envelope = cls.extract_json(raw_error_body)
            if envelope is None:
                logger.warning("repair_failed_json: envelope de erro não é JSON válido")
                return None

        # ── navega até o campo de erro ─────────────────────────────────────
        error_obj = envelope.get("error", envelope)
        if not isinstance(error_obj, dict):
            logger.warning("repair_failed_json: campo 'error' ausente ou inválido")
            return None

        code = error_obj.get("code", "")
        if code != "json_validate_failed":
            logger.debug(f"repair_failed_json: code={code!r} (não é json_validate_failed)")
            return None

        failed_gen: Optional[str] = error_obj.get("failed_generation")
        if not failed_gen:
            logger.warning("repair_failed_json: failed_generation ausente")
            return None

        logger.info(f"repair_failed_json: failed_generation encontrado ({len(failed_gen)} chars)")

        # ── tenta parse direto ─────────────────────────────────────────────
        recovered: Optional[Dict] = cls._try_parse(failed_gen)
        if recovered is None:
            logger.info("repair_failed_json: parse direto falhou — tentando extract_json")
            recovered = cls.extract_json(failed_gen)

        if recovered is None:
            logger.error("repair_failed_json: não foi possível parsear failed_generation")
            return None

        # ── injeta defaults para campos required ausentes ──────────────────
        #if schema:
        #   try:
        #        recovered = cls.inject_defaults(recovered, schema, wrapper_key)
        #    except Exception as e:
        #        logger.warning(f"repair_failed_json: inject_defaults falhou: {e} — usando sem inject")

        # ── serializa e retorna ────────────────────────────────────────────
        try:
            result_str = json.dumps(recovered, ensure_ascii=False)
            logger.info("repair_failed_json: ✅ recuperação bem-sucedida")
            return result_str
        except Exception as e:
            logger.error(f"repair_failed_json: json.dumps falhou: {e}")
            return None

    # ──────────────────────────────────────────────────────────────────────────
    #  6. VALIDAÇÃO ESTRUTURAL
    # ──────────────────────────────────────────────────────────────────────────

    @classmethod
    def validate_structure(cls, text: str) -> bool:
        """
        Valida se {} e [] estão balanceados e aspas fechadas.
        Não valida JSON completo — apenas estrutura superficial.
        """
        depth_curly  = 0
        depth_square = 0
        in_string    = False
        escaped      = False

        for ch in text:
            if escaped:
                escaped = False
                continue
            if ch == '\\' and in_string:
                escaped = True
                continue
            if ch == '"':
                in_string = not in_string
                continue
            if in_string:
                continue
            if ch == '{':
                depth_curly  += 1
            elif ch == '}':
                depth_curly  -= 1
            elif ch == '[':
                depth_square += 1
            elif ch == ']':
                depth_square -= 1

            if depth_curly < 0 or depth_square < 0:
                return False

        return depth_curly == 0 and depth_square == 0 and not in_string

    # ──────────────────────────────────────────────────────────────────────────
    #  MÉTODOS PRIVADOS
    # ──────────────────────────────────────────────────────────────────────────

    @staticmethod
    def _try_parse(text: Optional[str]) -> Optional[Dict]:
        """Parse silencioso: retorna dict ou None."""
        if not text:
            return None
        try:
            result = json.loads(text)
            if isinstance(result, dict):
                return result
            return None
        except (json.JSONDecodeError, ValueError):
            return None

    @staticmethod
    def _isolate_outer_block(text: str) -> Optional[str]:
        """Encontra o primeiro { e o } correspondente, retorna o bloco."""
        start = text.find('{')
        if start == -1:
            return None

        depth    = 0
        in_str   = False
        escaped  = False

        for i, ch in enumerate(text[start:], start):
            if escaped:
                escaped = False
                continue
            if ch == '\\' and in_str:
                escaped = True
                continue
            if ch == '"':
                in_str = not in_str
                continue
            if in_str:
                continue
            if ch == '{':
                depth += 1
            elif ch == '}':
                depth -= 1
                if depth == 0:
                    return text[start:i + 1]

        return None  # bloco nunca fechado

    @classmethod
    def _extract_value_at(cls, text: str, pos: int) -> Optional[str]:
        """Extrai o valor JSON que começa na posição `pos` do texto."""
        if pos >= len(text):
            return None

        ch = text[pos]

        # String
        if ch == '"':
            end = pos + 1
            escaped = False
            while end < len(text):
                c = text[end]
                if escaped:
                    escaped = False
                elif c == '\\':
                    escaped = True
                elif c == '"':
                    return text[pos:end + 1]
                end += 1
            return None

        # Object ou Array
        if ch in ('{', '['):
            close = '}' if ch == '{' else ']'
            depth    = 0
            in_str   = False
            escaped  = False
            for i, c in enumerate(text[pos:], pos):
                if escaped:
                    escaped = False
                    continue
                if c == '\\' and in_str:
                    escaped = True
                    continue
                if c == '"':
                    in_str = not in_str
                    continue
                if in_str:
                    continue
                if c == ch:
                    depth += 1
                elif c == close:
                    depth -= 1
                    if depth == 0:
                        return text[pos:i + 1]
            return None

        # Primitivo (número, bool, null) — lê até vírgula, } ou ]
        end = pos
        while end < len(text) and text[end] not in (',', '}', ']', '\n'):
            end += 1
        raw = text[pos:end].strip()
        return raw if raw else None

    @classmethod
    def _extract_by_keys(cls, text: str, keys: List[str]) -> Optional[Dict]:
        """Tenta reconstruir dict extraindo cada key individualmente."""
        assembled: Dict[str, Any] = {}
        any_found = False

        for key in keys:
            raw_val = cls.extract_key(text, key)
            if raw_val is None:
                logger.debug(f"_extract_by_keys: key {key!r} não encontrada")
                continue
            try:
                assembled[key] = json.loads(raw_val)
                any_found = True
            except json.JSONDecodeError:
                sanitized_val = cls.sanitize_json(raw_val)
                try:
                    assembled[key] = json.loads(sanitized_val)
                    any_found = True
                except json.JSONDecodeError:
                    logger.debug(f"_extract_by_keys: falha ao parsear valor de {key!r}")

        return assembled if any_found else None

    @classmethod
    def _inject_recursive(cls, data: Any, schema: Dict) -> None:
        """Injeta defaults recursivamente in-place em `data` conforme `schema`."""
        if not isinstance(schema, dict):
            return

        schema_type = schema.get("type")

        # ── objeto ─────────────────────────────────────────────────────────
        if schema_type == "object" and isinstance(data, dict):
            properties = schema.get("properties", {})
            #required   = schema.get("required", [])

            #for field_name in required:
            #    if field_name not in data:
            #        field_schema = properties.get(field_name, {})
            #        data[field_name] = cls._default_for_schema(field_schema)
            #        logger.info(f"inject_defaults: injetou campo ausente {field_name!r}")

            # recursão nos campos presentes
            for field_name, field_schema in properties.items():
                if field_name in data:
                    cls._inject_recursive(data[field_name], field_schema)

        # ── array ──────────────────────────────────────────────────────────
        elif schema_type == "array" and isinstance(data, list):
            items_schema = schema.get("items", {})
            if items_schema and isinstance(items_schema, dict):
                for item in data:
                    if isinstance(item, dict):
                        cls._inject_recursive(item, items_schema)

    @classmethod
    def _default_for_schema(cls, schema: Dict) -> Any:
        """Retorna o valor default para um sub-schema."""
        if not isinstance(schema, dict):
            return None

        schema_type = schema.get("type")

        if schema_type == "object":
            obj: Dict = {}
            # preenche campos required do sub-objeto também
            properties = schema.get("properties", {})
            #required   = schema.get("required", [])
            #for field_name in required:
            #    field_schema = properties.get(field_name, {})
            #    obj[field_name] = cls._default_for_schema(field_schema)
            return obj

        if schema_type == "array":
            return []

        return _TYPE_DEFAULTS.get(schema_type, None)


# ═══════════════════════════════════════════════════════════════════════════════
#  INTEGRAÇÃO — patch para TrindadePipeline.api_call()
# ═══════════════════════════════════════════════════════════════════════════════
#
#  Substituir em external_worker.py, no bloco `elif result.code == 400:`:
#
#  ANTES:
#      elif result.code == 400:
#          logger.error(f"Erro 400 do provedor: {result.content}")
#          return None, None, False
#
#  DEPOIS:
#      elif result.code == 400:
#          recovered_str = JsonUtils.repair_failed_json(
#              raw_error_body = getattr(result, "raw_error_body", None) or result.content,
#              schema         = schema,           # schema já podado para o step
#              wrapper_key    = "manifestacao_juridica",
#          )
#          if recovered_str:
#              logger.warning(f"♻️  [{label}] 400 recuperado via repair_failed_json")
#              if stats:
#                  stats.errors_other -= 1        # não conta como erro
#                  stats.last_raw_response = recovered_str
#              return recovered_str, None, False  # trata como END_SUCCESS
#          logger.error(f"Erro 400 irrecuperável: {result.content}")
#          return None, None, False
#
# ─────────────────────────────────────────────────────────────────────────────
#  providers.py — APIResult precisa expor raw_error_body:
#
#  Adicionar campo em APIResult (dataclass):
#      raw_error_body: Optional[str] = None
#
#  No bloco de tratamento do 400 em providers.py, onde o `⚠️ output_keys paylpsd:`
#  é logado, salvar o body:
#      result.raw_error_body = response_text   # string do body HTTP 400
# ═══════════════════════════════════════════════════════════════════════════════


if __name__ == "__main__":
    # ── smoke tests ───────────────────────────────────────────────────────────
    logging.basicConfig(level=logging.DEBUG)

    # Caso 1: failed_generation com JSON válido mas campo ausente
    error_body_valid = json.dumps({
        "error": {
            "message": "jsonschema: missing properties: 'provas_valoradas'",
            "type": "invalid_request_error",
            "code": "json_validate_failed",
            "failed_generation": json.dumps({
                "manifestacao_juridica": {
                    "fundamentacao": {
                        "teses_relator": [{
                            "identificador_unico_fundamento_atomico": "inepcia",
                            "nucleo_argumentativo": "Pedido genérico.",
                            "resultado": "PROVIDO",
                            # provas_valoradas ausente
                        }],
                        "temas_nao_analisados": [],
                    },
                    "decisao_ementa": {
                        "resultado_global": "PROVIDO",
                        "mapa_resultados_por_pedido": [],
                    }
                }
            })
        }
    })

    schema_test = {
        "type": "object",
        "properties": {
            "manifestacao_juridica": {
                "type": "object",
                "required": ["fundamentacao", "decisao_ementa"],
                "properties": {
                    "fundamentacao": {
                        "type": "object",
                        "properties": {
                            "teses_relator": {
                                "type": "array",
                                "items": {
                                    "type": "object",
                                    "required": ["provas_valoradas", "resultado"],
                                    "properties": {
                                        "provas_valoradas": {"type": "array"},
                                        "resultado": {"type": "string"},
                                    }
                                }
                            },
                            "temas_nao_analisados": {"type": "array"},
                        }
                    },
                    "decisao_ementa": {"type": "object"},
                }
            }
        }
    }

    result = JsonUtils.repair_failed_json(error_body_valid, schema=schema_test)
    assert result is not None, "❌ Caso 1 falhou"
    parsed = json.loads(result)
    teses = parsed["manifestacao_juridica"]["fundamentacao"]["teses_relator"]
    assert teses[0].get("provas_valoradas") == [], f"❌ inject_defaults falhou: {teses[0]}"
    print("✅ Caso 1: JSON válido + inject_defaults OK")

    # Caso 2: failed_generation com malformação "campo",":",[]
    malformed = '{"manifestacao_juridica": {"fundamentacao": {"sintaxe": "ok", "temas_nao_analisados",":",[]], "decisao_ementa": {"resultado_global": "PROVIDO"}}}}'
    error_body_malformed = json.dumps({
        "error": {
            "code": "json_validate_failed",
            "failed_generation": malformed,
        }
    })

    result2 = JsonUtils.repair_failed_json(error_body_malformed)
    assert result2 is not None, "❌ Caso 2 falhou"
    parsed2 = json.loads(result2)
    assert parsed2["manifestacao_juridica"]["fundamentacao"]["temas_nao_analisados"] == []
    print("✅ Caso 2: malformação 'campo',':',[] sanitizada OK")

    # Caso 3: validate_structure
    assert JsonUtils.validate_structure('{"a": [1, 2]}') is True
    assert JsonUtils.validate_structure('{"a": [1, 2]}}') is False
    assert JsonUtils.validate_structure('{"a": "texto}"}') is True  # } dentro de string
    print("✅ Caso 3: validate_structure OK")

    # Caso 4: extract_key
    text4 = '{"alpha": {"x": 1}, "beta": [1,2,3], "gamma": "texto"}'
    assert json.loads(JsonUtils.extract_key(text4, "alpha")) == {"x": 1}
    assert json.loads(JsonUtils.extract_key(text4, "beta"))  == [1, 2, 3]
    assert json.loads(JsonUtils.extract_key(text4, "gamma")) == "texto"
    print("✅ Caso 4: extract_key OK")

    print("\n🎯 Todos os smoke tests passaram.")