# v1

import re
import os
import json

import pandas as pd

from typing import List, Dict, Any, Optional
from pathlib import Path

from utils.data_contracts import ProjectStructure

# from data_contracts import ProjectStructure


TOP_N = 20
CONTEXT_LIMIT = 2000
FEAT_LIMIT = 60000
SOURCE_FILENAME_REGEX = re.compile(r'^\s*source_filename\s*=\s*"([^"]+)"', re.MULTILINE)


import difflib


def find_closest_match(mangled_name, origin_names, threshold=0.7):
    """
    

    Args:
        mangled_name: mangled
        origin_names: 
        threshold: 

    Returns:
        None
    """
    if not origin_names:
        return None

    # 
    matches = []
    for name in origin_names:
        # SequenceMatcher
        similarity = difflib.SequenceMatcher(None, mangled_name, name).ratio()
        matches.append((name, similarity))

    # 
    matches.sort(key=lambda x: x[1], reverse=True)

    # 
    best_match, best_similarity = matches[0]

    if best_similarity >= threshold:
        return best_match
    else:
        return None


def parse_linked_ir_to_map(linked_ir_text: str) -> Dict[str, Dict[str, Any]]:
    """
     LLVM IR  IR 

    Args:
        linked_ir_text:  linked.ll 

    Returns:
         Mangled Name  'lines'  'content'
    """
    function_map: Dict[str, Dict[str, Any]] = {}
    current_func_name: Optional[str] = None
    current_func_ir: List[str] = []
    in_function_body = False

    for line in linked_ir_text.splitlines():
        line_stripped = line.strip()

        # 1.  (define)
        if line_stripped.startswith("define "):
            parts = line_stripped.split("@", 1)
            if len(parts) > 1:
                full_name_and_attrs = parts[1].split("(", 1)[0].strip()
                current_func_name = full_name_and_attrs.split(" ", 1)[0].strip()

                in_function_body = True
                current_func_ir = [line]  #  define 

                # 
                if line_stripped.endswith("}"):
                    function_map[current_func_name] = {"lines": 1, "content": line}
                    in_function_body = False
                    current_func_name = None
                    current_func_ir = []
                    continue

        # 2. 
        elif in_function_body:
            current_func_ir.append(line)

            # 3.  (})
            if line_stripped == "}":
                if current_func_name:
                    function_map[current_func_name] = {
                        "lines": len(current_func_ir),
                        "content": "\n".join(current_func_ir),
                    }

                # 
                in_function_body = False
                current_func_name = None
                current_func_ir = []

    return function_map


def get_source_filename_from_ir(ir_file_path: Path) -> Optional[str]:
    try:
        # 
        with open(ir_file_path, "r", encoding="utf-8") as f:
            head = f.read(2048)  #  2KB

        match = SOURCE_FILENAME_REGEX.search(head)
        if match:
            return match.group(1)

        #  source_filename fallback
        return ir_file_path.name.replace(".ll", "")

    except Exception as e:
        print(f":  {ir_file_path} : {e}")
        return None


