#!<python_interpreter_path>

import json
import argparse
import os
from collections import defaultdict

# PyYAML
try:
    import yaml
except ImportError:
    # PyYAML
    print("Merge Remarks Error: Need to install PyYAML lib. Run 'pip install PyYAML' to install.")
    # raisesys.exit
    raise ImportError("PyYAML library is required.")

# PWD = os.path.dirname(__file__)
# LINKED_JSON_FILE = f'{PWD}/jsons/origin.json'
# REMARKS_YAML_FILE = f'{PWD}/remarks/o3.remarks.yaml'
# OUTPUT_JSON_FILE = f'{PWD}/origin_with_o3remarks.json'

PATH = "./test"
PATH = "./bitcount"
LINKED_JSON_FILE = f"{PATH}/jsons/origin.json"
REMARKS_YAML_FILE = f"{PATH}/remarks/o3.remarks.yaml"
OUTPUT_JSON_FILE = f"{PATH}/jsons/origin_with_o3remarks_order.json"


def extract_reason_string(args_list):
    """
     LLVM Remark  Arg 

    [{"Callee": "free"}, {"String": " will not be inlined into "}]
    -> "free will not be inlined into "
    """
    reason_parts = []
    if args_list:
        for arg in args_list:
            #  Arg
            #
            for value in arg.values():
                reason_parts.append(str(value))
    return "".join(reason_parts)


def process_remarks_yaml(filepath):
    """
    O3Remarks YAML (analysis/missed/passed)  PassName 
     [start, end] 
    Args  ReasonString 

    :param filepath: O3YAML
    :return: 
    """
    # remarks_data : {func_name: {'analysis': {pass_name: [merged_remark_entry, ...]}, 'missed': {...}, 'passed': {...}}}
    # merged_remark_entry : {'count': N, 'Name': ..., 'ReasonString': ...}
    remarks_data = defaultdict(
        lambda: {
            "missed": defaultdict(list),
            "passed": defaultdict(list),
            "analysis": defaultdict(list),
        }
    )

    try:
        with open(filepath, "r", encoding="utf-8") as f:
            content = f.read()
    except FileNotFoundError:
        print(
            f"Merge Remarks Error: File {filepath} not found. Check '{filepath}' path"
        )
        return remarks_data

    # YAML '---'  !Analysis, !Missed  !Passed
    yaml_documents = content.split("---")
    order_counter = 0

    for doc_string in yaml_documents:
        #
        doc_string = doc_string.strip()
        if not doc_string:
            continue

        order_counter += 1
        remark_type = None

        #
        if doc_string.startswith("!Missed"):
            remark_type = "missed"
            doc_string = doc_string[len("!Missed"):].strip()
        elif doc_string.startswith("!Passed"):
            remark_type = "passed"
            doc_string = doc_string[len("!Passed"):].strip()
        elif doc_string.startswith("!AnalysisFPCommute"):
            remark_type = "analysis"
            doc_string = doc_string[len("!AnalysisFPCommute"):].strip()
        elif doc_string.startswith("!Analysis"):
            remark_type = "analysis"
            doc_string = doc_string[len("!Analysis"):].strip()
        else:
            #
            continue

        try:
            # yaml.safe_loadYAML
            remark = yaml.safe_load(doc_string)
            if remark is None:
                continue

            #  'Function' ()  'Pass' (Pass)
            if "Function" in remark and "Pass" in remark:
                func_name = remark["Function"]
                pass_name = remark["Pass"]

                # ---  Args  ---
                args_list = remark.get("Args", [])
                reason_string = extract_reason_string(args_list)

                #
                #  ReasonString
                current_remark_details = {
                    "Name": remark.get("Name"),
                    "ReasonString": reason_string,
                }

                #  (, , Pass)
                target_list = remarks_data[func_name][remark_type][pass_name]

                #
                continue_flag = False
                if target_list:
                    for item in target_list:
                        exist_comparison_details = {
                            "Name": item["Name"],
                            "ReasonString": item["ReasonString"],
                        }
                        if current_remark_details == exist_comparison_details:
                            #
                            item["count"] += 1
                            continue_flag = True
                            break
                    if continue_flag:
                        continue

                #
                new_merged_entry = {
                    "count": 1,
                    "Name": current_remark_details["Name"],
                    "ReasonString": current_remark_details[
                        "ReasonString"
                    ],  #
                }
                target_list.append(new_merged_entry)

        except yaml.YAMLError as e:
            print(doc_string)
            print(
                f"Merge Remarks Warning: Parse YAML file failed (order: {order_counter}): {e}")
            continue

    # defaultdictdictJSON
    final_remarks_data = {}
    for func, types in remarks_data.items():
        #  defaultdict  dict
        final_remarks_data[func] = {
            "analysis": dict(types["analysis"]),
            "missed": dict(types["missed"]),
            "passed": dict(types["passed"]),
        }

    return final_remarks_data


