from ..utils.convert_to_openai_messages import convert_to_openai_messages
from .setup_text_llm import setup_text_llm


def convert_to_coding_llm(text_llm, debug_mode=False):
    """
    Takes a text_llm
    returns an OI Coding LLM (a generator that takes OI messages and streams deltas with `message`, 'language', and `code`).
    """

    def coding_llm(messages):
        # First, tell it how to run code.

        # Old way:
        # system_message += "\nTo execute code on the user's machine, write a markdown code block *with the language*, i.e:\n\n```python\nprint('Hi!')\n```\n\nYou will receive the output ('Hi!'). Use any language."

        # New way:
        gaslight = None
        if messages[-1]["role"] == "user":
            # Last message came from the user.
            if messages[-1]["message"].lower() not in [
                "hello",
                "hi",
                "hey",
                "helo",
                "hii",
                "hi!",
            ]:  # :)
                gaslight = """Let's explore this. I can run code on your machine by writing the code in a markdown code block. This works if I put a newline after ```shell, ```python, ```applescript, etc. then write code. I'm going to try to do this for your task **after I make a plan**. I'll put the *correct* language after the "```"."""
        else:
            # Last message came from the assistant.

            # (The below should actually always be True in OI if last message came from the assistant)
            # I think we don't need this actually.
            """
            if "output" in messages[-1]:
                if messages[-1]["output"] != "No output":
                    gaslight = "(Thought: I see that the code I just ran produced an output. The next message I send will go to the user.)"
                elif messages[-1]["output"] == "No output":
                    gaslight = "(Thought: I see that the code I just ran produced no output. The next message I send will go to the user.)"
            """

        if gaslight:
            messages.append({"role": "assistant", "message": gaslight})

        # If it tried to use Jupyter, let it know.
        if "code" in messages[-1]:
            if any([line.startswith("!") for line in messages[-1]["code"].split("\n")]):
                if "syntax" in messages[-1]["output"].lower():  # Detect error
                    messages[-1][
                        "output"
                    ] += "\nRemember you are not in a Jupyter notebook. Run shell by writing a markdown shell codeblock, not '!'."

        messages = convert_to_openai_messages(messages, function_calling=False)

        inside_code_block = False
        accumulated_block = ""
        language = None

        for chunk in text_llm(messages):
            if debug_mode:
                print("Chunk in coding_llm", chunk)

            if "choices" not in chunk or len(chunk["choices"]) == 0:
                # This happens sometimes
                continue

            content = chunk["choices"][0]["delta"].get("content", "")

            accumulated_block += content

            if accumulated_block.endswith("`"):
                # We might be writing "```" one token at a time.
                continue

            # Did we just enter a code block?
            if "```" in accumulated_block and not inside_code_block:
                inside_code_block = True
                accumulated_block = accumulated_block.split("```")[1]

            # Did we just exit a code block?
            if inside_code_block and "```" in accumulated_block:
                return

            # If we're in a code block,
            if inside_code_block:
                # If we don't have a `language`, find it
                if language is None and "\n" in accumulated_block:
                    language = accumulated_block.split("\n")[0]

                    # Default to python if not specified
                    if language == "":
                        language = "python"
                    else:
                        # Removes hallucinations containing spaces or non letters.
                        language = "".join(char for char in language if char.isalpha())

                    output = {"language": language}

                    # If we recieved more than just the language in this chunk, send that
                    if content.split("\n")[1]:
                        output["code"] = content.split("\n")[1]

                    yield output

                # If we do have a `language`, send the output as code
                elif language:
                    yield {"code": content}

            # If we're not in a code block, send the output as a message
            if not inside_code_block:
                yield {"message": content}

    return coding_llm
