editscore/utils.py

import os
from typing import Union, List, Optional
import json
import regex as re
import ast
import random
import json_repair

def fix_json(input_str):
    # Add double quotes around keys using regex
    fixed_str = re.sub(r'(\w+):', r'"\1":', input_str)
    
    # Add double quotes around string values if necessary and wrap int/float values in []
    def format_value(match):
        key, value, comma = match.groups()
        value = value.strip()
        # Check if value is an integer or float
        if re.match(r'^-?\d+(\.\d+)?$', value):
            value = f'[{value}]'
        # Check if value is a boolean or null
        elif re.match(r'^(true|false|null)$', value, re.IGNORECASE):
            pass  # leave as is
        else:
            # Add quotes around string values
            value = f'"{value}"'
        return f'{key}: {value}{comma}'
    
    fixed_str = re.sub(r'(".*?"):(.*?)(,|})', format_value, fixed_str)
    
    return fixed_str

def repair_reasoning_field_robust(json_str: str) -> str:
    """
    Robustly repair unescaped double quotes inside the "reasoning" field of a JSON string.
    This function uses regular expressions and a lookahead assertion to locate
    the end of the "reasoning" value, even if it is not the last field in the JSON.

    Args:
        json_str (str): A possibly malformed JSON string that may contain
                        unescaped quotes within the "reasoning" field.

    Returns:
        str: A repaired JSON string with properly escaped quotes inside "reasoning".
    """
    # 1. Define a regex pattern that locates the reasoning value using a lookahead.
    # The re.DOTALL flag allows '.' to match newline characters.
    pattern = re.compile(
        # --- Group 1: prefix part including the "reasoning" key and opening quote ---
        r'("reasoning"\s*:\s*")'
        
        # --- Group 2: content inside the reasoning string (non-greedy) ---
        r'(.*?)'
        
        # --- Lookahead assertion ---
        # Match the ending quote of the "reasoning" value,
        # but only if it is followed by a comma or closing brace.
        r'(?="\s*[,}])',
        
        re.DOTALL
    )

    # 2. Define a replacement function to escape quotes inside the "reasoning" content.
    def replacer(match):
        prefix = match.group(1)   # e.g., '"reasoning": "'
        content = match.group(2)  # e.g., 'Overall building...'

        # Escape all unescaped double quotes inside the reasoning text.
        fixed_content = content.replace('"', '\\"')

        # Reassemble the full matched segment. The suffix is not consumed by the pattern,
        # so we just return the prefix + repaired content.
        return prefix + fixed_content

    # 3. Apply the regex substitution across the entire JSON string.
    repaired_str = pattern.sub(replacer, json_str)

    return repaired_str

