# v1

import os
import re
import json
import logging
import subprocess

# from crewai import Agent, Task, Crew, Process, LLM
from crewai.tools import BaseTool
from pydantic import BaseModel, Field
from typing import Dict, Union, Type, Set

from scripts.merge_remarks_no_order import merge_json as merge_no_order

from utils.data_contracts import (
    OptimizationInput,
    OptimizationDecision,
    EvaluationReport,
    PlatformInfo,
    RuntimePerformance,
    StaticFeatAndRemarksItem,
    FunctionFeatures,
    GenericFeatures,
    OptimizationInputFeedback,
    Remarks,
    FunctionRemarks,
    IRContent,
)

logger = logging.getLogger("CompilationCrew.Tools")

remote_configs_file = "scripts/remote_exe/remote_configs.json"


# --- Support functions ---
def _get_cross_flag(target):
    """Cross compile flags (target arch mainly) from remote configs"""
    with open(remote_configs_file, "r") as f:
        remote_configs = json.load(f)

    return remote_configs.get(target, {"cross_compile_target": ""})["cross_compile_target"]


def _get_platform_info(target):
    """Target platform info from remote configs"""
    with open(remote_configs_file, "r") as f:
        remote_configs = json.load(f)

    return remote_configs.get(target)["platform_info"]


def _get_opt_pipe_param(decision_json: str):
    """
    Get optimization parameters string.
    """
    parts = []

    # print(decision_json)

    data = json.loads(decision_json)

    # optimized_pipeline
    if data:
        pipeline = data["optimized_pipeline"]
        # parts.append(f'-passes="{pipeline}"'.replace(" ", "").replace('<>', ''))

        pattern = r"(?P<name>\b(?:inline|loop-unroll))<(?P<params>[^<>]*)>"

        def repl(match):
            name = match.group("name")
            params = match.group("params")
            # Remove both inline-threshold and unroll-threshold key-value pairs
            cleaned = re.sub(
                r"(?:inline|unroll)-threshold\s*=\s*[^;>]+;?", "", params)
            # Remove redundant semicolons
            cleaned = re.sub(r";{2,}", ";", cleaned)
            cleaned = cleaned.strip("; ")
            return f"{name}<{cleaned}>" if cleaned else f"{name}<>"

        parts.append(
            re.sub(
                pattern,
                repl,
                f'-passes="{pipeline}"'.replace(" ", "").replace("<>", ""),
            )
        )

        # param_adjustments
        for key, value in data.items():
            if key == "optimized_pipeline":
                continue
            # 0/1
            if isinstance(value, bool):
                value_str = "1" if value else "0"
            #
            elif isinstance(value, str):
                value_str = f'"{value}"'
            #
            else:
                value_str = str(value)

            #
            clean_key = key.replace("<>", "")
            if key != "mtriple":
                clean_key = clean_key.replace("_", "-")

            if data[key] == "auto":
                #
                continue

            if key in ["devirt_param"]:
                continue
            parts.append(f"-{clean_key}={value_str}")

    return " ".join(parts)


def _check_pipeline_brackets(pipeline: str) -> bool:
    """
    Check pipeline's brackets. 
     (), <>, []
    """
    stack = []
    bracket_map = {")": "(", ">": "<", "]": "["}

    for i, ch in enumerate(pipeline):
        if ch in "(<[":
            stack.append(ch)
        elif ch in ")>]":
            if not stack or stack[-1] != bracket_map[ch]:
                print(
                    f"Mismatched or unbalanced bracket at position {i}: {ch}")
                return False
            stack.pop()

    if stack:
        print(f"Unmatched opening brackets remaining: {stack}")
        return False

    return True


