# parse_llm_response.py
"""
This module handles parsing of LLM responses.
"""
from difflib import get_close_matches
from anndict.utils.anndictionary_ import normalize_string

def extract_dictionary_from_ai_string(
    ai_string : str
) -> str:
    """
    Cleans a generated string by removing everything before the first ``'{'``
    and everything after the first ``'}'``.

    Parameters
    -----------
    generated_string
        The string generated by retry_call_llm which includes unwanted characters or code.

    Returns
    -------
    A cleaned string that can be evaluated as a dictionary.
    """
    # Find the positions of the first '{' and the first '}'
    start = ai_string.find('{')
    end = ai_string.find('}')
    # Extract the first substring that starts with '{' and ends with '}'
    cleaned_string = ai_string[start:end+1]
    return cleaned_string


def extract_list_from_ai_string(
    ai_string
) -> str:
    """
    Cleans a generated string by removing everything before the first '{'
    and everything after the first '}'.

    Parameters
    ----------
    generated_string
        The string generated by retry_call_llm which includes unwanted characters or code.

    Returns
    -------
    A cleaned string that can be evaluated as a dictionary.
    """
    # Find the positions of the first '{' and the first '}'
    start = ai_string.find('[')
    end = ai_string.find(']')
    # Extract the first substring that starts with '{' and ends with '}'
    cleaned_string = ai_string[start:end+1]
    return cleaned_string


def process_llm_category_mapping(
    original_categories: list,
    llm_dict: dict[str, str]
    ) -> dict[str, str]:
    """
    Map original categories to simplified categories using LLM output.

    Normalizes strings, finds exact or close matches between original categories
    and LLM-provided keys, and creates a mapping. If no match is found, the 
    original category is used.

    Parameters
    -----------
    original_categories
        List of original category strings.

    llm_dict
        Dictionary of LLM-provided category mappings.

    Returns
    -------
    Mapping of original categories to simplified categories.
    """

    # Normalize original categories and LLM keys
    normalized_llm_keys = [normalize_string(key) for key in llm_dict.keys()]

    # Create mapping
    final_mapping = {}
    for original in original_categories:
        normalized = normalize_string(original)
        if normalized in normalized_llm_keys:
            # Direct match found
            index = normalized_llm_keys.index(normalized)
            final_mapping[original] = llm_dict[list(llm_dict.keys())[index]]
        else:
            # Find closest match
            matches = get_close_matches(normalized, normalized_llm_keys, n=1, cutoff=0.6)
            if matches:
                index = normalized_llm_keys.index(matches[0])
                final_mapping[original] = llm_dict[list(llm_dict.keys())[index]]
            else:
                # No match found, use original category
                final_mapping[original] = original

    return final_mapping
