"""
Smart JSON extractor for handling messy LLM responses.

Handles common issues with cheap/weak models:
- JSON inside markdown code blocks (```json ... ```)
- Extra text before/after JSON
- Truncated JSON (attempts auto-repair)
- Multiple JSON objects (extracts first valid one)
- Common syntax errors (trailing commas, single quotes, etc.)
"""

import json
import logging
import re
from typing import Any, Optional, Type, TypeVar

from pydantic import BaseModel

logger = logging.getLogger(__name__)

T = TypeVar("T", bound=BaseModel)


class JSONExtractor:
    """
    Smart JSON extractor that handles messy LLM responses.

    Attempts multiple extraction strategies in order:
    1. Direct JSON parse
    2. Extract from markdown code blocks
    3. Find JSON object/array in text
    4. Fix common syntax errors
    5. Repair truncated JSON
    """

    # Pattern for markdown code blocks: ```json ... ``` or ``` ... ```
    CODE_BLOCK_PATTERN = re.compile(
        r"```(?:json)?\s*\n?(.*?)\n?```",
        re.DOTALL | re.IGNORECASE
    )

    # Pattern to find JSON object { ... } in text
    JSON_OBJECT_PATTERN = re.compile(
        r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}",
        re.DOTALL
    )

    # Pattern to find JSON array [ ... ] in text
    JSON_ARRAY_PATTERN = re.compile(
        r"\[[^\[\]]*(?:\[[^\[\]]*\][^\[\]]*)*\]",
        re.DOTALL
    )

    @classmethod
    def extract(
        cls,
        text: str,
        schema: Optional[Type[T]] = None,
    ) -> tuple[Optional[dict[str, Any]], Optional[T], str]:
        """
        Extract JSON from text, optionally validating against Pydantic schema.

        Args:
            text: Raw LLM response text
            schema: Optional Pydantic model class for validation

        Returns:
            Tuple of (raw_dict, parsed_model, extraction_method)
            - raw_dict: Extracted JSON as dict, or None if failed
            - parsed_model: Validated Pydantic model, or None if no schema or failed
            - extraction_method: String describing how JSON was extracted
        """
        if not text or not text.strip():
            return None, None, "empty_input"

        text = text.strip()

        # Strategy 1: Direct parse
        result = cls._try_direct_parse(text)
        if result is not None:
            return cls._validate_and_return(result, schema, "direct_parse")

        # Strategy 2: Extract from markdown code blocks
        result = cls._extract_from_code_blocks(text)
        if result is not None:
            return cls._validate_and_return(result, schema, "code_block")

        # Strategy 3: Find JSON object in text
        result = cls._find_json_in_text(text)
        if result is not None:
            return cls._validate_and_return(result, schema, "text_search")

        # Strategy 4: Fix common syntax errors and retry
        fixed_text = cls._fix_common_errors(text)
        if fixed_text != text:
            result = cls._try_direct_parse(fixed_text)
            if result is not None:
                return cls._validate_and_return(result, schema, "syntax_fix")

            result = cls._find_json_in_text(fixed_text)
            if result is not None:
                return cls._validate_and_return(result, schema, "syntax_fix_search")

        # Strategy 5: Try to repair truncated JSON
        result = cls._repair_truncated_json(text)
        if result is not None:
            return cls._validate_and_return(result, schema, "truncation_repair")

        logger.warning("All JSON extraction strategies failed for: %.200s...", text)
        return None, None, "failed"

    @classmethod
    def _validate_and_return(
        cls,
        data: dict[str, Any],
        schema: Optional[Type[T]],
        method: str,
    ) -> tuple[dict[str, Any], Optional[T], str]:
        """Validate extracted data against schema if provided."""
        if schema is None:
            return data, None, method

        try:
            validated = schema.model_validate(data)
            return data, validated, method
        except Exception as e:
            logger.warning(
                "JSON extracted via %s but validation failed: %s",
                method, e
            )
            return data, None, f"{method}_validation_failed"

    @classmethod
    def _try_direct_parse(cls, text: str) -> Optional[dict[str, Any]]:
        """Try to parse text directly as JSON."""
        try:
            result = json.loads(text)
            if isinstance(result, dict):
                return result
            elif isinstance(result, list):
                # Wrap list in dict for consistency
                return {"items": result}
        except json.JSONDecodeError:
            pass
        return None

    @classmethod
    def _extract_from_code_blocks(cls, text: str) -> Optional[dict[str, Any]]:
        """Extract JSON from markdown code blocks."""
        matches = cls.CODE_BLOCK_PATTERN.findall(text)

        for match in matches:
            content = match.strip()
            if not content:
                continue

            # Try direct parse
            result = cls._try_direct_parse(content)
            if result is not None:
                return result

            # Try with syntax fixes
            fixed = cls._fix_common_errors(content)
            result = cls._try_direct_parse(fixed)
            if result is not None:
                return result

        return None

    @classmethod
    def _find_json_in_text(cls, text: str) -> Optional[dict[str, Any]]:
        """Find JSON object or array embedded in text."""
        # First try to find balanced braces/brackets
        result = cls._find_balanced_json(text)
        if result is not None:
            return result

        # Fallback to regex patterns
        # Try object pattern first
        for match in cls.JSON_OBJECT_PATTERN.finditer(text):
            candidate = match.group(0)
            result = cls._try_direct_parse(candidate)
            if result is not None:
                return result

            # Try with fixes
            fixed = cls._fix_common_errors(candidate)
            result = cls._try_direct_parse(fixed)
            if result is not None:
                return result

        # Try array pattern
        for match in cls.JSON_ARRAY_PATTERN.finditer(text):
            candidate = match.group(0)
            result = cls._try_direct_parse(candidate)
            if result is not None:
                return result

        return None

    @classmethod
    def _find_balanced_json(cls, text: str) -> Optional[dict[str, Any]]:
        """Find JSON by matching balanced braces/brackets."""
        # Find first { or [
        start_idx = -1
        start_char = None

        for i, char in enumerate(text):
            if char == '{':
                start_idx = i
                start_char = '{'
                break
            elif char == '[':
                start_idx = i
                start_char = '['
                break

        if start_idx == -1:
            return None

        # Match balanced delimiters
        end_char = '}' if start_char == '{' else ']'
        depth = 1
        in_string = False
        escape = False

        for i in range(start_idx + 1, len(text)):
            char = text[i]

            if escape:
                escape = False
                continue

            if char == '\\':
                escape = True
                continue

            if char == '"':
                in_string = not in_string
                continue

            if in_string:
                continue

            if char == start_char:
                depth += 1
            elif char == end_char:
                depth -= 1
                if depth == 0:
                    candidate = text[start_idx:i + 1]
                    result = cls._try_direct_parse(candidate)
                    if result is not None:
                        return result

                    # Try with fixes
                    fixed = cls._fix_common_errors(candidate)
                    result = cls._try_direct_parse(fixed)
                    if result is not None:
                        return result
                    break

        return None

    @classmethod
    def _fix_common_errors(cls, text: str) -> str:
        """Fix common JSON syntax errors from LLMs."""
        # Remove BOM and other invisible characters
        text = text.strip('\ufeff\u200b\u200c\u200d')

        # Fix single quotes to double quotes (careful with apostrophes)
        # Only replace when it looks like JSON string delimiter
        text = re.sub(r"(?<=[\[{,:\s])'(?=\s*[^'])", '"', text)
        text = re.sub(r"'(?=\s*[,}\]])", '"', text)

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

        # Fix missing quotes around property names
        text = re.sub(r"(\{|,)\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*:", r'\1"\2":', text)

        # Fix Python-style booleans and None
        text = re.sub(r'\bTrue\b', 'true', text)
        text = re.sub(r'\bFalse\b', 'false', text)
        text = re.sub(r'\bNone\b', 'null', text)

        # Fix unescaped newlines in strings (simple heuristic)
        # This is tricky - only do if it helps

        return text

    @classmethod
    def _repair_truncated_json(cls, text: str) -> Optional[dict[str, Any]]:
        """
        Attempt to repair truncated JSON by closing open brackets/braces.

        This is a best-effort repair for when models hit token limits.
        """
        # Find the start of JSON
        start_idx = -1
        for i, char in enumerate(text):
            if char in '{[':
                start_idx = i
                break

        if start_idx == -1:
            return None

        # Extract from start and analyze structure
        json_part = text[start_idx:]

        # Count unclosed delimiters
        stack: list[str] = []
        in_string = False
        escape = False

        for char in json_part:
            if escape:
                escape = False
                continue

            if char == '\\':
                escape = True
                continue

            if char == '"':
                in_string = not in_string
                continue

            if in_string:
                continue

            if char in '{[':
                stack.append(char)
            elif char == '}':
                if stack and stack[-1] == '{':
                    stack.pop()
            elif char == ']':
                if stack and stack[-1] == '[':
                    stack.pop()

        # If string is unclosed, try to close it
        if in_string:
            json_part = json_part.rstrip()
            if not json_part.endswith('"'):
                json_part += '"'

        # Close unclosed delimiters
        closing = ""
        for opener in reversed(stack):
            if opener == '{':
                closing += '}'
            elif opener == '[':
                closing += ']'

        if closing:
            # Try to make it valid by closing
            # First, try truncating at a reasonable point
            # (after last complete value)

            # Simple approach: just close what's open
            repaired = json_part + closing
            result = cls._try_direct_parse(repaired)
            if result is not None:
                logger.info("Repaired truncated JSON by adding: %s", closing)
                return result

            # Try removing partial content before closing
            # Look for last comma or colon and truncate there
            for truncate_pattern in [
                r',\s*"[^"]*"?\s*:\s*$',  # Key with colon but no value: ,"key":
                r',\s*"[^"]*$',  # Partial key after comma
                r':\s*"[^"]*$',  # Partial string value
                r':\s*$',  # Colon with no value
                r',\s*\d+\.?\d*$',  # Partial number
                r',\s*$',  # Trailing comma
            ]:
                truncated = re.sub(truncate_pattern, '', json_part)
                if truncated != json_part:
                    repaired = truncated + closing
                    result = cls._try_direct_parse(repaired)
                    if result is not None:
                        logger.info("Repaired truncated JSON by truncating and closing")
                        return result

        return None