def _init_features(bench_path, initial_pipeline):
    try:
        result0 = subprocess.run(
            [
                "./scripts/extract_features.sh",
                f"{bench_path}/IR/linked.ll",
                initial_pipeline,
            ],
            capture_output=True,
            text=True,
            check=True,
        )
        print(result0.stdout)

        result1 = subprocess.run(
            # ['./scripts/input_schema_gen.py', f"{bench_path}/features"],
            ["./scripts/input_json_gen.py", f"{bench_path}/features"],
            capture_output=True,
            text=True,
            check=True,
        )
        print(result1.stdout)

        result2 = subprocess.run(
            # ['./scripts/merge_remarks.py'],
            ["./scripts/merge_remarks_no_order.py", f"--bench_path", bench_path],
            capture_output=True,
            text=True,
            check=True,
        )
        print(result2.stdout)

        with open(
            f"{bench_path}/jsons/origin_with_o3remarks_order.json",
            "r",
            encoding="utf-8",
        ) as file:
            # data = json.load(file)
            data = file.read()
            # print(data)
            return data

    except subprocess.CalledProcessError as e:
        logger.error(f"Feature extraction script failed: {e.stderr}")
        return json.dumps({"error": f"Script failed: {e.stderr}"})


def _run_init(bench_path, ldflags, runflags, target=""):
    """Run initial round (baseline)."""
    cmds = []
    cur_round = "init"
    try:
        cmds = _generate_compile_cmd(
            cur_round,
            f"{bench_path}/IR/linked.ll",
            pipeline='-passes="default<O3>"',
            ldflags=ldflags,
            runflags=runflags,
            target=target,
        )

    except subprocess.CalledProcessError as e:
        logger.error(f"Feature extraction script failed: {e.stderr}")
        return json.dumps({"error": f"Script failed: {e.stderr}"})

    results = []
    for cmd in cmds:
        try:
            result = subprocess.run(
                cmd,
                shell=True,
                capture_output=True,
                text=True,
                executable="/bin/bash",  # bash
                check=True,
            )
            results.append(result)
            print(result.stdout)
            print(result.stderr)
        except subprocess.CalledProcessError as e:
            return json.dumps({"error": f"Script failed: {e.stderr}; \n{cmd}"})

    try:
        json_file_path = (
            f"{bench_path}/perf_jsons/perf_stats_avg_linked_{cur_round}.json"
        )
        with open(json_file_path, "r") as f:
            perf_str = json.load(f)
            print(perf_str)
        return json.dumps(perf_str)
    except Exception as e:
        return json.dumps({"error": f"Read perf json failed: {e}"})


def _get_IR(ir_file_path):
    """Get IR content"""
    try:
        with open(ir_file_path, "r") as ir:
            # print(ir.read()[:-50])
            return ir.read()
    except Exception as e:
        logger.error(f"Execution failed: {e}")
        return json.dumps({"error": f"Internal error during execution: {e}"})


def _get_new_remarks(cur_round: int, bench_path: str) -> str:
    # logger.info(f"--- Calling external script for features: ./scripts/extract_features.sh {ir_file_path}")

    remarks_json = f"{bench_path}/jsons/origin_with_remarks_order_{cur_round}.json"
    remarks = None

    try:
        remarks_yaml = (
            bench_path + "/remarks/" + "linked" + f"_{cur_round}.remarks.yaml"
        )

        merge_no_order(
            remarks_yaml, f"{bench_path}/jsons/origin.json", remarks_json)

        with open(remarks_json, "r", encoding="utf-8") as f:
            remarks = json.load(f)
        # print(remarks)
        return remarks

    except subprocess.CalledProcessError as e:
        logger.error(f"Feature extraction script failed: {e.stderr}")
        return json.dumps({"error": f"Script failed: {e.stderr}"})


def _get_new(bench_path: str, cur_round: int):
    performance_json = f"{bench_path}/perf_jsons/perf_stats_avg_linked_{cur_round}.json"
    remarks_json = f"{bench_path}/jsons/origin_with_remarks_order_{cur_round}.json"
    perf = None
    remarks = None
    try:
        with open(performance_json, "r") as f:
            perf = json.load(f)
    except Exception as e:
        perf = {"error": f"Read perf json failed: {e}"}
    try:
        with open(remarks_json, "r") as f:
            remarks = json.load(f)
    except Exception as e:
        remarks = {"error": f"Read remarks json failed: {e}"}

    return perf, remarks