def convert_llvm_dot_to_json(bench_path: str) -> Dict[str, Any]:
    """
     LLVM Call Graph (.dot )  JSON 
    JSON : { "function_name": ["callee_1", "callee_2", ...], ... }

    Args:
        dot_content:  DOT 

    Returns:
         JSON 
    """
    dot_name = f"{bench_path}/IR/linked.ll.callgraph.dot"
    dot_content = ""
    with open(dot_name, "r") as f:
        dot_content = f.read()

    # 1. 
    # ID: Node0x... [...,label="{FunctionName}"]
    NODE_PATTERN = re.compile(r'(Node0x[0-9a-f]+)\s+\[[^\]]*label="\{(.*?)\}"\];')
    # : Node0x... -> Node0x...;
    EDGE_PATTERN = re.compile(r"(Node0x[0-9a-f]+)\s*->\s*(Node0x[0-9a-f]+);")

    # 2. :
    # node_map:  Node ID 
    # { 'Node0x5608c2a4d000': 'BZ2_blockSort', ... }
    node_map: Dict[str, str] = {}

    # call_graph:  (JSON)
    # { 'BZ2_blockSort': ['fallbackSort', 'mainSort', ...], ... }
    call_graph: Dict[str, List[str]] = {}

    # ---  1:  () ---
    for match in NODE_PATTERN.finditer(dot_content):
        node_id = match.group(1)
        func_name = match.group(2)
        node_map[node_id] = func_name

        # 
        call_graph[func_name] = []

    # ---  2:  () ---
    for match in EDGE_PATTERN.finditer(dot_content):
        caller_id = match.group(1)
        callee_id = match.group(2)

        # 
        caller_name = node_map.get(caller_id)
        callee_name = node_map.get(callee_id)

        # 
        if caller_name and callee_name:
            if callee_name not in call_graph[caller_name]:
                call_graph[caller_name].append(callee_name)

    # 
    return {"call_graph": call_graph}


def get_structural_source_func_names(bench_path: str, function_list):

    all_func_names_set = set(function_list)
    ir_dir_path = f"{bench_path}/IR/base"
    all_structural_map = {"bench_path": bench_path, "proj_struct": []}

    call_graph = convert_llvm_dot_to_json(bench_path)

    # 1.  .ll 
    for file_name in os.listdir(ir_dir_path):
        if not file_name.endswith(".ll"):
            continue

        ir_file_path = Path(ir_dir_path) / file_name

        # 2.  (Source File Name)
        source_filename = get_source_filename_from_ir(ir_file_path)
        if not source_filename:
            source_filename = file_name  # 

        # 3.  IR 
        current_file_functions: List[str] = []
        try:
            ir_content = ir_file_path.read_text(encoding="utf-8")

            #  IR 
            for func_name in all_func_names_set:
                # :  IR  "define ... @<func_name>("
                #  C  "define" 
                search_term = f"define "  #  'define '

                #  'define'  '@<func_name>' 
                # if f"define {func_name}(" in ir_content or f"define dso_local {func_name}(" in ir_content:
                # current_file_functions.append(func_name)
                # : 
                # f"define .*?{func_name}\(" in ir_content
                pattern = rf"define\s+.*?@{func_name}\s*\("
                if re.search(pattern, ir_content):
                    func_call = {
                        "name": func_name,
                        "callee": call_graph["call_graph"][func_name],
                    }
                    current_file_functions.append(func_call)

            if current_file_functions:
                all_structural_map["proj_struct"].append(
                    {"filename": source_filename, "functions": current_file_functions}
                )

        except Exception as e:
            print(f":  {ir_file_path} : {e}")

    for func_name in all_func_names_set:
        all_structural_map

    return all_structural_map


def select_high_potential_functions(
    all_function_data: pd.DataFrame,
    # n: int = TOP_N,
    function_map: Dict[str, Dict[str, Any]],
    context_limit: int = CONTEXT_LIMIT,
    feature_json_limit: int = FEAT_LIMIT,
) -> List[Dict[str, Any]]:
    """
     Top N 
    """
    context_sum = 0
    origin_funcs = list(function_map.keys())

    # if context_sum < context_limit:
    #    return all_functions

    selected_functions = []
    selected_extend = []

    cur_context = 0

    # LLM
    all_function_data = all_function_data.sort_values(by="score", ascending=False)
    for i in range(len(all_function_data["name"])):
        func = all_function_data["name"][i]
        selected_functions.append(func)
        try:
            cur_context += function_map[func]["lines"]
        except Exception as e:
            print(e)
            find_closest_match(func, origin_funcs)
            i -= 1

        print(f"appended function: {func}, current length: {cur_context} lines. ")
        if cur_context > context_limit:
            break

    # 
    # names = all_function_data.loc[all_function_data["score"] > 70, "name"]
    names = all_function_data.loc[all_function_data["score"] > 50, "name"]
    selected_extend = names.tolist()
    selected_extend = list(set(selected_extend + selected_functions))

    return selected_functions, selected_extend