def merge_remarks_into_json(json_filepath, remarks_data, del_features: bool):
    """
    JSON
    """
    try:
        with open(json_filepath, "r", encoding="utf-8") as f:
            data = json.load(f)
    except FileNotFoundError:
        print(f"Merge Remarks Error: File {json_filepath} not found.")
        return None
    except json.JSONDecodeError as e:
        print(
            f"Merge Remarks Error: Parse JSON file {json_filepath} failed: {e}")
        return None

    # JSON
    if isinstance(data, list):
        for func_entry in data:
            if del_features:
                del func_entry["features"]
            func_name = func_entry.get("name")

            if func_name and func_name in remarks_data:
                #
                func_entry["remarks"] = remarks_data[func_name]
            else:
                #
                func_entry["remarks"] = {
                    "missed": {}, "passed": {}, "analysis": {}}

    return data


def merge_json(remarks: str, static_json: str, output_json):
    remarks_data = process_remarks_yaml(remarks)
    if not remarks_data:
        print("Merge Remarks Error: Failed to process the remarks file or the file is empty. Script terminated.")
        return
    merged_data = merge_remarks_into_json(static_json, remarks_data, True)

    if merged_data is None:
        print("Data merge failed. Script terminated.")
        return

    try:
        with open(output_json, "w", encoding="utf-8") as f:
            # JSON
            json.dump(merged_data, f, indent=2, ensure_ascii=False)
        print(
            f"Merge Remarks Operation successful! The processed data has been saved to the file: {output_json}")
        # print(merged_data)
        return json.dumps(merged_data)
    except IOError as e:
        print(
            f"Merge Remarks Error: Failed to write to file {OUTPUT_JSON_FILE}: {e}")


def main(remarks: str, static_json: str, output_json: str):
    """

    """
    print(f"--- STEP 1: Processing optimization remarks file: {remarks} ---")
    remarks_data = process_remarks_yaml(remarks)

    if not remarks_data:
        print("Failed to process the remarks file or the file is empty. Script terminated.")
        return

    #  Pass  Pass
    print(
        f"   The remarks file has been processed, involving {len(remarks_data)} functions. Adjacent duplicate entries have been merged.")

    print(
        f"--- STEP 2: Merging data into static feature files: {static_json} ---")
    merged_data = merge_remarks_into_json(static_json, remarks_data, False)

    if merged_data is None:
        print("Data merge failed. Script terminated.")
        return

    print(f"--- STEP 3: Saving results to: {output_json} ---")
    try:
        with open(output_json, "w", encoding="utf-8") as f:
            # JSON
            json.dump(merged_data, f, indent=2, ensure_ascii=False)
        print(
            f"Operation successful! The processed data has been saved to the file: {output_json}")
    except IOError as e:
        print(
            f"Merge Remarks Error: Failed to write to file {output_json}: {e}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Process and merge LLVM optimization remarks.")
    parser.add_argument("--bench_path", type=str, required=True,
                        help="Path to the benchmark directory.")
    args = parser.parse_args()

    static_json = f"{args.bench_path}/jsons/origin.json"
    remarks = f"{args.bench_path}/remarks/o3.remarks.yaml"
    output_json = f"{args.bench_path}/jsons/origin_with_o3remarks_order.json"

    main(remarks=remarks, static_json=static_json, output_json=output_json)
