import sys

import llm

from ag.parser import parser
from ag.sd import main


def prompt(message):
    model = llm.get_model()
    resp = model.prompt(message)

    def stream():
        buffer = ""
        for chunk in resp:
            buffer += chunk
            if len(buffer) > 100:
                # print("\r", flush=True, end="")
                buffer = ""
            if "\n" not in chunk:
                pass
                # print(f"{chunk}", flush=True, end="")
            # yield chunk.encode("utf-8")
            yield chunk

    parser.stream(stream())
    print("")
    # main(stream())


def _hook(lines):
    import ast

    code = "".join(lines)
    if not code.strip():
        return lines
    try:
        compile(code, "<input>", "exec")
        return lines
    except SyntaxError:
        prompt(code)
        # print(f"📢 Echo: {code.strip()}")

        # full_text = ""
        # console = Console()
        # with Live(
        #     console=console, refresh_per_second=12, vertical_overflow="visible"
        # ) as live:
        #     for chunk in resp:
        #         # print(chunk, end="", flush=True)
        #         full_text += chunk
        #         live.update(Markdown(full_text))
        return ["# silent\n"]


def _clean(lines):
    new = []
    for line in lines:
        if not "??" in line:
            line = line.replace("?", "\?")
        new.append(line)
    return new


def load():
    try:
        from IPython import get_ipython

        ip = get_ipython()
        if ip:
            if _clean not in ip.input_transformers_cleanup:
                ip.input_transformers_cleanup.append(_clean)
            if _hook not in ip.input_transformers_post:
                ip.input_transformers_post.append(_hook)
                print("✅ AG is active.")
    except ImportError:
        pass


def unload():
    ip = get_ipython()
    if _hook in ip.input_transformers_post:
        ip.input_transformers_post.remove(_hook)
        print("🛑 AG is disabled.")


class AG:
    def __truediv__(self, other):
        prompt(other)

    def __rtruediv__(self, other):
        prompt(other)


load()
