# v1

import os
import json
import cxxfilt
import re
import logging
import shutil
import argparse
import sys
import datetime
import time
import subprocess

import pandas as pd
import fasteners

from pathlib import Path
from crewai import Task, Crew, Process, TaskOutput

from utils.select_potential import *
from utils.agents import *
from utils.robust_control import extract_clean_json_string, correct_llvm_pipeline

from utils.tools import (
    _init_features,
    _get_IR,
    _run_init,
    _get_new_remarks,
    _check_pipeline_brackets,
    _compile_and_run,
    _compile_and_run_2_stage,
    _get_platform_info
)

from utils.tasks import (
    create_function_score_tasks,
    create_evaluation_tasks,
    create_optimization_tasks,
    create_initial_analysis_tasks,
)

#  Pydantic
# import Pydantic data contract
from utils.data_contracts import (
    ProjectStructure,
    DevAttentionScores,
    OptimizationInput,
    OptimizationDecision,
    EvaluationReport,
    PlatformInfo,
    RuntimePerformance,
    StaticFeatAndRemarksItem,
    FunctionFeatures,
    GenericFeatures,
    OptimizationInputFeedback,
    Remarks,
    FunctionRemarks,
    IRContent,
)


parser = argparse.ArgumentParser(description="Main script to run AutoPass")

parser.add_argument(
    "-p", "--benchpath", type=str, required=True, help="| Path of benchmarks."
)
parser.add_argument(
    "-f", "--ldflags", type=str, required=True, help="| Linking flags for ld."
)
parser.add_argument(
    "-r", "--runflag", type=str, required=True, help="| Flags of running programs (e.g. inputs, outputs)."
)
parser.add_argument(
    "-s", "--savehistory", type=str, required=True, help="| Path of history to be saved."
)
parser.add_argument(
    "--target", type=str, required=False, help="Build target.", default=""
)

args = parser.parse_args()
BENCH_PATH = args.benchpath
FLAGS = args.ldflags
RUNFLAG = args.runflag
LOG_DIR = args.savehistory
TARGET = args.target

# <--- LOCK SET IN LOG_DIR
EXECUTION_LOCK_FILE = f"{LOG_DIR}/.autopass_executor.lock"

dirs = [
    "history",
    "perf_jsons",
    "IR",
    "IR/base",
    "bin",
    "obj",
    "remarks",
    "features",
    "jsons",
]

if TARGET == "":
    PLATFORM_INFO = PlatformInfo(
        target_arch="GNU/Linux",
        compiler="clang 17.0.6",
        cpu="Cortex-A76",
        freq="3.5 GHz",
        L1dcache="256 KiB",
        L1icache="256 KiB",
        L2cache="2 MiB",
        L3cache="2 MiB"
    )
else:
    PLATFORM_INFO = PlatformInfo.model_validate(
        _get_platform_info(target=TARGET))

MAX_ROUNDS = 3  # optimization times
MAX_DECISION_RETRY = 5  # retry times of Reasoning Agent
CONTEXT_LIMIT = 3000
# FEAT_LIMIT = 3000

