"""
Parser utilities for extracting JSON from LLaMA responses
"""

import json
import re
import logging
from typing import Dict, Optional

logger = logging.getLogger(__name__)


def parse_json_response(response_text: str) -> Optional[Dict[str, float]]:
    """
    Parse JSON from LLaMA response text
    
    The model might return extra text before/after the JSON, so we need to extract it.
    
    Args:
        response_text: Raw text response from LLaMA
        
    Returns:
        Parsed dictionary of probabilities, or None if parsing fails
    """
    try:
        # Try direct JSON parsing first
        return json.loads(response_text.strip())
    except json.JSONDecodeError:
        pass
    
    # Try to find JSON object in the text
    try:
        # Look for content between curly braces
        json_match = re.search(r'\{[^{}]*\}', response_text)
        if json_match:
            json_str = json_match.group(0)
            return json.loads(json_str)
    except (json.JSONDecodeError, AttributeError) as e:
        logger.error(f"Failed to parse JSON from response: {e}")
        logger.debug(f"Response text: {response_text}")
    
    return None


def normalize_probabilities(probs: Dict[str, float]) -> Dict[str, float]:
    """
    Normalize probability distribution to sum to 1.0
    
    Args:
        probs: Dictionary of class probabilities
        
    Returns:
        Normalized probabilities
    """
    total = sum(probs.values())
    
    if total == 0:
        # If all zeros, return uniform distribution
        n = len(probs)
        return {k: 1.0 / n for k in probs.keys()}
    
    # Normalize to sum to 1.0
    return {k: v / total for k, v in probs.items()}


def ensure_all_classes(probs: Dict[str, float], all_classes: list) -> Dict[str, float]:
    """
    Ensure all classes are present in probability dictionary
    
    Args:
        probs: Partial probability dictionary
        all_classes: List of all possible classes
        
    Returns:
        Complete probability dictionary with all classes
    """
    complete_probs = {cls: 0.0 for cls in all_classes}
    complete_probs.update(probs)
    return complete_probs


def validate_and_fix_probabilities(
    probs: Optional[Dict[str, float]],
    all_classes: list,
    default_class: str = None
) -> Dict[str, float]:
    """
    Validate and fix probability distribution

    Args:
        probs: Probability dictionary (might be None or invalid)
        all_classes: List of all possible classes
        default_class: Default class to assign high probability if parsing fails

    Returns:
        Valid, normalized probability distribution
    """
    # If parsing failed, return default distribution
    if probs is None:
        logger.warning(f"Parsing failed, using default class: {default_class}")
        if default_class and default_class in all_classes:
            result = {cls: 0.0 for cls in all_classes}
            result[default_class] = 1.0
            return result
        else:
            # Uniform distribution (FALLBACK - indicates model failure)
            logger.error("No default class provided, falling back to uniform distribution")
            n = len(all_classes)
            return {cls: 1.0 / n for cls in all_classes}

    # If model returned partial results (e.g., top 3), that's OK
    # Just ensure all classes are present (missing ones get 0.0)
    probs = ensure_all_classes(probs, all_classes)

    # Normalize
    probs = normalize_probabilities(probs)

    # Ensure values are between 0 and 1
    probs = {k: max(0.0, min(1.0, v)) for k, v in probs.items()}

    logger.info(f"Validated probabilities - Top class: {max(probs, key=probs.get)} ({max(probs.values()):.2%})")

    return probs

