# v1

import re
import json

from difflib import get_close_matches
from crewai import TaskOutput


def extract_clean_json_string(task_result: TaskOutput) -> str | None:
    """
     CrewAI TaskOutput  JSON 

    1.  CrewAI  json_dict ()
    2.  raw  Markdown 
    3.  JSON  None

    Args:
        task_result: CrewAI  TaskOutput 

    Returns:
         JSON  None
    """

    if task_result.json_dict:
        #  A (): CrewAI 
        try:
            #  json.dumps  JSON 
            print("Standard json extracted.")
            return json.dumps(task_result.json_dict)
        except Exception:
            # 
            pass

    #  B:  raw 
    raw_output = task_result.raw.strip()

    #  ```json ... ```  {...} 
    #  Markdown 
    #  re.DOTALL 
    match = re.search(r"```(?:json)?\s*(.+)\s*```", raw_output, re.DOTALL)

    if match:
        # 
        potential_json_str = match.group(1).strip()
        print("Markdown code block extracted.")
    else:
        #  {  } 
        start = raw_output.find("{")
        end = raw_output.rfind("}")
        if start != -1 and end != 0 and end > start:
            print("Brackets content extracted.")
            potential_json_str = raw_output[start : end + 1].strip()
        else:
            #  JSON 
            print("Can not extract json.")
            return None

    # 
    try:
        #  JSON
        json.loads(potential_json_str)
        print("Loaded from string.")
        return potential_json_str
    except json.JSONDecodeError:
        #  JSON
        print("Can not loadd from string.")
        return None


def correct_llvm_pipeline(
    initial_pipeline: str, generated_pipeline: str, cutoff: float = 0.8
) -> str:
    """
     LLVM Pass  LLM  Pass 
    

    Args:
        initial_pipeline:  LLVM Pass 
        generated_pipeline: LLM  LLVM Pass 
        cutoff:  (0.0  1.0)

    Returns:
         LLVM Pass 
    """

    temp_pipeline = re.sub(r"[<>()=;,]", " ", initial_pipeline)

    # 1****
    correct_tokens = set(
        word.strip()
        for word in temp_pipeline.split()
        if word.strip() and len(word.strip()) > 1
    )

    # Pass
    #  pass<...>, pass(...), pass, require<...>, function<...> 
    top_level_elements = re.findall(
        r"[^,<>]+(?:\<[^<>]*\>)?(?:\([^()]*\))?", initial_pipeline + ","
    )
    # 
    correct_tokens.update(
        element.strip() for element in top_level_elements if element.strip()
    )

    temp_generated = generated_pipeline

    temp_generated = re.sub(r"([<>(),;=])", r" \1 ", temp_generated)

    #  correct_tokens  Pass 
    sorted_correct_tokens = sorted(list(correct_tokens), key=len, reverse=True)

    #  Pass  ""
    tokens_to_check = re.split(r"([,<>():;=])", temp_generated)

    corrected_tokens_list = []
    for token in tokens_to_check:
        token = token.strip()
        if not token or re.match(r"[<>():;=,]", token):  # 
            corrected_tokens_list.append(token)
            continue

        is_split = False
        #  token
        for correct_token_1 in sorted_correct_tokens:
            if token.startswith(correct_token_1) and len(token) > len(correct_token_1):
                # 
                remaining_part = token[len(correct_token_1) :]

                #  Pass 
                for correct_token_2 in sorted_correct_tokens:
                    if remaining_part.startswith(correct_token_2):
                        # 
                        #  correct_token_1 + ',' + 
                        corrected_tokens_list.append(correct_token_1)
                        corrected_tokens_list.append(",")  #  Pass 
                        # 
                        corrected_tokens_list.extend(
                            correct_llvm_pipeline("", remaining_part, cutoff).split(",")
                        )
                        is_split = True
                        break
                if is_split:
                    break

        if not is_split:
            # 3
            corrected_tokens_list.append(token)

    # 
    temp_corrected_pipeline = "".join(corrected_tokens_list)
    temp_corrected_pipeline = re.sub(
        r",\s*,", ",", temp_corrected_pipeline
    )  # 

    #  ( difflib )
    #  token ,  ([, <, >, (, ), =, ;])  Pass  (,) 
    final_tokens = re.split(r"([,<>():;=])", temp_generated)
    final_corrected_tokens = []

    for token in final_tokens:
        token = token.strip()
        if not token:  # 
            continue

        # 
        if (
            re.match(r"[<>():;=,]", token)
            or token.isdigit()
            or re.match(r"^-?\d+(\.\d+)?$", token)
        ):
            final_corrected_tokens.append(token)
            continue

        # 
        matches = get_close_matches(token, sorted_correct_tokens, n=1, cutoff=cutoff)

        if matches:
            #  Pass/
            final_corrected_tokens.append(matches[0])
        else:
            #  (Pass)
            final_corrected_tokens.append(token)

    #  Pass 
    result_pipeline = ""
    for token in final_corrected_tokens:
        if re.match(r"[<>():;=,]", token):
            # 
            result_pipeline += token
        else:
            # Pass/
            if result_pipeline and not re.match(r"[<,(,]", result_pipeline[-1]):
                # Pass/
                #  PassName, PassName  PassName<...>, PassName<...> 
                pass
            result_pipeline += token

    # 
    result_pipeline = re.sub(r",+", ",", result_pipeline)
    
    result_pipeline = re.sub(r"\s*([<>():;=,])\s*", r"\1", result_pipeline)

    return result_pipeline.strip()