def fallback_repair_json(input_str: str) -> str:
    """
    Last-resort JSON repair that tries to preserve the 'reasoning' text
    even when it contains unescaped quotes or other corruption.

    Target output:
      {"reasoning": "<text>", "score": [float, float]}

    Approach:
      1. Locate 'reasoning' key position and 'score' key position.
      2. Extract the raw substring between them (reasoning_raw).
      3. Clean only the outer noise (leading/trailing quotes, commas, braces),
         but preserve internal punctuation.
      4. Unescape common escape sequences and normalize quotes.
      5. Extract numeric scores robustly.
      6. Return a valid JSON string.
    """

    s = input_str

    # Normalize whitespace for easier searching (but keep original for slicing)
    lowered = s.lower()

    # 1) find the start of reasoning key (case-insensitive)
    m_reason = re.search(r'"?reasoning"?\s*[:：]', lowered)
    m_score = re.search(r'"?score"?\s*[:：]', lowered)

    reasoning_text = ""
    scores = []

    if m_reason and m_score:
        # compute the real indices in the original string
        start_idx = m_reason.end()  # right after colon in 'reasoning:'
        score_start_idx = m_score.start()

        # 2) slice the original string between reasoning value start and score key start
        reasoning_raw = s[start_idx:score_start_idx]

        # 3) clean outer noise but preserve inner content:
        #    - strip whitespace and outer commas/braces
        reasoning_raw = reasoning_raw.strip()
        # remove leading commas/braces/colons
        reasoning_raw = re.sub(r'^[\s,{\[]+', '', reasoning_raw)
        # remove trailing commas/braces/colons (but keep inner punctuation)
        reasoning_raw = re.sub(r'[\s,}\]]+$', '', reasoning_raw)

        # If the reasoning starts with a quote char, drop it (we'll re-escape later).
        if reasoning_raw.startswith(("'", '"')):
            reasoning_raw = reasoning_raw[1:]
        # If it ends with a quote char (common), drop it.
        if reasoning_raw.endswith(("'", '"')):
            reasoning_raw = reasoning_raw[:-1]

        # 4) normalize escapes:
        # Replace common escaped sequences (\" -> "), but avoid creating unbalanced quotes.
        reasoning_raw = reasoning_raw.replace('\\"', '"').replace("\\'", "'")
        # Replace fancy quotes with straight quotes (optional)
        reasoning_raw = re.sub(r'[“”]', '"', reasoning_raw)
        reasoning_raw = re.sub(r"[‘’]", "'", reasoning_raw)

        # Trim again
        reasoning_text = reasoning_raw.strip()
    else:
        # If we couldn't find both keys, try a looser regex capturing 'reasoning' value
        m_loose = re.search(r'"?reasoning"?\s*[:：]\s*["\']?(.*?)["\']?\s*(,|$)', s, re.DOTALL | re.IGNORECASE)
        if m_loose:
            reasoning_text = m_loose.group(1).strip()
            # normalize escapes as above
            reasoning_text = reasoning_text.replace('\\"', '"').replace("\\'", "'")
            reasoning_text = re.sub(r'[“”]', '"', reasoning_text)
            reasoning_text = re.sub(r"[‘’]", "'", reasoning_text)

    # 5) Extract two numeric scores anywhere after the 'score' key (robust)
    if m_score:
        # slice from score key to the end
        score_slice = s[m_score.end():]
        # find numbers (integers or floats)
        nums = re.findall(r'-?\d+(?:\.\d+)?', score_slice)
        try:
            scores = [float(n) for n in nums[:2]]
        except Exception:
            scores = []
    else:
        # fallback: try to find any two numbers in the whole string
        nums = re.findall(r'-?\d+(?:\.\d+)?', s)
        try:
            scores = [float(n) for n in nums[:2]]
        except Exception:
            scores = []

    # Ensure we always return two floats
    if len(scores) < 2:
        scores += [0.0] * (2 - len(scores))

    # 6) Construct final object. Let json.dumps handle escaping inside the reasoning.
    repaired_obj = {
        "reasoning": reasoning_text,
        "score": scores
    }

    return json.dumps(repaired_obj, ensure_ascii=False)

def robust_json_fix(s: str):
    try:
        return json_repair.loads(s)
    except Exception:
        pass
    
    for fixer in [fix_json, repair_reasoning_field_robust]:
        s = fixer(s)
        try:
            return json_repair.loads(s)
        except Exception:
            print(f"Error: Cannot fix {fixer.__name__} {s=}")
            continue

    try:
        repaired_str = fallback_repair_json(s)
        return json_repair.loads(repaired_str)
    except Exception as e:
        print(f"Error: Cannot fix fallback_repair_json {s=} {e=}")
        return False
        
