import gradio as gr
from models.genz_model import to_genz
from models.formal_model import to_formal


def style_transfer(text, mode):
    if mode == "Gen Z":
        return to_genz(text)
    else:
        return to_formal(text)


with gr.Blocks(theme=gr.themes.Soft(), title="Text Style Transfer") as demo:
    gr.Markdown(
        "# 🎨 Text Style Transfer\n"
        "Transform your text into **Gen Z slang** or **formal language** using AI-powered models."
    )

    with gr.Row(equal_height=True):
        # LEFT COLUMN
        with gr.Column(scale=1):
            input_text = gr.Textbox(
                lines=6,
                label="Input Text",
                placeholder="Enter your text here...",
                info="Type the text you want to transform."
            )

            style_choice = gr.Radio(
                choices=["Gen Z", "Formal"],
                value="Gen Z",
                label="Target Style",
                info="Choose the style to convert to."
            )

            submit_btn = gr.Button("Transform", variant="primary")

        # RIGHT COLUMN
        with gr.Column(scale=1):
            output_text = gr.Textbox(
                label="Transformed Text",
                interactive=False,
                info="The transformed text will appear here.",
                lines=6
            )

    gr.Examples(
        examples=[
            ["Hello, how are you?", "Gen Z"],
            ["I am fine, thank you.", "Formal"],
            ["What's up?", "Gen Z"],
            ["This is a test message.", "Formal"]
        ],
        inputs=[input_text, style_choice],
        outputs=output_text,
        fn=style_transfer,
        cache_examples=True
    )

    submit_btn.click(
        fn=style_transfer,
        inputs=[input_text, style_choice],
        outputs=output_text
    )

    gr.Markdown("---\n*Powered by Hugging Face Transformers and Gradio.*")


if __name__ == "__main__":
    demo.launch()