INITIAL_PIPELINE = "annotation2metadata,forceattrs,inferattrs,coro-early,function<eager-inv>(lower-expect,simplifycfg<bonus-inst-threshold=1;no-forward-switch-cond;no-switch-range-to-icmp;no-switch-to-lookup;keep-loops;no-hoist-common-insts;no-sink-common-insts;speculate-blocks;simplify-cond-branch>,sroa<modify-cfg>,early-cse<>,callsite-splitting),openmp-opt,ipsccp,called-value-propagation,globalopt,function<eager-inv>(mem2reg,instcombine<max-iterations=1000;no-use-loop-info>,simplifycfg<bonus-inst-threshold=1;no-forward-switch-cond;switch-range-to-icmp;no-switch-to-lookup;keep-loops;no-hoist-common-insts;no-sink-common-insts;speculate-blocks;simplify-cond-branch>),require<globals-aa>,function(invalidate<aa>),require<profile-summary>,cgscc(devirt<4>(inline<only-mandatory>,inline<inline-threshold=500>,function-attrs<skip-non-recursive>,argpromotion,openmp-opt-cgscc,function<eager-inv;no-rerun>(sroa<modify-cfg>,early-cse<memssa>,speculative-execution,jump-threading,correlated-propagation,simplifycfg<bonus-inst-threshold=1;no-forward-switch-cond;switch-range-to-icmp;no-switch-to-lookup;keep-loops;no-hoist-common-insts;no-sink-common-insts;speculate-blocks;simplify-cond-branch>,instcombine<max-iterations=1000;no-use-loop-info>,aggressive-instcombine,constraint-elimination,libcalls-shrinkwrap,tailcallelim,simplifycfg<bonus-inst-threshold=1;no-forward-switch-cond;switch-range-to-icmp;no-switch-to-lookup;keep-loops;no-hoist-common-insts;no-sink-common-insts;speculate-blocks;simplify-cond-branch>,reassociate,loop-mssa(loop-instsimplify,loop-simplifycfg,licm<no-allowspeculation>,loop-rotate<header-duplication;no-prepare-for-lto>,licm<allowspeculation>,simple-loop-unswitch<nontrivial;trivial>),simplifycfg<bonus-inst-threshold=1;no-forward-switch-cond;switch-range-to-icmp;no-switch-to-lookup;keep-loops;no-hoist-common-insts;no-sink-common-insts;speculate-blocks;simplify-cond-branch>,instcombine<max-iterations=1000;no-use-loop-info>,loop(loop-idiom,indvars,loop-deletion,loop-unroll-full),sroa<modify-cfg>,vector-combine,mldst-motion<no-split-footer-bb>,gvn<>,sccp,bdce,instcombine<max-iterations=1000;no-use-loop-info>,jump-threading,correlated-propagation,adce,memcpyopt,dse,move-auto-init,loop-mssa(licm<allowspeculation>),coro-elide,simplifycfg<bonus-inst-threshold=1;no-forward-switch-cond;switch-range-to-icmp;no-switch-to-lookup;keep-loops;hoist-common-insts;sink-common-insts;speculate-blocks;simplify-cond-branch>,instcombine<max-iterations=1000;no-use-loop-info>),function-attrs,function(require<should-not-run-function-passes>),coro-split)),deadargelim,coro-cleanup,globalopt,globaldce,elim-avail-extern,rpo-function-attrs,recompute-globalsaa,function<eager-inv>(float2int,lower-constant-intrinsics,chr,loop(loop-rotate<header-duplication;no-prepare-for-lto>,loop-deletion),loop-distribute,inject-tli-mappings,loop-vectorize<no-interleave-forced-only;no-vectorize-forced-only;>,loop-load-elim,instcombine<max-iterations=1000;no-use-loop-info>,simplifycfg<bonus-inst-threshold=1;forward-switch-cond;switch-range-to-icmp;switch-to-lookup;no-keep-loops;hoist-common-insts;sink-common-insts;speculate-blocks;simplify-cond-branch>,slp-vectorizer,vector-combine,instcombine<max-iterations=1000;no-use-loop-info>,loop-unroll<O3>,transform-warning,sroa<preserve-cfg>,instcombine<max-iterations=1000;no-use-loop-info>,loop-mssa(licm<allowspeculation>),alignment-from-assumptions,loop-sink,instsimplify,div-rem-pairs,tailcallelim,simplifycfg<bonus-inst-threshold=1;no-forward-switch-cond;switch-range-to-icmp;no-switch-to-lookup;keep-loops;no-hoist-common-insts;no-sink-common-insts;speculate-blocks;simplify-cond-branch>),globaldce,constmerge,cg-profile,rel-lookup-table-converter,function(annotation-remarks),verify"


# log configuration
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("CompilationCrew")