def read_file_to_string(file_path):
    """
    Reads the contents of a text file and returns it as a string.

    :param file_path: The path to the text file.
    :return: A string containing the contents of the file.
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            return file.read()
    except FileNotFoundError:
        print(f"The file {file_path} was not found.")
        return None
    except Exception as e:
        print(f"An error occurred: {e}")
        return None

def read_files_to_string(file_paths):
    """
    Reads the contents of multiple text files and returns them as a single string,
    with each file's contents separated by a newline.

    :param file_paths: A list of paths to text files.
    :return: A string containing the concatenated contents of the files.
    """
    all_contents = []  # List to hold the contents of each file

    for file_path in file_paths:
        try:
            with open(file_path, 'r', encoding='utf-8') as file:
                all_contents.append(file.read())
        except FileNotFoundError:
            print(f"The file {file_path} was not found.")
        except Exception as e:
            print(f"An error occurred while reading {file_path}: {e}")

    # Join all the contents with a newline character
    return "\n".join(all_contents)

def get_file_path(filename: Union[str, os.PathLike], search_from: Union[str, os.PathLike] = "."):
    """
    Search for a file across a directory and return its absolute path.

    Args:
        filename (Union[str, os.PathLike]): The name of the file to search for.
        search_from (Union[str, os.PathLike], optional): The directory from which to start the search. Defaults to ".".

    Returns:
        str: Absolute path to the found file.

    Raises:
        FileNotFoundError: If the file is not found.
    """
    for root, dirs, files in os.walk(search_from):
        for name in files:
            if name == filename:
                return os.path.abspath(os.path.join(root, name))
    raise FileNotFoundError(filename, "not found.")



#+=========================================================================================
def verify(s, target_sequence):
    # Count the occurrences of the target sequence
    count = s.count(target_sequence)

    # Check if the target sequence appears exactly twice
    return count == 2


def is_int_between_0_and_10(s):
    try:
        num = int(s)
        return 0 <= num <= 10
    except ValueError:
        return False
    
def is_str_a_list_of_ints_0_to_10(s):
    try:
        # Attempt to parse the string as a Python literal (list, dict, etc.)
        parsed = ast.literal_eval(s)

        # Check if the parsed object is a list
        if not isinstance(parsed, list):
            return False

        # Check if all elements are integers and between 0 to 10
        return all(isinstance(item, int) and 0 <= item <= 10 for item in parsed)

    except (ValueError, SyntaxError):
        # If parsing fails or any other error occurs
        return False
    
def is_str_valid_score_format_brackets(s):
    try:
        # Removing brackets and splitting the string by commas
        content = s.strip("[]").split(',')

        length = len(content)

        # Parsing each element and checking the format and range
        scores = {}
        for item in content:
            key, value = item.split(':')
            key = key.strip()
            value = int(value.strip())

            # Check if the key starts with 'score' and the value is in the correct range
            if not key.startswith("score") or not 0 <= value <= 10:
                return False

            scores[key] = value

        fetch_words = [f"score{i+1}" for i in range(length)]
        # Check if at least 'score1' and 'score2' are present
        return all(key in scores for key in fetch_words)

    except (ValueError, SyntaxError):
        # If any parsing error occurs
        return False

def normalize_quotes(s: str) -> str:
    """
    Replace curly/smart quotes with normal ASCII quotes.
    """
    # 常见的几种智能引号 U+201C U+201D U+2018 U+2019
    return s.replace("“", '"').replace("”", '"').replace("‘", "'").replace("’", "'")


#+=========================================================================================
def mllm_output_to_dict(input_string, give_up_parsing=False, text_prompt=None, score_range: int = 10):
    """
    Args:
        input_string (str): actually the output of the mllm model to be parsed
        output_file_name (str): The name of the output file.
    """
    # Catch for gpt4v rate_limit_exceeded error
    if input_string == "rate_limit_exceeded":
        return "rate_limit_exceeded"

    if give_up_parsing:
        guessed_value = random.randint(0, score_range)
        json_content = {'score': [guessed_value, guessed_value], "reasoning": f"guess_if_cannot_parse | {input_string}"}
        return json_content
    
    # Define the delimiters
    delimiter = '||V^=^V||'

    if input_string.count(delimiter) == 2:
        if not verify(input_string, delimiter):
            print("The required delimiters were not found correctly in the string.", flush=True)
            return False
        # Extract the content between the delimiters
        start_index = input_string.find(delimiter) + len(delimiter)
        end_index = input_string.rfind(delimiter)
    else:
        # find the json mannually
        # some mllm tends not to output the delimiters, but it does output the json contents
        # so we will find the json content mannually
        start_index = input_string.find('{')
        end_index = input_string.rfind('}') + 1
        if start_index == -1 or end_index == 0:
            # json not found
            # some mllm tends to output only a list of scores like [6, 0], 
            # this time we will just get the scores and ignore the reasoning (other part of the json)
            start_index = input_string.find('[')
            end_index = input_string.rfind(']') + 1
            if re.match(r'^\[\d+, ?\d+\]$', input_string[start_index:end_index]):
                scores = json.loads(input_string[start_index:end_index])
                if not isinstance(scores, list):
                    scores = [scores]
                json_content = {'score': scores, "reasoning": "System: output is simply a list of scores"}
                json_str = json.dumps(json_content)
                input_string = json_str
                start_index = 0
                end_index = len(json_str)
            elif is_int_between_0_and_10(input_string): # if output is simply a number
                scores = [int(input_string)]
                json_content = {'score': scores, "reasoning": "System: output is simply a number"}
                json_str = json.dumps(json_content)
                input_string = json_str
                start_index = 0
                end_index = len(json_str)
            else:
                print(f"22 222 Failed to find the json content in the string. {text_prompt=} {input_string=}", flush=True)
                return False
    
    # Check if we found two delimiters
    if start_index != -1 and end_index != -1 and start_index != end_index:
        # Extract the JSON string
        json_str = input_string[start_index:end_index].strip()
        json_str = json_str.replace("\n", "")
        # Parse the JSON string into a dictionary
        try:
            json_str = normalize_quotes(json_str)
            new_data = json.loads(json_str)
            if not isinstance(new_data['score'], list):
                new_data['score'] = [new_data['score']]
        except Exception as e1:
            print(f"Now fixing: {e1=} {json_str=}")
            
            new_data = robust_json_fix(json_str)
        return new_data
    else:
        print("The required delimiters were not found correctly in the string.")
        return False

def write_entry_to_json_file(input_string, uid, prompt_input, vision_input, output_file_name, give_up_parsing=False):
    """
    Args:
        input_string (str): actually the output of the mllm model to be parsed
        uid (str): The unique identifier for the each item in the test data
        prompt_input (str): The prompt input for the entry. text prompt.
        vision_input (str): The vision input for the entry. image links.
        output_file_name (str): The name of the output file.
    """
    # Catch for gpt4v rate_limit_exceeded error
    if input_string == "rate_limit_exceeded":
        return "rate_limit_exceeded"

    # Define the delimiters
    delimiter = '||V^=^V||'

    if input_string.count(delimiter) == 2:
        if not verify(input_string, delimiter):
            print("The required delimiters were not found correctly in the string.")
            return False
        # Extract the content between the delimiters
        start_index = input_string.find(delimiter) + len(delimiter)
        end_index = input_string.rfind(delimiter)
    else:
        # find the json mannually
        # some mllm tends not to output the delimiters, but it does output the json contents
        # so we will find the json content mannually
        start_index = input_string.find('{')
        end_index = input_string.rfind('}') + 1
        if start_index == -1 or end_index == 0:
            # json not found
            # some mllm tends to output only a list of scores like [6, 0], 
            # this time we will just get the scores and ignore the reasoning (other part of the json)
            start_index = input_string.find('[')
            end_index = input_string.rfind(']') + 1
            if give_up_parsing: # if we want to give up parsing
                guessed_value = random.randint(0, 10)
                print(f"Failed to find the json content in the string. Guess a value : {guessed_value}.")
                json_content = {'score': [guessed_value], "reasoning": f"guess_if_cannot_parse | {input_string}"}
                json_str = json.dumps(json_content)
                input_string = json_str
                start_index = 0
                end_index = len(json_str)
            elif re.match(r'^\[\d+, ?\d+\]$', input_string[start_index:end_index]):
                scores = json.loads(input_string[start_index:end_index])
                json_content = {'score': scores, "reasoning": None}
                json_str = json.dumps(json_content)
                input_string = json_str
                start_index = 0
                end_index = len(json_str)
            elif is_int_between_0_and_10(input_string): # if output is simply a number
                scores = [int(input_string)]
                json_content = {'score': scores, "reasoning": None}
                json_str = json.dumps(json_content)
                input_string = json_str
                start_index = 0
                end_index = len(json_str)
            else:
                print("Failed to find the json content in the string.")
                return False
    
    # Check if we found two delimiters
    if start_index != -1 and end_index != -1 and start_index != end_index:
        # Extract the JSON string
        json_str = input_string[start_index:end_index].strip()
        json_str = json_str.replace("\n", "")
        try:
            # Parse the JSON string into a dictionary
            new_data = json.loads(json_str)
            
            # Ensure the directory exists
            os.makedirs(os.path.dirname(output_file_name), exist_ok=True)
            
            # Initialize or load existing data
            if os.path.exists(output_file_name):
                with open(output_file_name, 'r') as json_file:
                    data = json.load(json_file)
            else:
                data = {}
            
            # If the additional key is already in the data, add or update notes
            if uid in data:
                data[uid].update(new_data)  # Update with new data
                if prompt_input:  # If there are new notes, update or add them
                    data[uid]['prompt_input'] = prompt_input
                if vision_input:  # If there are new notes, update or add them
                    data[uid]['vision_input'] = vision_input
            else:
                # If it's a new key, add the entry to the dictionary
                data[uid] = new_data
                if prompt_input:
                    data[uid]['prompt_input'] = prompt_input
                if vision_input:
                    data[uid]['vision_input'] = vision_input
            
            # Write the updated data to the file
            with open(output_file_name, 'w') as json_file:
                json.dump(data, json_file, indent=4)
            
            print(f"Data was successfully updated in {output_file_name}")
            return True
        except json.JSONDecodeError as e:
            print(f"An error occurred while parsing the JSON content: {e}")
            return False
    else:
        print("The required delimiters were not found correctly in the string.")
        return False


def check_key_in_json(file_path, key):
    try:
        with open(file_path, 'r') as json_file:
            data = json.load(json_file)
        
        # Check if the key exists at the top level of the JSON structure
        if key in data:
            return True
        else:
            return False
    except FileNotFoundError:
        print(f"The file {file_path} was not found.")
    except json.JSONDecodeError as e:
        print(f"Error reading {file_path}: {e}")
    except Exception as e:
        print(f"An error occurred with {file_path}: {e}")
    return False