def _generate_compile_cmd(
    cur_round: int, ir_file: str, pipeline: str, ldflags: str, runflags: str, target=""
) -> str:
    """
    Generate shell commands, compile and run program. 
    """
    llvm_path = "<llvm_path>"
    clang = f"{llvm_path}/bin/clang"
    opt = f"{llvm_path}/bin/opt"
    llc = f"{llvm_path}/bin/llc"

    ir_dir = os.path.dirname(ir_file)
    bench_dir = os.path.dirname(ir_dir)
    benchname = os.path.splitext(os.path.basename(ir_file))[0]

    out_ir = f"{bench_dir}/IR/{benchname}_{cur_round}.ll"
    opt_obj = f"{bench_dir}/obj/{benchname}_{cur_round}.o"
    opt_bin = f"{bench_dir}/bin/{benchname}_{cur_round}.out"
    remarks_path = f"{bench_dir}/remarks/{benchname}_{cur_round}.remarks.yaml"

    cross_flag = _get_cross_flag(target=target)

    cmds = [
        f"""
    set -e
    {opt} {ir_file} -S -o {out_ir} {pipeline} -pass-remarks='.*' -pass-remarks-missed='.*' -pass-remarks-analysis='.*' -pass-remarks-output="{remarks_path}" 2> >(grep -v 'remark:')
    """,
        f"""{llc} {out_ir} -filetype=obj -o {opt_obj}""",
        f"""{clang} {opt_obj} {cross_flag} -flto -o {opt_bin} {ldflags}""",
    ]
    if target == "":
        cmds.append(
            f"""./scripts/perf_.sh {opt_bin} {runflags}""",
        )
    else:
        cmds.append(
            f"""python ./scripts/remote_exe/remote_exe.py --curround {cur_round} --benchpath {bench_dir} --runflags '{runflags}' --target {target}"""
        )
    return cmds


def _compile_and_run(
    cur_round: int,
    ir_file_path: str,
    decision_json: str,
    ldflags: str,
    runflags: str,
    target=""
) -> str:
    logger.info(
        f"--- Calling external script for compilation. IR file: {ir_file_path}")

    try:
        pipeline_str = _get_opt_pipe_param(decision_json)
        # print("Pipeline:", pipeline_str)

        # check pipeline
        if not _check_pipeline_brackets(pipeline_str):
            return json.dumps({"error": "Pipeline has unbalanced brackets"})

    except Exception as e:
        return f"error: Create command failed: {e}; {ir_file_path}; \n{decision_json}"

    try:
        cmds = _generate_compile_cmd(
            cur_round, ir_file_path, pipeline_str, ldflags, runflags, target
        )
        print("\n======= Running shell pipeline =======")
        # print(cmds)
        # print("=====================================")

    except Exception as e:
        return json.dumps({"error": f"Compilation: {e}"})

    results = []
    for cmd in cmds:
        try:
            result = subprocess.run(
                cmd,
                shell=True,
                capture_output=True,
                text=True,
                executable="/bin/bash",  # bash
                check=True,
            )
            results.append(result)
            print(result.stdout)
            print(result.stderr)
        except subprocess.CalledProcessError as e:
            return json.dumps({"error": f"Script failed: {e.stderr}; \n{cmd}"})

    try:
        json_file_path = f"{os.path.dirname(os.path.dirname(ir_file_path))}/perf_jsons/perf_stats_avg_linked_{cur_round}.json"
        with open(json_file_path, "r") as f:
            perf_str = json.load(f)
            print(perf_str)
        return json.dumps(perf_str)
    except Exception as e:
        return json.dumps({"error": f"Read perf json failed: {e}"})


