"""
cortex/extraction/json_utils.py — JSON repair utilities for LLM output.

Consolidates the inline _repair_json / _repair_truncated_json functions
that were previously in extractor.py. The standalone json_repair.py (which
was never imported anywhere) is deleted as part of this refactor.
"""

from __future__ import annotations

import json
import logging
import re

logger = logging.getLogger(__name__)


def _strip_reasoning_tags(text: str) -> str:
    """Strip <think>/<reasoning>/<thought> blocks from reasoning model output.

    Handles: closed tags, nested tags, unclosed tags (fallback to first JSON char).
    """
    # Strip closed reasoning tags (greedy — handles nested tags too)
    _TAG_PATTERN = re.compile(
        r'<(?:think|thinking|reasoning|thought)>.*?</(?:think|thinking|reasoning|thought)>',
        re.DOTALL,
    )
    result = _TAG_PATTERN.sub('', text)

    # Check for unclosed opening tags — strip everything up to first JSON char
    _OPEN_TAG = re.compile(r'<(?:think|thinking|reasoning|thought)>', re.IGNORECASE)
    if _OPEN_TAG.search(result):
        # Unclosed tag — find first { or [
        for i, ch in enumerate(result):
            if ch in ('{', '['):
                result = result[i:]
                break

    return result.strip()


def _strip_prose_preamble(text: str) -> str:
    """Strip prose text before the first JSON object/array character.

    Only strips if there's a clear { or [ that starts JSON content.
    Preserves the original text if no JSON structure is found.
    """
    stripped = text.strip()
    # Already starts with JSON — nothing to do
    if stripped and stripped[0] in ('{', '['):
        return stripped

    # Find first { or [ that could be JSON start
    for i, ch in enumerate(stripped):
        if ch in ('{', '['):
            return stripped[i:]

    # No JSON structure found — return as-is for downstream handling
    return stripped


def _repair_json(raw: str, session_id: str = "") -> str:
    """Best-effort repair of malformed LLM JSON output.

    Handles common failure modes from LLM structured output:
      0. Reasoning model tags (<think>...</think> and variants)
      0b. Prose preamble before JSON content
      1. Markdown code fences (```json ... ```)
      2. Trailing commas before ] or }
      3. Truncated JSON (output cut off mid-object due to max_tokens)

    Returns the repaired JSON string (may still be invalid if damage is
    too severe — caller should handle JSONDecodeError).
    """
    _stripped = raw.strip()

    # 0. Strip reasoning model tags (think/reasoning/thought)
    _stripped = _strip_reasoning_tags(_stripped)

    # 0b. Strip prose preamble before JSON (e.g. "Let me analyze...\n\n{...}")
    _stripped = _strip_prose_preamble(_stripped)

    # 1. Strip markdown code fences
    if _stripped.startswith("```"):
        first_nl = _stripped.find("\n")
        if first_nl != -1:
            _stripped = _stripped[first_nl + 1:]
        if _stripped.rstrip().endswith("```"):
            _stripped = _stripped.rstrip()[:-3].rstrip()

    # 2. Remove trailing commas before ] or }
    _stripped = re.sub(r',\s*([}\]])', r'\1', _stripped)

    # 3. Try parsing — if it works, no truncation repair needed
    try:
        json.loads(_stripped)
        return _stripped
    except json.JSONDecodeError:
        pass

    # 4. Truncation repair
    return _repair_truncated_json(_stripped, session_id)


def _repair_truncated_json(text: str, session_id: str = "") -> str:
    """Attempt to salvage claims from truncated JSON output.

    Strategy: find the last complete JSON object boundary (closing '}'),
    truncate there, and close any open array/object wrappers.
    """
    is_wrapped = False
    inner = text.strip()
    if inner.startswith("{"):
        claims_match = re.match(r'\{\s*"claims"\s*:\s*\[', inner)
        if claims_match:
            is_wrapped = True
            prefix = inner[:claims_match.end()]
            inner = inner[claims_match.end():]
        else:
            return text
    elif inner.startswith("["):
        prefix = "["
        inner = inner[1:]
    else:
        return text

    # Find last position where brace_depth returns to 0 (complete object)
    last_complete = -1
    brace_depth = 0
    in_string = False
    escape_next = False

    for i, ch in enumerate(inner):
        if escape_next:
            escape_next = False
            continue
        if ch == '\\' and in_string:
            escape_next = True
            continue
        if ch == '"' and not escape_next:
            in_string = not in_string
            continue
        if in_string:
            continue
        if ch == '{':
            brace_depth += 1
        elif ch == '}':
            brace_depth -= 1
            if brace_depth == 0:
                last_complete = i

    if last_complete == -1:
        return text

    salvaged_inner = inner[:last_complete + 1].rstrip().rstrip(',')

    if is_wrapped:
        repaired = prefix + salvaged_inner + "]}"
    else:
        repaired = prefix + salvaged_inner + "]"

    try:
        data = json.loads(repaired)
        claims_list = data if isinstance(data, list) else data.get("claims", [])
        salvaged_count = len(claims_list) if isinstance(claims_list, list) else 0
        logger.warning(
            "session=%s JSON truncation repaired: salvaged %d claims from truncated output",
            session_id, salvaged_count,
        )
        return repaired
    except json.JSONDecodeError:
        return text
