"""Find the line of code generated by the model using the block in the version code"""

import json
import os
import random
import re


def process_line_mask(code_snippet, core_token):
    if not core_token:
        return None, None

    replaced_lines = {}
    lines = code_snippet.split('\n')

    in_multi_line_comment = False

    for i, line in enumerate(lines):
        if in_multi_line_comment:
            if ('"""' in line or "'''" in line) and not re.findall(
                r"'''(.*?)'''|\"\"\"(.*?)\"\"\"", line
            ):
                in_multi_line_comment = False
            continue
        elif line.strip().startswith('#'):
            continue
        elif re.findall(r"'''(.*?)'''|\"\"\"(.*?)\"\"\"", line):
            continue
        elif ('"""' in line or "'''" in line) and not re.findall(
            r"'''(.*?)'''|\"\"\"(.*?)\"\"\"", line
        ):
            in_multi_line_comment = True
            continue
        else:
            if re.search(r'\bdef\s+task_function\b', line):
                continue

            if re.search(r'\b{}\b(?!\s*=)'.format(re.escape(core_token)), line):
                replaced_lines.update({i: line})

    if replaced_lines:
        random_line_location = random.choice(list(replaced_lines.keys()))

        masked_line = lines[random_line_location]
        leading_spaces = re.match(r'^\s*', masked_line).group(0)
        masked_line = masked_line.strip()
        lines[random_line_location] = leading_spaces + '<line_mask>'

        masked_code = '\n'.join(lines)

        return masked_code, masked_line

    return None, None


def load_json(file_path):
    with open(file_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    return data


def save_json(file_path, data):
    with open(file_path, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=4)


if __name__ == '__main__':
    model_list = os.listdir('../data/result_data/block_completion')
    for model in model_list:
        input_json_file = f'../data/result_data/block_completion/{model}/VersiCode_block_completion.json'
        output_json_file = input_json_file
        data = load_json(input_json_file)

        for item in data:
            core_token = item['core_token']
            code = item['code']

            _, core_line_in_code = process_line_mask(code, core_token)
            if core_line_in_code:
                item['core_line_in_code'] = core_line_in_code
            else:
                item['core_line_in_code'] = 'N/A'

            model_output_clear = item['model_output_clear']
            core_line_in_output_list = []

            for entry in eval(model_output_clear):
                _, core_line_in_output = process_line_mask(entry, core_token)
                if core_line_in_output:
                    core_line_in_output_list.append(core_line_in_output)
                else:
                    core_line_in_output_list.append('N/A')

            item['core_line_in_output_clear'] = core_line_in_output_list

        save_json(output_json_file, data)
        print('Done!')
