# utils/safe_json.py — V1.0 (CANONICAL JSON EXTRACTOR)
"""
CTO Directive: One canonical JSON extractor for the entire server.
Replaces 3 ad-hoc implementations in orchestrator.py, strategy_manager.py,
problem_understanding.py, and ocr_strip_engine.py.

Rules:
1. LOG the RAW LLM string BEFORE any processing (critical for debugging hallucinations).
2. LaTeX backslash shield applied BEFORE extraction (prevents json.loads crashes).
3. json_repair applied before json.loads (handles LLM-malformed JSON).
4. Fail-CLOSED: any failure returns a safe error dict, never raises.
5. Supports both {} (dict) and [] (array) top-level JSON blocks.
"""

import re
import json
import logging
from typing import Union

logger = logging.getLogger(__name__)

try:
    from json_repair import repair_json
except ImportError:
    logger.warning("[SAFE_JSON] json_repair not available. Using raw json.loads fallback.")
    repair_json = lambda x: x  # noqa: E731

# Sentinel returned on any unrecoverable failure
_PARSE_FAILURE = {
    "logic_error": True,
    "error_type": "PARSING_FAILURE",
    "final_answer": "מצטערת, התשובה שיצרתי התבלגנה קצת בדרך. נוכל לנסות שוב? 🔄"
}


def safe_extract_json(
    llm_response_text: str,
    caller: str = "UNKNOWN",
    allow_array: bool = True
) -> Union[dict, list]:
    """
    Generic, robust JSON extraction from any LLM response string.

    Args:
        llm_response_text: Raw text from LLM (may include markdown, prose, LaTeX).
        caller: Identifier for the calling module (used in log tags).
        allow_array: If True, accepts top-level JSON arrays [...] as well as objects {...}.

    Returns:
        Parsed dict or list on success.
        _PARSE_FAILURE dict on any failure (never raises).
    """
    # ── Step 1: LOG RAW INPUT (CTO requirement) ────────────────────────────────
    logger.info(
        f"[SAFE_JSON:{caller}] RAW LLM STRING ({len(llm_response_text)} chars): "
        f"{llm_response_text[:600]!r}"
    )

    if not llm_response_text or not llm_response_text.strip():
        logger.error(f"[SAFE_JSON:{caller}] Empty LLM response.")
        return dict(_PARSE_FAILURE)

    # ── Step 2: Strip markdown fencing (```json ... ```) ──────────────────────
    # This is the most common wrapper the LLM adds despite instructions
    text = re.sub(r'```(?:json)?\s*', '', llm_response_text)
    text = text.replace('```', '')

    # ── Step 3: LaTeX Backslash Shield (CRITICAL FOR MATH) ─────────────────────
    # LLMs frequently output raw backslashes for LaTeX (\frac, \ln) which are
    # invalid in JSON strings. We escape them here unless already escaped.
    # Note: We only escape backslashes that are NOT followed by a valid JSON escape char
    # or another backslash.
    shielded = re.sub(r'\\(?![\\"/bfnrtu])', r'\\\\', text)

    # ── Step 4: Locate the outermost JSON block (REGEX MANDATE) ────────────────
    # V286.2: We use a robust re.DOTALL search to find the outermost {} or [] block.
    # This ignores any conversational preamble or postamble added by the LLM.
    pattern = r'(\{.*\}|\[.*\])'
    match = re.search(pattern, shielded, re.DOTALL)
    
    if not match:
        logger.error(
            f"[SAFE_JSON:{caller}] No JSON block found via Regex. "
            f"First 200 chars: {llm_response_text[:200]!r}"
        )
        return dict(_PARSE_FAILURE)

    extracted = match.group(0)

    # ── Step 5: Try raw json.loads first ─────────────────────────────────────
    try:
        parsed = json.loads(extracted)
        logger.info(f"[SAFE_JSON:{caller}] ✅ Parsed successfully (direct).")
        return _sanitize_math_newlines(parsed)
    except json.JSONDecodeError:
        pass  # Fall through to json_repair

    # ── Step 5b: json_repair + json.loads (for malformed LLM output) ─────────
    try:
        repaired = repair_json(extracted)
        parsed = json.loads(repaired)
        logger.info(f"[SAFE_JSON:{caller}] ✅ Parsed successfully (repair).")
        return _sanitize_math_newlines(parsed)

    except Exception as e:
        logger.error(f"[SAFE_JSON:{caller}] json_repair fail: {type(e).__name__} - {e}. Extracted: {extracted[:200]!r}")

    # ── Step 6: Last resort — try raw extraction without LaTeX shield ─────────
    # We look for anything that looks like a JSON block
    patterns = [
        r'(\{.*\}|\[.*\])',  # Greedy match for outermost block
        r'(\{[^{}]*\}|\[[^\[\]]*\])'  # Non-nested match as fallback
    ]

    try:
        for pattern in patterns:
            match = re.search(pattern, llm_response_text, re.DOTALL)
            if match:
                parsed = json.loads(repair_json(match.group(0)))
                logger.warning(f"[SAFE_JSON:{caller}] ✅ Recovered via raw extraction (no LaTeX shield).")
                return _sanitize_math_newlines(parsed)
    except Exception:
        pass

    logger.error(f"[SAFE_JSON:{caller}] 🚨 All extraction strategies failed. Returning PARSE_FAILURE.")
    return dict(_PARSE_FAILURE)


def _sanitize_math_newlines(data: Union[dict, list, str]) -> Union[dict, list, str]:
    """Recursively cleans up newline characters and restores swallowed LaTeX backslashes."""
    if isinstance(data, list):
        return [_sanitize_math_newlines(i) for i in data]
    elif isinstance(data, dict):
        new_dict = {}
        for k, v in data.items():
            if isinstance(v, str):
                # V286.1: Fix swallowed backslashes! 
                # When LLM outputs "\frac", JSON parses \f as form feed (\x0c). \beta as backspace (\x08).
                v = v.replace('\x0c', r'\f')
                v = v.replace('\x08', r'\b')
                v = v.replace('\t', r'\t')
                v = v.replace('\r', r'\r')
                
                # \n might be intended for \neq, \nu, \nabla, \normal, etc.
                import re
                v = re.sub(r'\n(eq|u|abla|ormal| left| right)', r'\\n\1', v)

                # Special case for math-heavy keys: ensure no \n breaks rendering
                if k in ['block_math', 'math_block', 'latex', 'equation', 'content'] or '$' in v:
                    v = v.replace('\n', ' ')
                
                new_dict[k] = v
            else:
                new_dict[k] = _sanitize_math_newlines(v)
        return new_dict
    return data