from radon.complexity import cc_visit
import os
import json
import textwrap

def get_cyclomatic_complexity(nodeid, source_code: str) -> int:
    """Returns the cyclomatic complexity of the first function defined in the code."""
    if nodeid.split('::')[0].endswith('.rst'):
        return -1
    try:
        blocks = cc_visit(source_code)
        if len(blocks) == 1:
            return blocks[0].complexity
        else:
            print(f'blocks more than 1: {nodeid}')
            complex_sum = 0
            for block in blocks:
                complex_sum += block.complexity
            return complex_sum / len(blocks)
    except Exception as e:
        print(f'Exception when calculating complexity: {nodeid}: {e}')
        return -1  # fallback if no function is found


if __name__ == "__main__":
    V4_save_path = 'repocod_data_v7'
    for idx, task_id in enumerate(os.listdir(V4_save_path)):
        if "pylint" not in task_id:
            continue
        print(f'============{idx}. {task_id} =============')
        task_json_path = os.path.join(V4_save_path, task_id)
        with open(task_json_path, 'r') as f:
            task_obj = json.load(f)
        call_distance_test = task_obj['call_distance_test']
        for call_dis, test_list in call_distance_test.items():
            for test in test_list:
                nodeid = test['base_nodeid']
                src_code = test['src_code']
                src_code = textwrap.dedent(src_code)
                # src_code = f"def _():\n{textwrap.indent(textwrap.dedent(src_code), '    ')}"
                complexity = get_cyclomatic_complexity(nodeid, src_code)
                test['complexity'] = complexity
        with open(task_json_path, 'w') as f:
            json.dump(task_obj, f, indent=4)
        # exit()