def extract_json(
    text: str,
    schema: Optional[Type[T]] = None,
) -> tuple[Optional[dict[str, Any]], Optional[T], str]:
    """
    Convenience function for JSONExtractor.extract().

    Args:
        text: Raw LLM response text
        schema: Optional Pydantic model class for validation

    Returns:
        Tuple of (raw_dict, parsed_model, extraction_method)
    """
    return JSONExtractor.extract(text, schema)


def extract_json_or_raise(
    text: str,
    schema: Type[T],
    error_class: Type[Exception] = ValueError,
) -> T:
    """
    Extract and validate JSON, raising on failure.

    Args:
        text: Raw LLM response text
        schema: Pydantic model class for validation
        error_class: Exception class to raise on failure

    Returns:
        Validated Pydantic model instance

    Raises:
        error_class: If extraction or validation fails
    """
    raw, parsed, method = JSONExtractor.extract(text, schema)

    if parsed is not None:
        return parsed

    if raw is not None:
        raise error_class(
            f"JSON extracted via {method} but failed schema validation. "
            f"Raw data: {raw}"
        )

    raise error_class(
        f"Failed to extract JSON from response. Method: {method}. "
        f"Response: {text[:500]}..."
    )


__all__ = [
    "JSONExtractor",
    "extract_json",
    "extract_json_or_raise",
]
