import gradio as gr
from huggingface_hub import InferenceClient
import os
import spaces
import textwrap
from dotenv import load_dotenv

# https://huggingface.co/learn/cookbook/en/enterprise_cookbook_gradio
load_dotenv(override=True)

TEXT_MODEL_NAME = "microsoft/Phi-3-mini-4k-instruct:featherless-ai"

CODE_LANGUAGES = [("Python", "python"), ("C++", "cpp"), ("TypeScript", "typescript")]

CODE_EXAMPLE = textwrap.dedent("""\
    def factorial(num: int) -> int:
        n = 1
        for idx in range(1, num + 1)
            n *= idx
            
        return n
""")

CLIENT = InferenceClient(
    api_key=os.environ["HF_TOKEN"],
)


@spaces.GPU
def code_converter(code: str, from_lang: str, to_lang: str) -> str:
    system_input = """You are a code-only generator. Do not introduce the code. Do not explain the code. 
    Do not include markdown code block backticks. Output only raw, valid code."""
    user_input = f"""Please convert this {code} from {from_lang} to {to_lang}."""
    messages = [
        {"role": "system", "content": system_input},
        {"role": "user", "content": user_input},
    ]

    response = CLIENT.chat_completion(
        messages=messages, model=TEXT_MODEL_NAME, max_tokens=500, seed=430
    )
    return response.choices[0].message.content


with gr.Blocks() as demo:
    gr.HTML("<h1 style='text-align: center;'>My Code Converter</h1>")
    gr.HTML(
        """<h3 style='text-align: center;'>This app uses Microsoft's Phi-3-mini-4k-instruct
            to convert code to another language.</h3>"""
    )  # Adds styled HTML

    with gr.Row():
        with gr.Column():
            from_lang_input = gr.Dropdown(
                choices=CODE_LANGUAGES,
                value="python",
                label="Select Source Language",
                interactive=True,
            )
            code_input = gr.Code(
                language="python",
                value=CODE_EXAMPLE,
                label="Input Code",
                buttons=["copy"],
            )

        with gr.Column():

            to_lang_input = gr.Dropdown(
                choices=CODE_LANGUAGES,
                value="cpp",
                label="Select Target Language",
                interactive=True,
            )
            code_output = gr.Code(
                label="Converted Code",
                language="typescript",
                buttons=["copy"],
            )

    submit_btn = gr.Button("Click Here To Convert Code")

    from_lang_input.change(
        fn=lambda lang: gr.Code(language=lang),
        inputs=from_lang_input,
        outputs=code_input,
    )

    to_lang_input.change(
        fn=lambda lang: gr.Code(language=lang),
        inputs=to_lang_input,
        outputs=code_output,
    )

    submit_btn.click(
        fn=code_converter,
        inputs=[code_input, from_lang_input, to_lang_input],
        outputs=code_output,
    )
demo.launch()