if BENCH_PATH and FLAGS and RUNFLAG and LOG_DIR:
    print(
        f"\n\n\n\n\nbench path: {BENCH_PATH}\nldflags: {FLAGS}\nrun flag: {RUNFLAG}\n\n\n\n"
    )
    #
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )
    logger = logging.getLogger(
        f"CompilationCrew_{BENCH_PATH.split('/')[-1]}"
    )

    # initialize execution lock file
    logger.info(f"--- : {EXECUTION_LOCK_FILE} ---")
    try:
        exec_lock = fasteners.InterProcessLock(EXECUTION_LOCK_FILE)
    except Exception as e:
        sys.exit(f": {e}")

    if not Path(BENCH_PATH).is_dir():
        print(f"error: Benchmark directory does not exist: {BENCH_PATH}")
        exit(1)

    for d in dirs:
        try:
            # recursive=True
            # exist_ok=True
            target = Path(f"{BENCH_PATH}/{d}")
            if not target.is_dir():
                target.mkdir(exist_ok=True)

        except OSError as e:
            print(f"{d} mkdir error: {e}")
            exit(1)

    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    HISTORY = f"{BENCH_PATH}/history/history_{timestamp}.json"
    IR_FILE_PATH = f"{BENCH_PATH}/IR/linked.ll"  # IR file path

    RESULT = []

    logger.info(f"=======================================================")
    logger.info(
        f"---  The history will be stored in file: <bench>/history/history_{timestamp}.json  ---"
    )
    logger.info(f"=======================================================\n")

    current_round = 0
    initial_opt = "default<O3>"
    potential_funcs = None
    exceed_context_limit = False
    potential_irs = []
    other_irs = []
    best_time = 999999999999999999

    #
    # get execution lock
    logger.info(
        f"--- {BENCH_PATH.split('/')[-1]}:  (for _run_init)... ---"
    )
    with exec_lock:
        logger.info(
            f"--- {BENCH_PATH.split('/')[-1]}:  _run_init... ---"
        )

        logger.info(
            f"--- Round {current_round}: Initial Feature Extraction using {IR_FILE_PATH} ---"
        )

        initial_remarks_str = _init_features(BENCH_PATH, initial_opt)
        initial_performance = _run_init(BENCH_PATH, FLAGS, RUNFLAG, TARGET)

    logger.info(f"--- {BENCH_PATH.split('/')[-1]}: _run_init  ---")

    best_time = json.loads(initial_performance)["average_time_elapsed"]

    # print(initial_performance)
    lastround_performance = initial_performance

    initial_remarks_json = json.loads(initial_remarks_str)
    # print(type(initial_remarks_json))
    # print(initial_remarks_model)

    # Analysis Agent
    with exec_lock:
        ir_content_text = _get_IR(IR_FILE_PATH)

        function_map = parse_linked_ir_to_map(ir_content_text)

    with open(HISTORY, "w", encoding="utf-8") as f:
        json.dump(RESULT, f, indent=2, ensure_ascii=False)

    try:
        archive_path = f"{LOG_DIR}/history_{BENCH_PATH.split('/')[-1]}.json"
        shutil.copy2(HISTORY, archive_path)
        print(f"----- log saved to {archive_path}")
    except Exception as e:
        print(f"----- : {e}")
        pass

    # exit(0)

    if "-lstdc++" in FLAGS:
        is_cpp = True
    else:
        is_cpp = False

    if ir_content_text.count("\n") > CONTEXT_LIMIT:
        # long program branch
        exceed_context_limit = True
        logger.info(
            f"--- IR text is too long to process. Beginning potential functions selection. ---"
        )
        # exit(1)

        # developer attention score agent
        structural_sources = get_structural_source_func_names(
            BENCH_PATH, list(function_map.keys())
        )
        struct_bk = structural_sources.copy()
        structural_model = ProjectStructure.model_validate_json(
            json.dumps(structural_sources)
        )

        # print(f"{len(json.dumps(structural_sources))}")
        if len(json.dumps(structural_sources)) > 500000:
            for project in structural_sources["proj_struct"]:
                for func in project["functions"]:
                    del func["callee"]
                    # we can skip standard lib (external) functions
                    if is_cpp and ("_ZNS" in func["name"] or "_ZNSt" in func["name"] or "_ZSt" in func["name"] or "_ZNKSt" in func["name"] or "_cxx" in func["name"]):
                        del func["name"]

        RESULT.append(
            {
                "bench_path": BENCH_PATH,
                "proj_struct": (structural_sources["proj_struct"]),
            }
        )
        # print(structural_sources)
        # exit(1)

        funcs_attention_tasks = create_function_score_tasks(
            bench_path=BENCH_PATH, project_struct=json.dumps(
                structural_sources)
        )

        funcs_attention_crew = Crew(
            agents=[score_agent],
            tasks=funcs_attention_tasks,
            process=Process.sequential,
            verbose=1,
        )

        logger.info(
            "--- Starting Crew for Function Developer Attention Score Task. ---"
        )
        try:
            starttime = time.time()
            funcs_scores = funcs_attention_crew.kickoff()
            endtime = time.time()

            RESULT.append(
                {
                    "task_name": "function_potential",
                    "token_usage": [dict(funcs_scores.token_usage)],
                    "time_elapsed": [endtime - starttime],
                }
            )
            pass
        except Exception as e:
            logger.error(f"Initial Analysis Tasks (Task 1A) failed: {e}")
            exit(1)
        funcs_scores_json = extract_clean_json_string(funcs_scores)
        # print(json.loads(funcs_scores_json))

        with exec_lock:
            data = json.loads(funcs_scores_json)
            output_data = pd.DataFrame(data["scores"])
            # print(output_data)

            potential_funcs, extend_funcs = select_high_potential_functions(
                output_data, function_map, CONTEXT_LIMIT
            )
            # print(extend_funcs)
            # exit(0)
            _, potential_irs, other_irs, all_irs = select_irs(
                BENCH_PATH, struct_bk["proj_struct"], extend_funcs
            )

        if not other_irs:
            exceed_context_limit = False

        RESULT.append(funcs_scores_json)
        print(potential_funcs)

        RESULT.append(
            {
                "potential_funcs": list(potential_funcs),
                "potential_irs": list(potential_irs),
                "other_irs": list(other_irs)
            }
        )

        with exec_lock:
            ir_content_text_input = extract_function_ir(
                potential_funcs, function_map)
            # print(ir_content_text_input)

            # Analysis agent's report + static feat + remarks
            # get features and generate inputs for Reasoning Agent
            initial_remarks_json = extract_function_json(
                potential_funcs, initial_remarks_json
            )
            # print(initial_remarks_json)
            initial_remarks_input = json.dumps(initial_remarks_json)

    else:
        ir_content_text_input = ir_content_text
        initial_remarks_input = initial_remarks_str

    # ---  LLM :  OptimizationInput (Agent A) ---
    #  Analysis Task
    initial_tasks = create_initial_analysis_tasks(
        # ir_file=IR_FILE_PATH,
        ir_content=ir_content_text_input
    )

    initial_crew = Crew(
        agents=[analysis_agent],
        tasks=initial_tasks,
        process=Process.sequential,
        verbose=1,
    )

    logger.info("--- Starting Crew for Initial IR Content Analyzing ---")
    try:
        # print(ir_content_text)
        starttime = time.time()
        ir_analysis = initial_crew.kickoff()
        endtime = time.time()

        RESULT.append(
            {
                "task_name": "initial_analysis",
                "token_usage": [dict(ir_analysis.token_usage)],
                "time_elapsed": [endtime - starttime],
            }
        )
        pass

    except Exception as e:
        logger.error(f"Initial Analysis Tasks (Task 1A) failed: {e}")
        exit(1)

    logger.info(
        f"Initial OptimizationInput successfully generated. Starting iterative loop."
    )
    optimization_result = None
    performance = None

    with exec_lock:
        initial_input_model = OptimizationInput(
            ir_analysis_summary=ir_analysis.raw,  #
            static_feat_and_remarks=json.loads(initial_remarks_input),
            platform_info=PLATFORM_INFO,
            runtime_performance=json.loads(initial_performance),
            pipeline_params_to_optimize={
                "-passes": INITIAL_PIPELINE,
                "-unroll-count": 8,
                "-unroll-allow-partial": 0,
                "-unroll-threshold": 300,
                "-inline-threshold": 300,
                "-slp-threshold": -2,
                "-devirt<>": 4,
                "-mtriple": "",
                "-mcpu": "",
            },
        )
    current_input_json = initial_input_model.json()
    # print(current_input_json)
    best_parameters = {
        "-unroll-count": 8,
        "-unroll-allow-partial": 0,
        "-unroll-threshold": 300,
        "-inline-threshold": 300,
        "-slp-threshold": -2,
        "-devirt<>": 4,
    }  # initial

    RESULT.append(
        {
            "round": current_round,
            "initial_pipeline": initial_opt,
            **initial_input_model.model_dump()["runtime_performance"],
        }
    )
    with open(HISTORY, "w", encoding="utf-8") as f:
        json.dump(RESULT, f, indent=2, ensure_ascii=False)

    try:
        archive_path = f"{LOG_DIR}/history_{BENCH_PATH.split('/')[-1]}.json"
        shutil.copy2(HISTORY, archive_path)
        print(f"----- {archive_path}")
    except Exception as e:
        print(f"----- : {e}")
        pass

    # Feedback Loop
    decision_token_usage = []
    decision_time_usage = []
    for i in range(MAX_ROUNDS):
        current_round = i
        logger.info(f"=======================================================")
        logger.info(f"--- STARTING OPTIMIZATION ROUND {current_round + 1} ---")

        logger.info(
            f"Current input:\n{current_input_json[:20]}...{current_input_json[-20:]}"
        )
        logger.info(
            f"=======================================================\n")

        execution_success = False
        retry_attempt = 0
        optimization_decision = None

        while (not execution_success) and retry_attempt < MAX_DECISION_RETRY:

            # create task
            o_tasks = create_optimization_tasks(
                bench_path=BENCH_PATH,
                current_input=current_input_json,
                round_num=current_round + 1,
            )

            crew = Crew(
                agents=[reasoning_agent],
                tasks=o_tasks,
                process=Process.sequential,
                verbose=1,
            )

            try:
                # LLM generation
                starttime = time.time()
                optimization_result = crew.kickoff()
                endtime = time.time()

                decision_token_usage.append(
                    dict(optimization_result.token_usage))
                decision_time_usage.append(endtime - starttime)

            except Exception as e:
                logger.error(
                    f"Crew execution failed in Round {current_round + 1}: {e}")
                break

            optimization_decision = extract_clean_json_string(
                optimization_result)
            if not optimization_decision:
                logger.info(
                    f"\nError: CAN NOT extract json structure from Agent output:\n{optimization_decision}"
                )

            # print(optimization_decision)
            # print(type(optimization_decision))

            if not _check_pipeline_brackets(
                json.loads(optimization_decision)["passes_param_adjustments"][
                    "optimized_pipeline"
                ]
            ):
                continue

            #  1 ,  2
            can_execute = 2
            while can_execute > 0:
                execution_success = True

                #
                # execution
                logger.info(
                    f"--- {BENCH_PATH.split('/')[-1]}:  (for executor_agent Round {current_round + 1})... ---"
                )
                with exec_lock:
                    logger.info(
                        f"--- {BENCH_PATH.split('/')[-1]}:  execution... ---"
                    )
                    if exceed_context_limit:
                        execution_result = _compile_and_run_2_stage(
                            current_round + 1,
                            f"{BENCH_PATH}/IR/linked.ll",
                            json.dumps(
                                json.loads(optimization_decision)[
                                    "passes_param_adjustments"
                                ]
                            ),
                            FLAGS,
                            RUNFLAG,
                            potential_irs,
                            other_irs,
                            TARGET,
                        )

                    else:
                        # if potential_irs and not other_irs:
                        #    pass

                        execution_result = _compile_and_run(
                            current_round + 1,
                            f"{BENCH_PATH}/IR/linked.ll",
                            json.dumps(
                                json.loads(optimization_decision)[
                                    "passes_param_adjustments"
                                ]
                            ),
                            FLAGS,
                            RUNFLAG,
                            TARGET,
                        )

                performance = execution_result
                print(execution_result)
                print(json.loads(performance))

                if not performance:
                    logger.info(
                        f"\nError: CAN NOT extract json structure from Agent output:\n{performance}."
                    )
                    execution_success = False
                    break

                if "error" in performance:
                    logger.info(
                        f"\nError: Compilation error. Trying to repair pipeline once."
                    )
                    execution_success = False

                    if can_execute == 1:
                        # can not repair, re-generate
                        break

                    optimization_decision_json = json.loads(
                        optimization_decision)
                    repaired_pipeline = correct_llvm_pipeline(
                        INITIAL_PIPELINE,
                        optimization_decision_json["passes_param_adjustments"][
                            "optimized_pipeline"
                        ],
                    )
                    optimization_decision_json["passes_param_adjustments"][
                        "optimized_pipeline"
                    ] = repaired_pipeline
                    optimization_decision = json.dumps(
                        optimization_decision_json)
                    can_execute -= 1

                    logger.info(
                        f"--- {BENCH_PATH.split('/')[-1]}: executor_agent  ---"
                    )

                    # logger.info(f"\n\n{execution_result.json_dict}")
                    # logger.info(f"\n\n{execution_result.raw.strip()}")

                else:
                    # ,
                    execution_success = True
                    break

            if execution_success:
                # ,
                # execution finished, exit retry loop
                break

            logger.info(
                f"=======================================================")
            logger.info(
                f"---    Generated invalid pipeline or parameters.    ---")
            logger.info(
                f"---                  Retrying: {retry_attempt+1}/{MAX_DECISION_RETRY}                  ---"
            )
            logger.info(
                f"=======================================================\n")

            if retry_attempt + 1 == MAX_DECISION_RETRY:
                exit(1)

            retry_attempt += 1

        RESULT.append(
            {
                "task_name": "decision",
                "token_usage": decision_token_usage,
                "time_elapsed": decision_time_usage,
            }
        )

        with exec_lock:
            remarks = _get_new_remarks(current_round + 1, BENCH_PATH)
            if exceed_context_limit:
                remarks = extract_function_json(potential_funcs, remarks)

        if json.loads(performance)["average_time_elapsed"] < best_time:
            best_time = json.loads(performance)["average_time_elapsed"]
            best_parameters["-unroll-count"] = json.loads(optimization_decision)[
                "passes_param_adjustments"]["unroll_count"]
            best_parameters["-unroll-allow-partial"] = json.loads(
                optimization_decision)["passes_param_adjustments"]["unroll_allow_partial"]
            best_parameters["-unroll-threshold"] = json.loads(
                optimization_decision)["passes_param_adjustments"]["unroll_threshold"]
            best_parameters["-inline-threshold"] = json.loads(
                optimization_decision)["passes_param_adjustments"]["inline_threshold"]
            best_parameters["-slp-threshold"] = json.loads(optimization_decision)[
                "passes_param_adjustments"]["slp_threshold"]
            best_parameters["-devirt"] = json.loads(optimization_decision)[
                "passes_param_adjustments"]["devirt_param"]

        e_tasks = create_evaluation_tasks(
            round_num=current_round + 1,
            new_remarks=json.dumps(remarks),
            currentround_performance=performance,
            lastround_performance=lastround_performance,
            base_performance=initial_performance,
            best_parameters=json.dumps(best_parameters),
            dcs_ctx=o_tasks[0],
        )

        crew = Crew(
            agents=[evaluation_agent],
            tasks=e_tasks,
            process=Process.sequential,
            verbose=1,
        )

        try:
            starttime = time.time()
            optimization_analysis = crew.kickoff()
            endtime = time.time()

            # LLM usage
            RESULT.append(
                {
                    "task_name": "evaluation_feedback",
                    "token_usage": dict(optimization_analysis.token_usage),
                    "time_elapsed": [endtime - starttime],
                }
            )

        except Exception as e:
            logger.error(
                f"Crew execution failed in Round {current_round + 1}: {e}")
            break

        # save results
        RESULT.append(
            {
                "round": current_round + 1,
                **json.loads(optimization_decision),
                **json.loads(performance),
            }
        )

        with open(HISTORY, "w", encoding="utf-8") as f:
            json.dump(RESULT, f, indent=2, ensure_ascii=False)
            logger.info(
                f"=======================================================")
            logger.info(
                f"---                 History stored.                 ---")
            logger.info(
                f"=======================================================\n")

        try:
            archive_path = f"{LOG_DIR}/history_{BENCH_PATH.split('/')[-1]}.json"
            shutil.copy2(HISTORY, archive_path)
            print(f"----- {archive_path}")
        except Exception as e:
            print(f"----- : {e}")
            pass

        with exec_lock:
            next_input_model = OptimizationInputFeedback(
                opt_analysis_summary=optimization_analysis.raw,
                new_remarks=remarks,
                platform_info=PLATFORM_INFO,
                runtime_performance=json.loads(performance),
                previous_pipeline_params=json.loads(optimization_decision)[
                    "passes_param_adjustments"
                ],
            )

        lastround_performance = performance
        try:
            current_input_json = next_input_model.json(exclude_none=True)
        except Exception as e:
            logger.error(
                f"Failed to parse feedback loop output into OptimizationInputFeedback: {e}"
            )
            break

        # logger.info(f"=======================================================")
        logger.info(
            f"---    Successfully generated input for Round {current_round + 2}.    ---"
        )

        if i > MAX_ROUNDS:
            logger.info(
                f"=======================================================")
            logger.info(
                f"---     Reached MAX_ROUNDS ({MAX_ROUNDS}). Stopping loop.     ---"
            )
            logger.info(
                f"=======================================================\n")
            break
        # i += 1

    logger.info(f"=======================================================")
    logger.info(
        "--- Final Optimization Loop Result (Last Round Input for Next) ---\n\n"
    )
    if optimization_result:
        print(optimization_result)
    else:
        print("Optimization loop finished without a final result.")

    try:
        archive_path = f"{LOG_DIR}/history_{BENCH_PATH.split('/')[-1]}.json"
        shutil.copy2(HISTORY, archive_path)
        print(f"----- {archive_path}")
    except Exception as e:
        print(f"----- : {e}")
        pass
