"""
Centralized utilities for extracting and parsing JSON from LLM responses.

Provides two public functions:
- get_llm_response_text: Extract text from LLM response objects
- extract_json: Parse JSON from LLM-generated text with robust cleaning
"""

import json
import re
from typing import Optional, Type, Union

from loguru import logger


def get_llm_response_text(response) -> str:
    """Extract text content from an LLM response object.

    Handles LangChain AIMessage (.content), plain text responses (.text),
    and arbitrary objects (via str()). Removes <think> tags from the output.

    Args:
        response: LLM response object, string, or None.

    Returns:
        Extracted text with think tags removed. Empty string for None input.
    """
    if response is None:
        return ""
    if hasattr(response, "content") and response.content is not None:
        raw = response.content
    elif hasattr(response, "text") and response.text is not None:
        raw = response.text
    else:
        raw = str(response)
    if not isinstance(raw, str):
        raw = str(raw)
    return _remove_think_tags(raw)


def extract_json(
    text: str,
    expected_type: Optional[Type] = None,
) -> Optional[Union[dict, list]]:
    """Extract and parse JSON from LLM-generated text.

    Applies a cleaning pipeline to handle common LLM output patterns:
    code fences, think tags, prose surrounding JSON, and minor artifacts.

    Args:
        text: Raw text potentially containing JSON.
        expected_type: Expected JSON type (dict or list). If specified,
            bracket extraction is ordered to prefer the matching type.
            None accepts either type.

    Returns:
        Parsed dict or list, or None if no valid JSON found.
    """
    if not text or not text.strip():
        return None

    text = text.strip()
    text = _strip_code_fences(text)
    text = _remove_think_tags(text)

    # Step 1: Try direct parse
    try:
        result = json.loads(text)
        if isinstance(result, (dict, list)):
            if expected_type is None or isinstance(result, expected_type):
                return result
            # Type mismatch — fall through to bracket extraction
    except (json.JSONDecodeError, ValueError):
        pass

    # Step 2: Bracket extraction ordered by expected_type
    if expected_type is list:
        bracket_pairs = [("[", "]"), ("{", "}")]
    elif expected_type is dict:
        bracket_pairs = [("{", "}"), ("[", "]")]
    else:
        bracket_pairs = [("{", "}"), ("[", "]")]

    for open_char, close_char in bracket_pairs:
        extracted = _extract_by_brackets(text, open_char, close_char)
        if extracted is None:
            continue

        # Try parsing the extracted substring
        try:
            result = json.loads(extracted)
            if isinstance(result, (dict, list)):
                if expected_type is None or isinstance(result, expected_type):
                    return result
        except (json.JSONDecodeError, ValueError):
            pass

        # Try cleaning LLM artifacts and retrying
        cleaned = _clean_llm_json_artifacts(extracted)
        if cleaned != extracted:
            try:
                result = json.loads(cleaned)
                if isinstance(result, (dict, list)):
                    if expected_type is None or isinstance(
                        result, expected_type
                    ):
                        return result
            except (json.JSONDecodeError, ValueError):
                pass

    logger.debug("No valid JSON found in text")
    return None


# ---------------------------------------------------------------------------
# Private helpers
# ---------------------------------------------------------------------------


def _remove_think_tags(text: str) -> str:
    """Remove <think>...</think> tags from text.

    Duplicated from search_utilities to keep this module free of
    internal package dependencies (follows type_utils.py convention).
    """
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
    text = re.sub(r"</think>", "", text)
    text = re.sub(r"<think>", "", text)
    return text.strip()


def _strip_code_fences(text: str) -> str:
    """Remove markdown code fences from text.

    Uses split-based extraction (not startswith) to handle fences
    appearing mid-text with surrounding prose.
    """
    if "```json" in text:
        parts = text.split("```json")
        if len(parts) > 1:
            return parts[1].split("```")[0].strip()
    elif "```" in text:
        parts = text.split("```")
        if len(parts) >= 3:
            return parts[1].strip()
    return text


def _extract_by_brackets(
    text: str, open_char: str, close_char: str
) -> Optional[str]:
    """Extract substring between outermost matching brackets.

    Uses find()/rfind() which is equivalent to re.search with re.DOTALL
    for bracket matching.
    """
    start = text.find(open_char)
    end = text.rfind(close_char)
    if start >= 0 and end > start:
        return text[start : end + 1]
    return None


def _clean_llm_json_artifacts(text: str) -> str:
    """Clean common LLM JSON artifacts from malformed JSON text.

    Only called after json.loads has already failed, so the text is
    already malformed. These regexes cannot corrupt valid JSON.

    Handles:
    - Trailing commas before ] or }
    - Inline // comments
    - Ellipsis entries (... or "...")
    """
    # Remove trailing commas
    text = re.sub(r",\s*([}\]])", r"\1", text)
    # Remove // line comments
    text = re.sub(r"//[^\n]*", "", text)
    # Remove ellipsis entries, preserving comma separator when between items
    text = re.sub(
        r',\s*"?\.\.\.+"?\s*,', ",", text
    )  # between items: keep one comma
    return re.sub(r',?\s*"?\.\.\.+"?\s*', "", text)  # trailing/leading ellipsis