def _generate_compile_cmd_2_stage(
    cur_round: int,
    ir_file: str,
    pipeline: str,
    ldflags: str,
    runflags: str,
    potential_irs: Set[str],
    other_irs: Set[str],
    target=""
) -> str:
    """
    Generate shell commands. 
    """
    llvm_path = "<llvm_path>"
    clang = f"{llvm_path}/bin/clang"
    opt = f"{llvm_path}/bin/opt"
    llc = f"{llvm_path}/bin/llc"
    link = f"{llvm_path}/bin/llvm-link"

    ir_dir = os.path.dirname(ir_file)
    bench_dir = os.path.dirname(ir_dir)
    benchname = os.path.splitext(os.path.basename(ir_file))[0]

    link_ir_potential = f"{bench_dir}/IR/{benchname}_potential_{cur_round}_base.ll"
    link_ir_others = f"{bench_dir}/IR/{benchname}_others_{cur_round}_base.ll"

    out_ir_potential = f"{bench_dir}/IR/{benchname}_potential_{cur_round}.ll"
    out_ir_others = f"{bench_dir}/IR/{benchname}_others_{cur_round}.ll"

    opt_obj_potential = f"{bench_dir}/obj/{benchname}_potential_{cur_round}.o"
    opt_obj_others = f"{bench_dir}/obj/{benchname}_others_{cur_round}.o"

    opt_bin = f"{bench_dir}/bin/{benchname}_{cur_round}.out"
    remarks_path = f"{bench_dir}/remarks/{benchname}_{cur_round}.remarks.yaml"
    other_remarks_path = f"{bench_dir}/remarks/{benchname}_other_{cur_round}.remarks.yaml"

    cross_flag = _get_cross_flag(target=target)

    cmds = [
        f"""{link} {" ".join(potential_irs)} -S -o {link_ir_potential} """,
        f"""{link} {" ".join(other_irs)} -S -o {link_ir_others} """,
        f"""
    set -e
    {opt} {link_ir_potential} -S -o {out_ir_potential} {pipeline} -pass-remarks='.*' -pass-remarks-missed='.*' -pass-remarks-analysis='.*' -pass-remarks-output="{remarks_path}" 2> >(grep -v 'remark:')
    """,
        f"""
    set -e
    {opt} {link_ir_others} -S -o {out_ir_others} -passes="default<O3>" -pass-remarks='.*' -pass-remarks-missed='.*' -pass-remarks-analysis='.*' -pass-remarks-output="{other_remarks_path}" 2> >(grep -v 'remark:')
    """,
        f"""{llc} {out_ir_potential} -filetype=obj -o {opt_obj_potential}""",
        f"""{llc} {out_ir_others} -filetype=obj -o {opt_obj_others}""",
        f"""{clang} {opt_obj_potential} {opt_obj_others} -flto -o {opt_bin} {ldflags}""",
    ]
    if target == "":
        cmds.append(
            f"""./scripts/perf_.sh {opt_bin} {runflags}""",
        )
    else:
        cmds.append(
            f"""python ./scripts/remote_exe/remote_exe.py --curround {cur_round} --benchpath {bench_dir} --runflags '{runflags}' --target {target}"""
        )
    return cmds


def _compile_and_run_2_stage(
    cur_round: int,
    ir_file_path: str,
    decision_json: str,
    ldflags: str,
    runflags: str,
    potential_irs: Set[str],
    other_irs: Set[str],
    target=""
) -> str:
    logger.info(
        f"--- Calling external script for compilation. IR file: {ir_file_path}")

    try:
        pipeline_str = _get_opt_pipe_param(decision_json)
        # print("Pipeline:", pipeline_str)

        #  pipeline
        if not _check_pipeline_brackets(pipeline_str):
            return json.dumps({"error": "Pipeline has unbalanced brackets"})

    except Exception as e:
        return f"error: Create command failed: {e}; \n{decision_json}"

    try:
        cmds = _generate_compile_cmd_2_stage(
            cur_round, ir_file_path, pipeline_str, ldflags, runflags, potential_irs, other_irs, target
        )
        print("\n======= Running shell pipeline =======")
        # print(cmds)
        # print("=====================================")

    except Exception as e:
        return json.dumps({"error": f"Compilation: {e}"})

    results = []
    for cmd in cmds:
        try:
            result = subprocess.run(
                cmd,
                shell=True,
                capture_output=True,
                text=True,
                executable="/bin/bash",
                check=True,
            )
            results.append(result)
            print(result.stdout)
            print(result.stderr)
        except subprocess.CalledProcessError as e:
            return json.dumps({"error": f"Script failed: {e.stderr}; \n{cmd}"})

    try:
        json_file_path = f"{os.path.dirname(os.path.dirname(ir_file_path))}/perf_jsons/perf_stats_avg_linked_{cur_round}.json"
        with open(json_file_path, "r") as f:
            perf_str = json.load(f)
            print(perf_str)
        return json.dumps(perf_str)
    except Exception as e:
        return json.dumps({"error": f"Read perf json failed: {e}"})