def select_irs(
    bench_path: str,
    structural_sources: Dict[str, Dict[str, Any]],
    potential_funcs: List[str],
) -> List[Dict[str, Any]]:
    """
    IR
    """
    base_ir_folder = Path(f"{bench_path}/IR/base")
    all_ir_files = set(str(filepath) for filepath in base_ir_folder.glob("*.ll"))
    # all_c_files = set()
    found_functions_map = {}

    #  IR 
    potential_cs = set()  # IR ()
    other_cs = None  # IR ()

    # 
    print(structural_sources)
    for file_info in structural_sources:
        filename = file_info.get("filename")
        functions = file_info.get("functions", [])

        # 
        # all_c_files.add(filename)

        # 
        is_target_file = False

        # 
        for func in functions:
            func_name = func.get("name")

            # 
            if func_name in potential_funcs:
                # 
                if func_name not in found_functions_map:
                    found_functions_map[func_name] = filename

                # 
                is_target_file = True

        if is_target_file:
            # 
            potential_cs.add(filename)

    # other_cs = all_c_files - target_ir_files
    # other_cs = all_c_files.difference(potential_cs)

    potential_irs = set()
    other_irs = set()

    for cfile in potential_cs:
        potential_irs.add(
            f"{bench_path}/IR/base/{(cfile.split('/')[-1]).split('.')[0]}.ll"
        )

    # 3. 
    other_irs = all_ir_files.difference(potential_irs)

    return found_functions_map, potential_irs, other_irs, all_ir_files


def extract_function_ir(
    potential_funcs: List[str], ir_function_map: Dict[str, Dict[str, Any]]
) -> str:
    """
     Map  IR 

    Args:
        potential_funcs:  ( IR  Mangled Name)
        ir_function_map:  IR 

    Returns:
         IR  Analysis Agent 
    """

    ir_snippets = []

    for name in potential_funcs:
        data = ir_function_map.get(name)
        if data:
            ir_snippets.append(
                f"; --- Function: {name} (Lines: {data['lines']}) ---\n{data['content']}"
            )
        else:
            print(
                f"Warning: function '{func_name}' is not contained in the function map."
            )

    return "\n\n".join(ir_snippets)


def extract_function_json(
    potential_funcs: List[str], origin_json: List[Dict[str, Any]]
) -> List[Dict[str, Any]]:
    """
    /remarks JSON 

    Args:
        potential_funcs:  (e.g., ["getbit", "bit_count"]).
        origin_json: remarks JSON  (1  2).

    Returns:
        remarks JSON 
    """

    # 1.  JSON O(1)
    #  O(N) 
    function_map = {
        item.get("name"): item
        for item in origin_json
        if item.get("name")  #  name 
    }

    # 2. 
    potential_data = []

    for func_name in potential_funcs:
        data = function_map.get(func_name)
        if data:
            potential_data.append(data)
        else:
            # 
            print(
                f"Warning: function '{func_name}' is not contained in the origin json."
            )

    return potential_data


