"""
llm_json_fixer: A simple utility to fix malformed JSON from LLM outputs
"""

import json
import re
from typing import Any, Union


def fix_json(json_string: str, return_dict: bool = False) -> Union[Any, dict]:
    """
    Fix malformed JSON strings commonly generated by LLMs.
    
    Args:
        json_string (str): The malformed JSON string to fix
        return_dict (bool): If True, returns a dict with detailed info. If False, returns just the parsed JSON.
    
    Returns:
        If return_dict=False: The parsed JSON object, or None if fixing failed
        If return_dict=True: A dictionary with keys: 'success', 'data', 'error', 'fixes_applied'
    
    Examples:
        >>> fix_json("{'name': 'John', 'age': 30}")
        {'name': 'John', 'age': 30}
        
        >>> fix_json('{"name": "John", "age": 30,}')
        {'name': 'John', 'age': 30}
        
        >>> fix_json('{"active": True, "data": None}')
        {'active': True, 'data': None}
    """
    if not isinstance(json_string, str):
        if return_dict:
            return {'success': False, 'data': None, 'error': 'Input must be a string', 'fixes_applied': []}
        return None
    
    fixes_applied = []
    original_string = json_string
    
    # Try parsing as-is first
    try:
        result = json.loads(json_string)
        if return_dict:
            return {'success': True, 'data': result, 'error': None, 'fixes_applied': []}
        return result
    except json.JSONDecodeError:
        pass
    
    # Extract JSON from surrounding text
    json_patterns = [
        r'\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}',  # Nested objects
        r'\[[^\[\]]*(?:\[[^\[\]]*\][^\[\]]*)*\]',  # Nested arrays
        r'\{.*\}',  # Simple objects
        r'\[.*\]'   # Simple arrays
    ]
    
    for pattern in json_patterns:
        matches = re.findall(pattern, json_string, re.DOTALL)
        if matches:
            json_string = max(matches, key=len)  # Get longest match
            if json_string != original_string:
                fixes_applied.append("Extracted JSON from text")
            break
    
    # Remove markdown code blocks
    if json_string.startswith('```') and json_string.endswith('```'):
        json_string = re.sub(r'^```(?:json)?\s*', '', json_string)
        json_string = re.sub(r'\s*```$', '', json_string)
        fixes_applied.append("Removed markdown code blocks")
    
    # Basic cleanup
    json_string = json_string.strip()
    
    # Fix common issues
    original = json_string
    
    # Fix single quotes to double quotes
    json_string = re.sub(r"'([^']*)'", r'"\1"', json_string)
    
    # Fix Python booleans and None
    json_string = re.sub(r'\bTrue\b', 'true', json_string)
    json_string = re.sub(r'\bFalse\b', 'false', json_string)
    json_string = re.sub(r'\bNone\b', 'null', json_string)
    
    # Fix unquoted keys
    json_string = re.sub(r'(\w+)(?=\s*:)', r'"\1"', json_string)
    
    # Remove trailing commas
    json_string = re.sub(r',\s*}', '}', json_string)
    json_string = re.sub(r',\s*]', ']', json_string)
    
    # Remove multiple commas
    json_string = re.sub(r',\s*,+', ',', json_string)
    
    # Fix missing commas
    json_string = re.sub(r'}\s*{', '}, {', json_string)
    json_string = re.sub(r']\s*\[', '], [', json_string)
    json_string = re.sub(r'"\s*"(?=\s*:)', '", "', json_string)
    
    # Balance brackets
    open_braces = json_string.count('{')
    close_braces = json_string.count('}')
    open_brackets = json_string.count('[')
    close_brackets = json_string.count(']')
    
    if open_braces > close_braces:
        json_string += '}' * (open_braces - close_braces)
        fixes_applied.append("Added missing closing braces")
    elif close_braces > open_braces:
        json_string = json_string.rsplit('}', close_braces - open_braces)[0]
        fixes_applied.append("Removed extra closing braces")
        
    if open_brackets > close_brackets:
        json_string += ']' * (open_brackets - close_brackets)
        fixes_applied.append("Added missing closing brackets")
    elif close_brackets > open_brackets:
        json_string = json_string.rsplit(']', close_brackets - open_brackets)[0]
        fixes_applied.append("Removed extra closing brackets")
    
    # Fix unclosed strings
    quote_count = json_string.count('"') - json_string.count('\\"')
    if quote_count % 2 == 1:
        json_string += '"'
        fixes_applied.append("Closed unclosed string")
    
    # Remove trailing comma at end
    if json_string.rstrip().endswith(','):
        json_string = json_string.rstrip()[:-1]
        fixes_applied.append("Removed trailing comma at end")
    
    if original != json_string:
        fixes_applied.append("Applied common JSON fixes")
    
    # Try parsing after fixes
    try:
        result = json.loads(json_string)
        if return_dict:
            return {'success': True, 'data': result, 'error': None, 'fixes_applied': fixes_applied}
        return result
    except json.JSONDecodeError as e:
        if return_dict:
            return {'success': False, 'data': None, 'error': str(e), 'fixes_applied': fixes_applied}
        return None


def is_valid_json(json_string: str) -> bool:
    """
    Check if a string is valid JSON.
    
    Args:
        json_string (str): The string to check
        
    Returns:
        bool: True if valid JSON, False otherwise
        
    Examples:
        >>> is_valid_json('{"name": "John"}')
        True
        >>> is_valid_json("{'name': 'John'}")
        False
    """
    try:
        json.loads(json_string)
        return True
    except (json.JSONDecodeError, TypeError):
        return False




