from functools import lru_cache
import json
from environs import Env
from openai import OpenAI

env = Env()
env.read_env()


def clean_json_markdown(text: str) -> str:
    # remove ```json
    if text.lstrip().startswith("```"):
        first_newline = text.find("\n")
        if first_newline != -1:
            text = text[first_newline + 1 :]

    # remove ```
    text = text.rstrip()
    if text.endswith("```"):
        text = text[:-3]

    # remove leading and trailing whitespace
    return text.strip()


@lru_cache(maxsize=32)
def format_json_string_by_ai(json_like_string: str) -> str:
    json_formatter = OpenAI(max_retries=10)
    completion = json_formatter.chat.completions.create(
        model="gpt-4o-2024-08-06",
        messages=[
            {
                "role": "system",
                "content": "You are an excellent JSON formatter. please format the user input JSON-like string with indent=2 carefully. PAY ATTENTION: The input is not sure to be an exact JSON string but a JSON-like string. You must always output a valid JSON string. Please do not return any other non-JSON string. Please do not add any other information. Please do not add any comment. Please do not add any newline.",
            },
            {"role": "user", "content": json_like_string},
        ],
        # response_format="json_object",
        # name="JSONFormatter",
        temperature=0.0,
        # user_id="JSONFormatter",
    )

    json_string = completion.choices[0].message.content
    return clean_json_markdown(json_string)


def json_string_to_dict_by_ai(json_like_string: str) -> dict:
    try:
        return json.loads(format_json_string_by_ai(json_like_string))
    except json.JSONDecodeError:
        raise ValueError(f"Invalid JSON-like string from llm: {json_like_string}")