if __name__ == "__main__":
    # ---  ---
    # 1:  features  remarks 
    #    json_structure_1 = [
    #      { "name": "alloc_bit_array", "features": {"callee": {...}}, "remarks": {...} },
    #      { "name": "getbit", "features": {"callee": {...}}, "remarks": {...} },
    #      { "name": "setbit", "features": {"callee": {...}}, "remarks": {} },
    #    ]
    #
    #    # 2:  remarks 
    #    json_structure_2 = [
    #      { "name": "alloc_bit_array", "remarks": {...} },
    #      { "name": "getbit", "remarks": {...} },
    #      { "name": "setbit", "remarks": {} },
    #      { "name": "flipbit", "remarks": {} },
    #      { "name": "bit_count", "remarks": {"analysis": {"loop-vectorize": [...]}} },
    #    ]
    #
    #    ir_content = None
    ir_file_path = "<project_path>/agents/test/bitcount/IR/linked.ll"
    with open(ir_file_path, "r") as ir:
        # print(ir.read()[:-50])
        ir_content = ir.read()
    #
    #    #  IR 
    function_map = parse_linked_ir_to_map(ir_content)
    #    #print(function_map)
    #    for func in function_map.keys():
    #        print(function_map[func]["lines"])
    #        print(function_map[func]["content"])
    #
    #    # / Top N 
    top_n_funcs = ["getbit", "alloc_bit_array"]
    #    # 
    #
    #    # 
    #    result_1 = extract_function_json(top_n_funcs, json_structure_1)
    #    print("---  1  ---")
    #    print(f" {len(json_structure_1)}  {len(result_1)} :")
    #    print(result_1)
    #    print("\n" + "="*50 + "\n")
    #
    #    result_2 = extract_function_json(top_n_funcs, json_structure_2)
    #    print("---  2  ---")
    #    print(f" {len(json_structure_2)}  {len(result_2)} :")
    #    print(result_2)
    #
    #    #  IR 
    potential_ir = extract_function_ir(top_n_funcs, function_map)
    #    print(potential_ir)
    #
    #    print(len(function_map.keys()))
    #
    #    struct = get_structural_source_func_names("<project_path>/agents/bitcount", list(function_map.keys()))
    #    struct = ProjectStructure.model_validate_json(json.dumps(struct))
    #    print(type(struct.json()))

    tmp_json = [
        {
            "filename": "./src_work/bitcnt_4.c",
            "functions": [
                {"name": "ntbl_bitcnt", "callee": ["ntbl_bitcnt"]},
                {"name": "btbl_bitcnt", "callee": ["btbl_bitcnt"]},
            ],
        },
        {
            "filename": "./src_work/bitarray.c",
            "functions": [
                {"name": "setbit", "callee": []},
                {"name": "flipbit", "callee": []},
                {"name": "alloc_bit_array", "callee": ["calloc"]},
                {"name": "getbit", "callee": []},
            ],
        },
        {
            "filename": "./src_work/bitcnt_3.c",
            "functions": [
                {"name": "BW_btbl_bitcount", "callee": []},
                {"name": "ntbl_bitcount", "callee": []},
                {"name": "AR_btbl_bitcount", "callee": []},
            ],
        },
        {
            "filename": "./src_work/bitcnt_2.c",
            "functions": [{"name": "bitcount", "callee": []}],
        },
        {
            "filename": "./src_work/bitstrng.c",
            "functions": [{"name": "bitstring", "callee": []}],
        },
        {
            "filename": "./src_work/bstr_i.c",
            "functions": [{"name": "bstr_i", "callee": ["strchr"]}],
        },
        {
            "filename": "./src_work/loop-wrap.c",
            "functions": [
                {
                    "name": "main",
                    "callee": [
                        "fopen",
                        "fprintf",
                        "__isoc99_fscanf",
                        "fclose",
                        "main1",
                    ],
                }
            ],
        },
        {
            "filename": "./src_work/bitfiles.c",
            "functions": [
                {"name": "bfopen", "callee": ["malloc", "fopen", "free"]},
                {"name": "bfread", "callee": ["fgetc"]},
                {"name": "bfclose", "callee": ["fclose", "free"]},
                {"name": "bfwrite", "callee": ["fputc"]},
            ],
        },
        {
            "filename": "./src_work/bitcnts.c",
            "functions": [
                {
                    "name": "main1",
                    "callee": ["fprintf", "exit", "atoi", "puts", "printf"],
                },
                {"name": "bit_shifter", "callee": []},
            ],
        },
        {
            "filename": "./src_work/bitcnt_1.c",
            "functions": [{"name": "bit_count", "callee": []}],
        },
    ]

    aa, bb, cc, dd = select_irs(
        "<project_path>/agents/benchmarks/cbench/automotive_bitcount",
        tmp_json,
        ["alloc_bit_array"],
    )
    print(aa)
    print(bb)
    print(cc)
    print(dd)
