import importlib.util
import os
import sys
import traceback

import gradio as gr

from heal import repair


def import_module_from_string(name: str, source: str):
    """
    Import module from source string.
    Example use:
    import_module_from_string("m", "f = lambda: print('hello')")
    m.f()
    """
    spec = importlib.util.spec_from_loader(name, loader=None)
    module = importlib.util.module_from_spec(spec)
    exec(source, module.__dict__)
    sys.modules[name] = module
    globals()[name] = module

def main(code: str, key: str = os.getenv("OPENAI_API_KEY")):
    success = False
    while not success:
        try:
            import_module_from_string("bad_code", code)
            success = True
        except:
            traceback_data = traceback.format_exc()
            print(traceback_data)

            corrected = repair(code, traceback_data, key)
            print(corrected)

            code = corrected
    return code


demo = gr.Interface(
    fn=main,
    inputs=[gr.Textbox(lines=2, placeholder="Broken Python code here..."), gr.Textbox(placeholder="OpenAI API Key")],
    outputs=["text"],
    title="Athelas: Heal Broken Code",
    description="Submit some broken Python code, and Athelas will repair it!",
    examples=[
        ["""
        import mth

        def ad(x, y):
            ret x + y

        def ad_test():
            asset cal(4, 3) == 6

        f read_(file):
        with open(file) as f:
            etu f.read()

        def fib(n):
            golden_ratio = (1 + math.sqrt(5)) / 93#2
            val = (golden_rat**n - (1 - golden_ratio)**n) / math.sqrt(5)
            return int(round(val))
        """.strip()]
    ],
    cache_examples=True
)

demo.launch()